Formatting.

src/shadowssh/__init__.py

@@ -1,9 +1,8 @@
-import logging 
+import logging
+
 from mololog.client import patch
 
-logging.basicConfig(
-    level = logging.INFO
-)
+logging.basicConfig(level=logging.INFO)
 
 log = logging.getLogger()
 patch("https://mololog.molodetz.nl/")

src/shadowssh/__main__.py

@@ -1,25 +1,31 @@
-from shadowssh.app import Application
+import argparse
 import asyncio
-import argparse 
+
+from shadowssh.app import Application
+
 
 def parse_args():
     parser = argparse.ArgumentParser(description="Shadowssh")
 
-    parser.add_argument("--host",type=str,required=False,default="0.0.0.0")
-    parser.add_argument("--port",type=int,required=False,default=443)
-    parser.add_argument("--host-service",type=str,required=False,default="127.0.0.1")
-    parser.add_argument("--port-service",type=int,required=False,default=4433)
-    parser.add_argument("--host-ssh", type=str,required=False,default="127.0.0.1")
-    parser.add_argument("--port-ssh", type=int,required=False,default=22)
+    parser.add_argument("--host", type=str, required=False, default="0.0.0.0")
+    parser.add_argument("--port", type=int, required=False, default=443)
+    parser.add_argument("--host-service", type=str, required=False, default="127.0.0.1")
+    parser.add_argument("--port-service", type=int, required=False, default=4433)
+    parser.add_argument("--host-ssh", type=str, required=False, default="127.0.0.1")
+    parser.add_argument("--port-ssh", type=int, required=False, default=22)
     return parser.parse_args()
 
 
 def serve():
     args = parse_args()
     app = Application()
-    asyncio.run(app.serve(args.host,args.port,args.host_service,args.port_service,args.host_ssh,args.port_ssh))
-
-
-
-
-
+    asyncio.run(
+        app.serve(
+            args.host,
+            args.port,
+            args.host_service,
+            args.port_service,
+            args.host_ssh,
+            args.port_ssh,
+        )
+    )

src/shadowssh/app.py

@@ -1,56 +1,64 @@
-import asyncio 
+import asyncio
+
 from app.app import Application as BaseApplication
+
 from shadowssh import log
 
+
 class Application(BaseApplication):
 
-    def __init__(self, host_forward_service=None, port_forward_service=None, host_forward_ssh=None, port_forward_ssh=None, *args,**kwargs):
+    def __init__(
+        self,
+        host_forward_service=None,
+        port_forward_service=None,
+        host_forward_ssh=None,
+        port_forward_ssh=None,
+        *args,
+        **kwargs,
+    ):
         self.host_forward_service = host_forward_service
         self.port_forward_service = port_forward_service
         self.host_forward_ssh = host_forward_ssh
         self.port_forward_ssh = port_forward_ssh
-        super().__init__(*args,**kwargs)
+        super().__init__(*args, **kwargs)
 
     async def connect(self, host, port):
-        log.info("Connected to {}:{}.".format(host,port))
+        log.info(f"Connected to {host}:{port}.")
         reader, writer = await asyncio.open_connection(host, port)
-        return reader,writer 
-    
-    async def stream(self, reader, writer,port):
+        return reader, writer
+
+    async def stream(self, reader, writer, port):
         while True:
             chunk = await reader.read(4096)
-            if not chunk: 
+            if not chunk:
                 writer.close()
                 log.info("Forward connection closed.")
-                return 
+                return
             writer.write(chunk)
             await writer.drain()
 
-    async def forward(self, reader, writer): 
-
+    async def forward(self, reader, writer):
 
         data = await reader.read(10)
         if not data:
-            return 
+            return
         port = None
         host = None
-        if b'SSH-2.0' in data:
+        if b"SSH-2.0" in data:
             port = self.port_forward_ssh
             host = self.host_forward_ssh
         else:
             port = self.port_forward_service
             host = self.host_forward_service
-        
-        log.info("Streaming to port: {}.".format(port))
+
+        log.info(f"Streaming to port: {port}.")
         upstream_reader, upstream_writer = await self.connect(host, port)
         upstream_writer.write(data)
         tasks = []
-        tasks.append(asyncio.create_task(self.stream(reader, upstream_writer,port)))
-        tasks.append(asyncio.create_task(self.stream(upstream_reader, writer,port)))
+        tasks.append(asyncio.create_task(self.stream(reader, upstream_writer, port)))
+        tasks.append(asyncio.create_task(self.stream(upstream_reader, writer, port)))
         await asyncio.gather(*tasks)
 
-
-
     async def route(self, host, port):
 
         async def handle_client(reader, writer):
@@ -58,21 +66,23 @@ class Application(BaseApplication):
 
         server = await asyncio.start_server(handle_client, host, port)
         addr = server.sockets[0].getsockname()
-        log.info(f'Serving on {addr}')
+        log.info(f"Serving on {addr}")
 
         async with server:
             await server.serve_forever()
 
-    async def serve(self, host,port,host_forward_service, port_forward_service, host_forward_ssh, port_forward_ssh):
+    async def serve(
+        self,
+        host,
+        port,
+        host_forward_service,
+        port_forward_service,
+        host_forward_ssh,
+        port_forward_ssh,
+    ):
         self.host_forward_service = host_forward_service
         self.port_forward_service = port_forward_service
         self.host_forward_ssh = host_forward_ssh
         self.port_forward_ssh = port_forward_ssh
-        tasks = [
-            self.route(host,port)
-        ]
+        tasks = [self.route(host, port)]
         await asyncio.gather(*tasks)
-
-
-
-