pdm/sshsock.py

   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         self.proc.stdin.flush()
  50         return len(data)
  51
  52     def recv(self, buflen, flags = 0):
  53         if (flags & socket.MSG_DONTWAIT) == 0:
  54             select.select([self.proc.stdout], [], [])
  55         return self.proc.stdout.read(buflen)
  56
  57     def fileno(self):
  58         return self.proc.stdout.fileno()
  59
  60     def __del__(self):
  61         self.close()
  62
  63 def cli():
  64     fcntl.fcntl(sys.stdin.buffer, fcntl.F_SETFL, fcntl.fcntl(sys.stdin.buffer, fcntl.F_GETFL) | os.O_NONBLOCK)
  65     sk = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
  66     try:
  67         try:
  68             sk.connect(sys.argv[1])
  69         except socket.error as err:
  70             sys.stdout.write("SSOCK-connect: %s\n" % err)
  71             sys.stdout.flush()
  72             return
  73         sys.stdout.write("SSOCK+\n")
  74         sys.stdout.flush()
  75         buf1 = b""
  76         buf2 = b""
  77         while True:
  78             wfd = []
  79             if buf1: wfd.append(sk)
  80             if buf2: wfd.append(sys.stdout.buffer)
  81             rfd, wfd, efd = select.select([sk, sys.stdin.buffer], wfd, [])
  82             if sk in rfd:
  83                 ret = sk.recv(65536)
  84                 if ret == b"":
  85                     break
  86                 else:
  87                     buf2 += ret
  88             if sys.stdin.buffer in rfd:
  89                 ret = sys.stdin.buffer.read()
  90                 if ret == b"":
  91                     break
  92                 else:
  93                     buf1 = ret
  94             if sk in wfd:
  95                 ret = sk.send(buf1)
  96                 buf1 = buf1[ret:]
  97             if sys.stdout.buffer in wfd:
  98                 sys.stdout.buffer.write(buf2)
  99                 sys.stdout.buffer.flush()
 100                 buf2 = b""
 101     finally:
 102         sk.close()
 103
 104 if __name__ == "__main__":
 105     cli()