diff --git a/src/shadowssh/__init__.py b/src/shadowssh/__init__.py index db4920a..378ee35 100644 --- a/src/shadowssh/__init__.py +++ b/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/") diff --git a/src/shadowssh/__main__.py b/src/shadowssh/__main__.py index 69ccc2d..d38054c 100644 --- a/src/shadowssh/__main__.py +++ b/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, + ) + ) diff --git a/src/shadowssh/app.py b/src/shadowssh/app.py index 3e75e9e..bfbcc2c 100644 --- a/src/shadowssh/app.py +++ b/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) - - - -