Added an exit method to pdm.perf.
[pdm.git] / pdm / sshsock.py
CommitLineData
b4cc75dc
FT
1import sys, os
2import subprocess, socket, fcntl, select
3
4class 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)
34a3ccd6
FT
15 head = self.recv(5)
16 if head != "SSOCK":
17 raise socket.error("unexpected reply from %s: %r" % (host, head))
18 head = self.recv(1)
19 if head == "+":
e3250828
FT
20 buf = ""
21 while True:
22 r = self.recv(1)
23 if r == "":
24 raise socket.error("unexpected EOF in SSH socket stream")
25 elif r == "\n":
26 break
27 buf += r
34a3ccd6
FT
28 return
29 elif head == "-":
30 buf = ""
31 while True:
32 r = self.recv(1)
33 if r in ("\n", ""):
34 break
35 buf += r
36 raise socket.error(buf)
37 else:
38 raise socket.error("unexpected reply from %s: %r" % (host, head))
b4cc75dc
FT
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
62def cli():
63 fcntl.fcntl(sys.stdin, fcntl.F_SETFL, fcntl.fcntl(sys.stdin, fcntl.F_GETFL) | os.O_NONBLOCK)
64 sk = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
65 try:
34a3ccd6
FT
66 try:
67 sk.connect(sys.argv[1])
68 except socket.error as err:
69 sys.stdout.write("SSOCK-connect: %s\n" % err)
1a8ebe55 70 sys.stdout.flush()
34a3ccd6
FT
71 return
72 sys.stdout.write("SSOCK+\n")
1a8ebe55 73 sys.stdout.flush()
b4cc75dc
FT
74 buf1 = ""
75 buf2 = ""
76 while True:
77 wfd = []
78 if buf1: wfd.append(sk)
79 if buf2: wfd.append(sys.stdout)
80 rfd, wfd, efd = select.select([sk, sys.stdin], wfd, [])
81 if sk in rfd:
82 ret = sk.recv(65536)
83 if ret == "":
84 break
85 else:
86 buf2 += ret
87 if sys.stdin in rfd:
88 ret = sys.stdin.read()
89 if ret == "":
90 break
91 else:
92 buf1 = ret
93 if sk in wfd:
94 ret = sk.send(buf1)
95 buf1 = buf1[ret:]
96 if sys.stdout in wfd:
97 sys.stdout.write(buf2)
98 sys.stdout.flush()
99 buf2 = ""
100 finally:
101 sk.close()
102
103if __name__ == "__main__":
104 cli()