feb2f9fcf9c835d9496bb2b56542b62f06fd3eac
[pdm.git] / 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 += ["python", "-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, fcntl.F_SETFL, fcntl.fcntl(sys.stdin, 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 = ""
44         buf2 = ""
45         while True:
46             wfd = []
47             if buf1: wfd.append(sk)
48             if buf2: wfd.append(sys.stdout)
49             rfd, wfd, efd = select.select([sk, sys.stdin], wfd, [])
50             if sk in rfd:
51                 ret = sk.recv(65536)
52                 if ret == "":
53                     break
54                 else:
55                     buf2 += ret
56             if sys.stdin in rfd:
57                 ret = sys.stdin.read()
58                 if ret == "":
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 in wfd:
66                 sys.stdout.write(buf2)
67                 sys.stdout.flush()
68                 buf2 = ""
69     finally:
70         sk.close()
71
72 if __name__ == "__main__":
73     cli()