diff options
Diffstat (limited to 'pkctl.py')
| -rw-r--r-- | pkctl.py | 114 |
1 files changed, 69 insertions, 45 deletions
@@ -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(): |