Merge branch 'pty' into master

1 files changed, 69 insertions, 45 deletions

diff --git a/pkctl.py b/pkctl.py
index 6fa6058..caf8a20 100644
--- a/pkctl.py
+++ b/pkctl.py
@@ -1,5 +1,5 @@
 #!/usr/bin/python3
-import os, sys, signal, socket, threading, time
+import os, sys, signal, socket, selectors, time, tty
 
 # basic config
 SOCKET_FILE = "/run/pk/pk.sock"
@@ -53,62 +53,86 @@ def start_cmd():
 def stop_cmd():
     return signald(signal.SIGTERM)
 
-def attach_reader(sock, state):
-    while state['attached']:
-        try:
-            data = sock.recv(1024)
-        except:
-            data = b'\xde\xad'
-        if data == b'\xde\xad':
-            state['attached'] = False
-            break
-        if len(data) > 0:
-            pnnl(str(data, 'utf-8'))
-    sock.close()
-
 def attach_cmd():
     if not isd_running():
         return False
 
     sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
-    state = {}
-    reader_thread = threading.Thread(target=attach_reader, args=(sock, state))
-
     try:
         sock.connect(SOCKET_FILE)
     except:
         sock.close()
         return False
-    state['attached'] = True
-    reader_thread.start()
-    while state['attached']:
-        try:
-            line = bytes(input().strip(), 'utf-8')
-        except EOFError:
-            print('detach')
-            line = b'detach'
-        if line == b'detach':
-            try:
-                sock.sendall(b'\xde\xad')
-            except:
-                pass
-            state['attached'] = False
-        elif line == b'clear':
-            os.system('clear')
-            line = b'\xc0\xdeprompt'
-        elif len(line) < 1:
-            line = b'\xc0\xdeprompt'
-        if not state['attached']:
-            break
-        try:
-            sock.sendall(line)
-        except:
-            state['attached'] = False
-            break
+    sel = selectors.DefaultSelector()
+    sel.register(sys.stdin.fileno(), selectors.EVENT_READ, 0)
+    sel.register(sock, selectors.EVENT_READ, 1)
+    attached = True
+    pty_mode = False
+    stdin_mode = None
 
-    state['attached'] = False
+    while attached:
+        events = sel.select()
+        for event, mask in events:
+            if event.data == 0:
+                try:
+                    data = os.read(sys.stdin.fileno(), 1024)
+                except Exception:
+                    data = False
+                if not data:
+                    attached = False
+                    break
+                elif pty_mode:
+                    try:
+                        sock.sendall(data)
+                        continue
+                    except:
+                        attached = False
+                        break
+                elif data.strip() == b'detach':
+                    try:
+                        sock.sendall(b'\xde\xad')
+                    except:
+                        pass
+                    attached = False
+                    break
+                elif data.strip() == b'clear':
+                    os.system('clear')
+                    data = b'\xc0\xdeprompt'
+                elif len(data.strip()) < 1:
+                    data = b'\xc0\xdeprompt'
+                try:
+                    sock.sendall(data)
+                except:
+                    attached = False
+                    break
+            else:
+                try:
+                    data = sock.recv(1024)
+                except Exception:
+                    data = False
+                if not data or data == b'\xde\xad':
+                    if pty_mode:
+                        tty.tcsetattr(sys.stdin.fileno(), tty.TCSAFLUSH, stdin_mode)
+                        pty_mode = False
+                        print('turned off pty mode due to remote detach')
+                    attached = False
+                    break
+                elif pty_mode and data[:6] == b'\xc0\xdenpty':
+                    tty.tcsetattr(sys.stdin.fileno(), tty.TCSAFLUSH, stdin_mode)
+                    pty_mode = False
+                    print('turned off pty mode due to npty command')
+                    if len(data) > 6:
+                        pnnl(str(data[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:
+                    pnnl(str(data, 'utf-8'))
+    sel.close()
     sock.close()
-    reader_thread.join()
+    if pty_mode:
+        tty.tcsetattr(sys.stdin.fileno(), tty.TCSAFLUSH, stdin_mode)
     return True
 
 def main():