From c663d99624882d464b259f989fe9edb79a150989 Mon Sep 17 00:00:00 2001
From: Wesley Kerfoot <wjak56@gmail.com>
Date: Mon, 17 Jun 2019 00:37:49 -0400
Subject: [PATCH] Various improvements, add handler to list pids

---
 client.py  | 39 ++++++++++++++++++++++++++++++---------
 server.rkt | 10 +++++++++-
 2 files changed, 39 insertions(+), 10 deletions(-)

diff --git a/client.py b/client.py
index 0fa1414..ceec11c 100755
--- a/client.py
+++ b/client.py
@@ -7,26 +7,39 @@ from base64 import b64encode, b64decode
 from contextlib import closing
 from json import loads
 
+def read_socket(sock, chunk_size=1024):
+    buf = []
+
+    while True:
+        chunk = sock.recv(1024)
+        buf.append(chunk)
+        if b"\n" in chunk:
+            break
+
+    return b"".join(buf).decode("utf-8")
+
 def read_log(pid):
     with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as s:
         s.connect("/tmp/shelltalk.sock")
         s.send(("read {0}\n".format(pid)).encode("utf-8"))
 
-        buf = []
+        buf = read_socket(s)
 
-        while True:
-            chunk = s.recv(1024)
-            buf.append(chunk)
-            if b"\n" in chunk:
-                break
-
-    return [b64decode(l) for l in loads(b"".join(buf).decode("utf-8"))]
+        return [b64decode(l) for l in loads(buf)]
 
 def spawn(pid):
     with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as s:
         s.connect("/tmp/shelltalk.sock")
         s.send(("spawn {0}\n".format(pid)).encode("utf-8"))
 
+
+def list_pids():
+    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as s:
+        s.connect("/tmp/shelltalk.sock")
+        s.send("list\n".encode("utf-8"))
+
+        return loads(read_socket(s))
+
 def write(pid, msg):
     with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as s:
         s.connect("/tmp/shelltalk.sock")
@@ -50,6 +63,11 @@ parser.add_argument("-R",
                     nargs=1,
                     help="Would you like to read some logs?")
 
+parser.add_argument("-L",
+                    "--list",
+                    action="store_true",
+                    help="List all of our shells")
+
 args = parser.parse_args()
 
 if args.spawn:
@@ -58,6 +76,9 @@ if args.spawn:
 if args.write:
     write(*args.write)
 
+if args.list:
+    print(list_pids())
+
 if args.read:
     for line in read_log(*args.read):
-        print(line)
+        print(line.decode("utf-8"))
diff --git a/server.rkt b/server.rkt
index c227ea4..a6b8ee2 100755
--- a/server.rkt
+++ b/server.rkt
@@ -69,8 +69,13 @@
      (handle-messages
       (hash-remove loggers pid))]
 
+    [(cons 'list out)
+     (displayln (hash-keys loggers))
+     (write-to
+       (hash-keys loggers) out)
+     (handle-messages loggers)]
+
     [(list 'read pid out)
-     (displayln "got read message")
      ; Reads all the logs for a given pid
      (logger-send loggers pid (cons 'read out))
      (handle-messages loggers)]))
@@ -106,6 +111,9 @@
         (thread-send message-handler (cons 'kill pid))
         (close-socket in out)]
 
+       [(list "list")
+        (thread-send message-handler (cons 'list out))]
+
        [other
         (displayln other)
         (handle-connection in out)])]))