From 15e06d9b0cf16fb2b4c710de4b20020cc4db410c Mon Sep 17 00:00:00 2001
From: Kelvin Ly <kelvin.ly1618@gmail.com>
Date: Sun, 14 May 2023 09:34:21 -0400
Subject: [PATCH] Add TCP piping util
---
shroom_pipe.py | 30 ++++++++++++++++++++++++++++++
1 file changed, 30 insertions(+)
create mode 100644 shroom_pipe.py
diff --git a/shroom_pipe.py b/shroom_pipe.py
new file mode 100644
index 0000000..6b84acb
--- /dev/null
+++ b/shroom_pipe.py
@@ -0,0 +1,30 @@
+import asyncio
+import sys
+
+# copied from https://stackoverflow.com/questions/64303607/python-asyncio-how-to-read-stdin-and-write-to-stdout
+async def connect_stdin_stdout():
+ loop = asyncio.get_event_loop()
+ reader = asyncio.StreamReader()
+ protocol = asyncio.StreamReaderProtocol(reader)
+ await loop.connect_read_pipe(lambda: protocol, sys.stdin)
+ w_transport, w_protocol = await loop.connect_write_pipe(asyncio.streams.FlowControlMixin, sys.stdout)
+ writer = asyncio.StreamWriter(w_transport, w_protocol, reader, loop)
+ return reader, writer
+
+async def pipe_loop(reader, writer):
+ while True:
+ res = await reader.read(1024)
+ if not res:
+ break
+ writer.write(res)
+ await writer.drain()
+
+async def main():
+ reader, writer = await connect_stdin_stdout()
+ tcp_reader, tcp_writer = await asyncio.open_connection('localhost', 9876)
+ task_totcp = asyncio.create_task(pipe_loop(reader, tcp_writer))
+ task_fromtcp = asyncio.create_task(pipe_loop(tcp_reader, writer))
+ await asyncio.wait([task_totcp, task_fromtcp])
+
+if __name__ == "__main__":
+ asyncio.run(main())