From 9c3765f9b87ee2941c6e330fbb3377c1e61fae0e Mon Sep 17 00:00:00 2001 From: Carson Fleming Date: Fri, 24 Mar 2023 07:08:56 -0700 Subject: Synchronize PTY stop sequence --- crypto.py | 4 +-- pkcli_stub.py | 84 ++++++++++++++++++++++++++++++++--------------------------- pkd_stub.py | 44 ++++++++++++++++++++----------- 3 files changed, 77 insertions(+), 55 deletions(-) diff --git a/crypto.py b/crypto.py index 662d5e6..77d795f 100644 --- a/crypto.py +++ b/crypto.py @@ -180,8 +180,8 @@ class PKSock: self.rpk = {'n': Crypto.b2i(self.recv()), 'e': Crypto.exp} return True - def handshake_server (self, server_rpk): - self.rpk = server_rpk + def handshake_server (self, server_pk): + self.rpk = server_pk self.raw_send(Crypto.i2b(self.nbytes)) rnbytes = Crypto.b2i(self.raw_recv(self.headsz)) if self.nbytes != rnbytes: diff --git a/pkcli_stub.py b/pkcli_stub.py index 476636d..09f7aae 100644 --- a/pkcli_stub.py +++ b/pkcli_stub.py @@ -36,7 +36,7 @@ def main(): port = int(port) print('working') - p,q,n,e,d = keygen(bits=bits) + p,q,n,e,d = Crypto.keygen(bits=bits) privkey = { 'n': n, 'd': d } refresh_hdb() print('done') @@ -62,16 +62,6 @@ def polymorph(): if os.fork() != 0: sys.exit(0) -def handshake(sock, privkey, rpubkey, bits): - nbytes, headsz = bits//8, 2 - sock.sendall(nbytes.to_bytes(headsz, 'big')) - rnbytes = int.from_bytes(sock.recv(headsz), 'big') - if rnbytes != nbytes: - return False - - send_encrypted(sock, privkey['n'].to_bytes(nbytes, 'big'), rpubkey['e'], rpubkey['n'], bits=bits) - return True - def refresh_hdb(): global hostkeys_url, hdb try: @@ -102,8 +92,22 @@ def get_hostkey(host): del hkdb[host] return False -def run_pty(sock, screen_is, screen_os): - term = screen_is.recv() +def pty_barrier(sock): + code = [0]*len(b'\xc0\xdeack') + + while bytes(code) != b'\xc0\xdeack': + buffer = sock.recv() + while len(buffer) > 0: + code = code[:-1]+buffer[0] + buffer = buffer[1:] + if bytes(code) == b'\xc0\xdeack': + break + + sock.stop_stream(len(buffer)) + +def run_pty(sock): + sock.start_stream() + term = sock.recv() sel = selectors.DefaultSelector() pid, shfd = pty.fork() if pid == 0: @@ -115,30 +119,36 @@ def run_pty(sock, screen_is, screen_os): try: sel.register(shfd, selectors.EVENT_READ, 0) - sel.register(sock, selectors.EVENT_READ, 1) + sel.register(sock.sock, selectors.EVENT_READ, 1) while True: events = sel.select() - for event, mask in events: + quit = False + + for event, _ in events: if event.data == 0: try: data = os.read(shfd, 1024) except: data = False if not data: - return True - screen_os.send(data) + quit = True + else: + sock.send(data) else: - try: - data = screen_is.recv() - except: - data = False + data = sock.recv() if not data: return False - elif data == b'\xc0\xdenpty': - return True - os.write(shfd, data) + elif data[:6] == b'\xc0\xdenpty': + quit = True + else: + os.write(shfd, data) + + if quit: + sock.send(b'\xc0\xdenpty') + pty_barrier(sock) + return True except: - return + return False finally: sel.close() try: @@ -152,29 +162,30 @@ def work(h_addr, port, privkey, bits): host = socket.gethostbyname(h_addr) except: host = h_addr - sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + raw_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) try: - sock.connect((host, port)) + raw_sock.connect((host, port)) + sock = PKSock(raw_sock, privkey, bits) hostkey = get_hostkey(host) if hostkey: rpubkey = hostkey else: rpubkey = {'n': server_modulus, 'e': exp} - if not handshake(sock, privkey, rpubkey, bits=bits): + if not sock.handshake_server(rpubkey): return True PS1 = '$ ' if 'PS1' in os.environ: PS1 = os.environ['PS1'] - send_encrypted(sock, PS1, rpubkey['e'], rpubkey['n'], bits=bits) + sock.send(PS1) while True: - cmd = recv_encrypted(sock, privkey['d'], privkey['n'], bits=bits) + cmd = sock.recv() if cmd == b'tunnel': - send_encrypted(sock, b'\xde\xad', rpubkey['e'], rpubkey['n'], bits=bits) + sock.send(b'\xde\xad') return True elif cmd == b'die': - send_encrypted(sock, b'\xde\xad', rpubkey['e'], rpubkey['n'], bits=bits) + sock.send(b'\xde\xad') return False elif cmd == b'refresh-hdb': if refresh_hdb(): @@ -182,11 +193,8 @@ def work(h_addr, port, privkey, bits): else: response = '[pk] Error: could not refresh host database.\n' elif cmd == b'pty': - screen_is = InStreamCipher(sock, privkey, bits=bits) - screen_os = OutStreamCipher(sock, rpubkey, bits=bits) - if not run_pty(sock, screen_is, screen_os): + if not run_pty(sock): return True - screen_os.send(b'\xc0\xdenpty') continue else: try: @@ -197,11 +205,11 @@ def work(h_addr, port, privkey, bits): response = str(response, 'utf-8') except Exception as e: response = '%s\n' % str(e) - send_encrypted(sock, '%s%s' % (response, PS1), rpubkey['e'], rpubkey['n'], bits=bits) + sock.send('%s%s' % (response, PS1)) except: return True finally: - sock.close() + raw_sock.close() if __name__ == '__main__': main() diff --git a/pkd_stub.py b/pkd_stub.py index f3ff1f7..c7209b2 100644 --- a/pkd_stub.py +++ b/pkd_stub.py @@ -139,14 +139,14 @@ def screens_pty(sel, screen, client): os.environ['TERM'] = 'xterm-256color' client['osc'].send(bytes(os.environ['TERM'], 'utf-8')) except: - tcp_unpty(sel, client) + tcp_unpty(sel, client, catchup=False) tcp_disconnect(sel, client) try: screen['sock'].sendall(b'\xc0\xdepty') except: screens_detach(sel, screen) - tcp_unpty(sel, client) + tcp_send_npty(sel, client) return def screens_read(sel, sock, screen): @@ -160,7 +160,7 @@ def screens_read(sel, sock, screen): if not data or data == b'\xde\xad': screens_detach(sel, screen) if screen['pty']: - tcp_unpty(sel, screen['pty'], npty_screen=False) + tcp_send_npty(sel, screen['pty']) return if screen['pty']: @@ -285,14 +285,15 @@ def register_screens(sel, socket_file): def tcp_disconnect(sel, client): global tcp_clients + if client not in tcp_clients: + return + sel.unregister(client['sock']) client['sock'].close() client['alive'] = False - - if client in tcp_clients: - idx = tcp_clients.index(client) - del tcp_clients[idx] - brint('[INFO] TCP Client %d disconnected.' % idx) + idx = tcp_clients.index(client) + del tcp_clients[idx] + brint('[INFO] TCP Client %d disconnected.' % idx) def tcp_dumpq(sel, client): global cmdq @@ -303,17 +304,30 @@ def tcp_dumpq(sel, client): except: tcp_disconnect(sel, client) -def tcp_unpty(sel, client, catchup=True, npty_screen=True): +def tcp_send_npty(sel, client): + try: + client['osc'].send(b'\xc0\xdenpty') + except: + tcp_disconnect(sel, client) + +def tcp_unpty(sel, client, catchup=True): if type(client['pty']) == dict: client['pty']['pty'] = False - if npty_screen and client['pty']['alive']: + if client['pty']['alive']: try: client['pty']['sock'].sendall(b'\xc0\xdenpty') except: screens_detach(sel, client['pty']) - del client['isc'] - del client['osc'] + + try: + client['osc'].send(b'\xc0\xdeack') + except: + tcp_disconnect(sel, client) + # this will become stop_stream(backtrack) + del client['isc'] + del client['osc'] client['pty'] = False + if catchup: tcp_dumpq(sel, client) @@ -333,15 +347,15 @@ def tcp_transport(sel, sock, client): return elif not client['pty']: brint('[%d]' % tcp_clients.index(client), data, end='', prompt=False) - elif data == b'\xc0\xdenpty': - tcp_unpty(sel, client) + elif data[:6] == b'\xc0\xdenpty': + tcp_unpty(sel, client, catchup=True) print('[INFO] npty acknowledged') else: try: client['pty']['sock'].sendall(data) except: screens_detach(sel, client['pty']) - tcp_unpty(sel, client, npty_screen=False) + tcp_send_npty(client) def tcp_handshake(sock): global privkey, bits, exp -- cgit v1.2.3