Formatting.

Browse Source
This commit is contained in:
retoor 2024-12-08 10:34:11 +01:00
parent 0dffc6bdd8
commit 415c16b4ef
3 changed files with 61 additions and 46 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
src/shadowssh/__init__.py
View File

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

22
src/shadowssh/__main__.py
View File

@ -1,6 +1,8 @@
from shadowssh.app import Application
import asyncio
import argparse import argparse
import asyncio
from shadowssh.app import Application
def parse_args(): def parse_args():
parser = argparse.ArgumentParser(description="Shadowssh") parser = argparse.ArgumentParser(description="Shadowssh")
@ -17,9 +19,13 @@ def parse_args():
def serve(): def serve():
args = parse_args() args = parse_args()
app = Application() 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,
)
)

42
src/shadowssh/app.py
View File

@ -1,10 +1,21 @@
import asyncio import asyncio
from app.app import Application as BaseApplication from app.app import Application as BaseApplication
from shadowssh import log from shadowssh import log
class Application(BaseApplication): 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.host_forward_service = host_forward_service
self.port_forward_service = port_forward_service self.port_forward_service = port_forward_service
self.host_forward_ssh = host_forward_ssh self.host_forward_ssh = host_forward_ssh
@ -12,7 +23,7 @@ class Application(BaseApplication):
super().__init__(*args, **kwargs) super().__init__(*args, **kwargs)
async def connect(self, host, port): 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) reader, writer = await asyncio.open_connection(host, port)
return reader, writer return reader, writer
@ -28,20 +39,19 @@ class Application(BaseApplication):
async def forward(self, reader, writer): async def forward(self, reader, writer):
data = await reader.read(10) data = await reader.read(10)
if not data: if not data:
return return
port = None port = None
host = None host = None
if b'SSH-2.0' in data: if b"SSH-2.0" in data:
port = self.port_forward_ssh port = self.port_forward_ssh
host = self.host_forward_ssh host = self.host_forward_ssh
else: else:
port = self.port_forward_service port = self.port_forward_service
host = self.host_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_reader, upstream_writer = await self.connect(host, port)
upstream_writer.write(data) upstream_writer.write(data)
tasks = [] tasks = []
@ -49,8 +59,6 @@ class Application(BaseApplication):
tasks.append(asyncio.create_task(self.stream(upstream_reader, writer, port))) tasks.append(asyncio.create_task(self.stream(upstream_reader, writer, port)))
await asyncio.gather(*tasks) await asyncio.gather(*tasks)
async def route(self, host, port): async def route(self, host, port):
async def handle_client(reader, writer): async def handle_client(reader, writer):
@ -58,21 +66,23 @@ class Application(BaseApplication):
server = await asyncio.start_server(handle_client, host, port) server = await asyncio.start_server(handle_client, host, port)
addr = server.sockets[0].getsockname() addr = server.sockets[0].getsockname()
log.info(f'Serving on {addr}') log.info(f"Serving on {addr}")
async with server: async with server:
await server.serve_forever() 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.host_forward_service = host_forward_service
self.port_forward_service = port_forward_service self.port_forward_service = port_forward_service
self.host_forward_ssh = host_forward_ssh self.host_forward_ssh = host_forward_ssh
self.port_forward_ssh = port_forward_ssh self.port_forward_ssh = port_forward_ssh
tasks = [ tasks = [self.route(host, port)]
self.route(host,port)
]
await asyncio.gather(*tasks) await asyncio.gather(*tasks)