Merge branch 'master' of ansgar.seatribe.se:/usr/local/src/pdm
[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]
16e7fd3d 12 args += ["python3", "-m", "pdm.sshsock", path]
b4cc75dc
FT
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 15 head = self.recv(5)
98ecc02e 16 if head != b"SSOCK":
34a3ccd6
FT
17 raise socket.error("unexpected reply from %s: %r" % (host, head))
18 head = self.recv(1)
98ecc02e
FT
19 if head == b"+":
20 buf = b""
e3250828
FT
21 while True:
22 r = self.recv(1)
98ecc02e 23 if r == b"":
e3250828 24 raise socket.error("unexpected EOF in SSH socket stream")
98ecc02e 25 elif r == b"\n":
e3250828
FT
26 break
27 buf += r
34a3ccd6 28 return
98ecc02e
FT
29 elif head == b"-":
30 buf = b""
34a3ccd6
FT
31 while True:
32 r = self.recv(1)
98ecc02e 33 if r in {b"\n", b""}:
34a3ccd6
FT
34 break
35 buf += r
98ecc02e 36 raise socket.error(buf.decode("utf-8"))
34a3ccd6
FT
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)
d14461c7 49 self.proc.stdin.flush()
b4cc75dc
FT
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
63def cli():
16e7fd3d 64 fcntl.fcntl(sys.stdin.buffer, fcntl.F_SETFL, fcntl.fcntl(sys.stdin.buffer, fcntl.F_GETFL) | os.O_NONBLOCK)
b4cc75dc
FT
65 sk = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
66 try:
34a3ccd6
FT
67 try:
68 sk.connect(sys.argv[1])
69 except socket.error as err:
70 sys.stdout.write("SSOCK-connect: %s\n" % err)
1a8ebe55 71 sys.stdout.flush()
34a3ccd6
FT
72 return
73 sys.stdout.write("SSOCK+\n")
1a8ebe55 74 sys.stdout.flush()
16e7fd3d
FT
75 buf1 = b""
76 buf2 = b""
b4cc75dc
FT
77 while True:
78 wfd = []
79 if buf1: wfd.append(sk)
16e7fd3d
FT
80 if buf2: wfd.append(sys.stdout.buffer)
81 rfd, wfd, efd = select.select([sk, sys.stdin.buffer], wfd, [])
b4cc75dc
FT
82 if sk in rfd:
83 ret = sk.recv(65536)
16e7fd3d 84 if ret == b"":
b4cc75dc
FT
85 break
86 else:
87 buf2 += ret
16e7fd3d
FT
88 if sys.stdin.buffer in rfd:
89 ret = sys.stdin.buffer.read()
90 if ret == b"":
b4cc75dc
FT
91 break
92 else:
93 buf1 = ret
94 if sk in wfd:
95 ret = sk.send(buf1)
96 buf1 = buf1[ret:]
16e7fd3d
FT
97 if sys.stdout.buffer in wfd:
98 sys.stdout.buffer.write(buf2)
99 sys.stdout.buffer.flush()
100 buf2 = b""
b4cc75dc
FT
101 finally:
102 sk.close()
103
104if __name__ == "__main__":
105 cli()