From 9d11c29091555e72dbeba51bfc71840ae802d602 Mon Sep 17 00:00:00 2001
From: Kai Vogelgesang <kai@leafbla.de>
Date: Sun, 11 Oct 2020 00:40:40 +0200
Subject: [PATCH] Implement log files and stdin comments

---
 .gitignore |  3 +++
 handler.py | 19 ++++++++++---------
 main.py    | 20 +++++++++++++++++---
 view.py    | 14 ++++++++++++++
 4 files changed, 44 insertions(+), 12 deletions(-)
 create mode 100644 view.py

diff --git a/.gitignore b/.gitignore
index a81c8ee..c88400d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,6 @@
+# captures
+capture_*.txt
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
diff --git a/handler.py b/handler.py
index 1c1223a..a1b2ff6 100644
--- a/handler.py
+++ b/handler.py
@@ -7,15 +7,16 @@ d_c = Diff(Color.GREEN, Color.RED)
 d_s = Diff(Color.BLUE, Color.YELLOW)
 
 
-def handle(tag: str, data: bytes):
-    msg = data.decode().strip()
-    if msg.startswith("0c") or msg.startswith("0a"):
+def handle(tag: str, data: str):
+    if tag == "stdin":
+        print(f"[stdin] {data}")
+    if data.startswith("0c") or data.startswith("0a"):
         if tag == "server":
-            print(f"{Color.YELLOW}{msg}{Color.RESET}")
-        else:
-            print(f"{Color.RED}{msg}{Color.RESET}")
+            print(f"{Color.YELLOW}{data}{Color.RESET}")
+        elif tag == "client":
+            print(f"{Color.RED}{data}{Color.RESET}")
     else:
         if tag == "server":
-            d_s.consume(msg)
-        else:
-            d_c.consume(msg)
+            d_s.consume(data)
+        elif tag == "client":
+            d_c.consume(data)
diff --git a/main.py b/main.py
index 80a556d..c2339ab 100644
--- a/main.py
+++ b/main.py
@@ -3,16 +3,26 @@ import os
 import importlib
 import traceback
 import signal
+import binascii
+import sys
+import datetime
 
 import handler
 
 SERVER_PORT = 22023
 
 
-async def queue_handler(queue: asyncio.Queue):
+def stdin_callback(queue: asyncio.Queue):
+    asyncio.ensure_future(queue.put(("stdin", sys.stdin.readline().rstrip("\n"))))
+
+
+async def queue_handler(queue: asyncio.Queue, log_filename: str = None):
     last_edit_timestamp = os.stat(handler.__file__).st_mtime_ns
     while True:
         tag, data = await queue.get()
+        if log_filename is not None:
+            with open(log_filename, "a") as f:
+                f.write(f"{tag} {data}\n")
         try:
             timestamp = os.stat(handler.__file__).st_mtime_ns
             if timestamp > last_edit_timestamp:
@@ -30,7 +40,7 @@ async def queue_handler(queue: asyncio.Queue):
 
 async def read_to_queue(reader: asyncio.StreamReader, tag: str, queue: asyncio.Queue):
     while True:
-        data = await reader.readline()
+        data = (await reader.readline()).decode().strip()
         if not data:
             break
         await queue.put((tag, data))
@@ -40,6 +50,8 @@ async def main():
 
     queue = asyncio.Queue()
 
+    asyncio.get_event_loop().add_reader(sys.stdin, stdin_callback, queue)
+
     proc_server = await asyncio.create_subprocess_shell(
         f'tshark -l -f "udp src port {SERVER_PORT}" -T fields -e data',
         stdout=asyncio.subprocess.PIPE,
@@ -64,7 +76,9 @@ async def main():
         signal.SIGINT, lambda: print("terminating...")
     )
 
-    handler_task = asyncio.create_task(queue_handler(queue))
+    now = datetime.datetime.now()
+    log_filename = f"capture_{now.strftime('%Y-%m-%d_%H:%M:%S')}.txt"
+    handler_task = asyncio.create_task(queue_handler(queue, log_filename=log_filename))
 
     await asyncio.gather(server_task, client_task)
 
diff --git a/view.py b/view.py
new file mode 100644
index 0000000..7f331c5
--- /dev/null
+++ b/view.py
@@ -0,0 +1,14 @@
+import handler
+import sys
+
+
+def main():
+    filename = sys.argv[1]
+    with open(filename, "r") as f:
+        for line in f.readlines():
+            tag, data = line.split()
+            handler.handle(tag, data)
+
+
+if __name__ == "__main__":
+    main()