| 1 | import sys, os |
| 2 | import subprocess, socket, fcntl, select |
| 3 | |
| 4 | class sshsocket(object): |
| 5 | def __init__(self, host, path, user = None, port = None): |
| 6 | args = ["ssh"] |
| 7 | if user is not None: |
| 8 | args += ["-u", str(user)] |
| 9 | if port is not None: |
| 10 | args += ["-p", str(int(port))] |
| 11 | args += [host] |
| 12 | args += ["python3", "-m", "pdm.sshsock", path] |
| 13 | self.proc = subprocess.Popen(args, stdin=subprocess.PIPE, stdout=subprocess.PIPE, close_fds=True) |
| 14 | fcntl.fcntl(self.proc.stdout, fcntl.F_SETFL, fcntl.fcntl(self.proc.stdout, fcntl.F_GETFL) | os.O_NONBLOCK) |
| 15 | head = self.recv(5) |
| 16 | if head != b"SSOCK": |
| 17 | raise socket.error("unexpected reply from %s: %r" % (host, head)) |
| 18 | head = self.recv(1) |
| 19 | if head == b"+": |
| 20 | buf = b"" |
| 21 | while True: |
| 22 | r = self.recv(1) |
| 23 | if r == b"": |
| 24 | raise socket.error("unexpected EOF in SSH socket stream") |
| 25 | elif r == b"\n": |
| 26 | break |
| 27 | buf += r |
| 28 | return |
| 29 | elif head == b"-": |
| 30 | buf = b"" |
| 31 | while True: |
| 32 | r = self.recv(1) |
| 33 | if r in {b"\n", b""}: |
| 34 | break |
| 35 | buf += r |
| 36 | raise socket.error(buf.decode("utf-8")) |
| 37 | else: |
| 38 | raise socket.error("unexpected reply from %s: %r" % (host, head)) |
| 39 | |
| 40 | def close(self): |
| 41 | if self.proc is not None: |
| 42 | self.proc.stdin.close() |
| 43 | self.proc.stdout.close() |
| 44 | self.proc.wait() |
| 45 | self.proc = None |
| 46 | |
| 47 | def send(self, data, flags = 0): |
| 48 | self.proc.stdin.write(data) |
| 49 | return len(data) |
| 50 | |
| 51 | def recv(self, buflen, flags = 0): |
| 52 | if (flags & socket.MSG_DONTWAIT) == 0: |
| 53 | select.select([self.proc.stdout], [], []) |
| 54 | return self.proc.stdout.read(buflen) |
| 55 | |
| 56 | def fileno(self): |
| 57 | return self.proc.stdout.fileno() |
| 58 | |
| 59 | def __del__(self): |
| 60 | self.close() |
| 61 | |
| 62 | def cli(): |
| 63 | fcntl.fcntl(sys.stdin.buffer, fcntl.F_SETFL, fcntl.fcntl(sys.stdin.buffer, fcntl.F_GETFL) | os.O_NONBLOCK) |
| 64 | sk = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) |
| 65 | try: |
| 66 | try: |
| 67 | sk.connect(sys.argv[1]) |
| 68 | except socket.error as err: |
| 69 | sys.stdout.write("SSOCK-connect: %s\n" % err) |
| 70 | return |
| 71 | sys.stdout.write("SSOCK+\n") |
| 72 | buf1 = b"" |
| 73 | buf2 = b"" |
| 74 | while True: |
| 75 | wfd = [] |
| 76 | if buf1: wfd.append(sk) |
| 77 | if buf2: wfd.append(sys.stdout.buffer) |
| 78 | rfd, wfd, efd = select.select([sk, sys.stdin.buffer], wfd, []) |
| 79 | if sk in rfd: |
| 80 | ret = sk.recv(65536) |
| 81 | if ret == b"": |
| 82 | break |
| 83 | else: |
| 84 | buf2 += ret |
| 85 | if sys.stdin.buffer in rfd: |
| 86 | ret = sys.stdin.buffer.read() |
| 87 | if ret == b"": |
| 88 | break |
| 89 | else: |
| 90 | buf1 = ret |
| 91 | if sk in wfd: |
| 92 | ret = sk.send(buf1) |
| 93 | buf1 = buf1[ret:] |
| 94 | if sys.stdout.buffer in wfd: |
| 95 | sys.stdout.buffer.write(buf2) |
| 96 | sys.stdout.buffer.flush() |
| 97 | buf2 = b"" |
| 98 | finally: |
| 99 | sk.close() |
| 100 | |
| 101 | if __name__ == "__main__": |
| 102 | cli() |