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
  16     def close(self):
  17         if self.proc is not None:
  18             self.proc.stdin.close()
  19             self.proc.stdout.close()
  20             self.proc.wait()
  21             self.proc = None
  22
  23     def send(self, data, flags = 0):
  24         self.proc.stdin.write(data)
  25         return len(data)
  26
  27     def recv(self, buflen, flags = 0):
  28         if (flags & socket.MSG_DONTWAIT) == 0:
  29             select.select([self.proc.stdout], [], [])
  30         return self.proc.stdout.read(buflen)
  31
  32     def fileno(self):
  33         return self.proc.stdout.fileno()
  34
  35     def __del__(self):
  36         self.close()
  37
  38 def cli():
  39     fcntl.fcntl(sys.stdin.buffer, fcntl.F_SETFL, fcntl.fcntl(sys.stdin.buffer, fcntl.F_GETFL) | os.O_NONBLOCK)
  40     sk = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
  41     try:
  42         sk.connect(sys.argv[1])
  43         buf1 = b""
  44         buf2 = b""
  45         while True:
  46             wfd = []
  47             if buf1: wfd.append(sk)
  48             if buf2: wfd.append(sys.stdout.buffer)
  49             rfd, wfd, efd = select.select([sk, sys.stdin.buffer], wfd, [])
  50             if sk in rfd:
  51                 ret = sk.recv(65536)
  52                 if ret == b"":
  53                     break
  54                 else:
  55                     buf2 += ret
  56             if sys.stdin.buffer in rfd:
  57                 ret = sys.stdin.buffer.read()
  58                 if ret == b"":
  59                     break
  60                 else:
  61                     buf1 = ret
  62             if sk in wfd:
  63                 ret = sk.send(buf1)
  64                 buf1 = buf1[ret:]
  65             if sys.stdout.buffer in wfd:
  66                 sys.stdout.buffer.write(buf2)
  67                 sys.stdout.buffer.flush()
  68                 buf2 = b""
  69     finally:
  70         sk.close()
  71
  72 if __name__ == "__main__":
  73     cli()