90 lines
2.5 KiB
Python
90 lines
2.5 KiB
Python
import asyncio
|
|
import os
|
|
import importlib
|
|
import traceback
|
|
import signal
|
|
import binascii
|
|
import sys
|
|
import datetime
|
|
|
|
import handler
|
|
|
|
SERVER_PORT = 22023
|
|
|
|
|
|
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:
|
|
print("[importlib.reload]")
|
|
importlib.reload(handler)
|
|
last_edit_timestamp = timestamp
|
|
|
|
if asyncio.iscoroutinefunction(handler.handle):
|
|
await handler.handle(tag, data)
|
|
else:
|
|
handler.handle(tag, data)
|
|
except Exception:
|
|
traceback.print_exc()
|
|
|
|
|
|
async def read_to_queue(reader: asyncio.StreamReader, tag: str, queue: asyncio.Queue):
|
|
while True:
|
|
data = (await reader.readline()).decode().strip()
|
|
if not data:
|
|
break
|
|
await queue.put((tag, data))
|
|
|
|
|
|
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,
|
|
stderr=asyncio.subprocess.DEVNULL,
|
|
)
|
|
proc_client = await asyncio.create_subprocess_shell(
|
|
f'tshark -l -f "udp dst port {SERVER_PORT}" -T fields -e data',
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.DEVNULL,
|
|
)
|
|
|
|
server_task = asyncio.create_task(
|
|
read_to_queue(proc_server.stdout, "server", queue)
|
|
)
|
|
|
|
client_task = asyncio.create_task(
|
|
read_to_queue(proc_client.stdout, "client", queue)
|
|
)
|
|
|
|
# This is just to supress the ugly KeyboardInterrupt print
|
|
asyncio.get_event_loop().add_signal_handler(
|
|
signal.SIGINT, lambda: print("terminating...")
|
|
)
|
|
|
|
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)
|
|
|
|
handler_task.cancel()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|