From 53bab3aafd4596f128cebc06dd27a9dfb27d5106 Mon Sep 17 00:00:00 2001
From: Jake Howard <git@theorangeone.net>
Date: Tue, 1 Jan 2019 19:07:00 +0000
Subject: [PATCH] Wrap it in a CLI

---
 scripts/test.sh  |  2 +-
 setup.cfg        |  6 ++++++
 tcp_nat_proxy.py | 46 +++++++++++++++++++++++++++++++++++++---------
 3 files changed, 44 insertions(+), 10 deletions(-)
 create mode 100644 setup.cfg

diff --git a/scripts/test.sh b/scripts/test.sh
index bfe744e..90d258d 100644
--- a/scripts/test.sh
+++ b/scripts/test.sh
@@ -4,5 +4,5 @@ set -e
 
 black --check tcp_nat_proxy.py
 isort -c tcp_nat_proxy.py
-flake8 tcp_nat_proxy.py
+flake8 tcp_nat_proxy.py --ignore=E128,E501
 mypy --strict-optional --ignore-missing-imports tcp_nat_proxy.py
diff --git a/setup.cfg b/setup.cfg
new file mode 100644
index 0000000..9eea40a
--- /dev/null
+++ b/setup.cfg
@@ -0,0 +1,6 @@
+[isort]
+multi_line_output=3
+include_trailing_comma=True
+force_grid_wrap=0
+use_parentheses=True
+line_length=88
diff --git a/tcp_nat_proxy.py b/tcp_nat_proxy.py
index 2b6a96b..e025b75 100644
--- a/tcp_nat_proxy.py
+++ b/tcp_nat_proxy.py
@@ -1,12 +1,24 @@
+import argparse
 import asyncio
+from collections import namedtuple
 
 import uvloop
 
+Route = namedtuple("Route", ["listen_port", "destination_host", "destination_port"])
+
 uvloop.install()
 
 BUFFER_SIZE = 4096
 
 
+def destination_host_display(route: Route):
+    return "{}:{}".format(route.destination_host, route.destination_port)
+
+
+def parse_argument(value):
+    return Route(*value.split(":"))
+
+
 async def pipe(reader, writer):
     try:
         while not reader.at_eof():
@@ -15,26 +27,42 @@ async def pipe(reader, writer):
         writer.close()
 
 
-async def create_proxy_pipe(listen_port, destination_host, destination_port):
+async def create_proxy_pipe(route: Route):
     async def handle_client(local_reader, local_writer):
         try:
             remote_reader, remote_writer = await asyncio.open_connection(
-                destination_host, destination_port
+                route.destination_host, route.destination_port
             )
-            pipe1 = pipe(local_reader, remote_writer)
-            pipe2 = pipe(remote_reader, local_writer)
-            await asyncio.gather(pipe1, pipe2)
+            await asyncio.gather(
+                pipe(local_reader, remote_writer), pipe(remote_reader, local_writer)
+            )
+        except ConnectionRefusedError:
+            print("Connection to {} refused".format(destination_host_display(route)))
+            pass
         finally:
             local_writer.close()
 
-    server = await asyncio.start_server(handle_client, "0.0.0.0", listen_port)
-    print("Serving on {}".format(server.sockets[0].getsockname()))
+    server = await asyncio.start_server(handle_client, "0.0.0.0", route.listen_port)
+    print(
+        "Routing from {} to {}".format(
+            route.listen_port, destination_host_display(route)
+        )
+    )
     await server.wait_closed()
 
 
 async def main():
-    await asyncio.gather(create_proxy_pipe(8888, "127.0.0.1", 8889))
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-R", "--route", action="append", required=True, type=parse_argument
+    )
+    args = parser.parse_args()
+    servers = [create_proxy_pipe(route) for route in args.route]
+    await asyncio.gather(*servers)
 
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    try:
+        asyncio.run(main())
+    except KeyboardInterrupt:
+        print("Process terminated")