From ceb6564185b5fe662d40db3ea04a960ddbc43a84 Mon Sep 17 00:00:00 2001 From: Tom van der Lee Date: Fri, 31 Dec 2021 16:56:24 +0100 Subject: Version 0.1.0 --- pyproject.toml | 3 ++ requirements.txt | 4 ++- setup.cfg | 13 +++++++ ttun/__main__.py | 102 +++++++++++++------------------------------------------ ttun/client.py | 76 +++++++++++++++++++++++++++++++++++++++++ ttun/types.py | 24 +++++++++++++ 6 files changed, 142 insertions(+), 80 deletions(-) create mode 100644 pyproject.toml create mode 100644 setup.cfg create mode 100644 ttun/client.py create mode 100644 ttun/types.py diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..9787c3b --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,3 @@ +[build-system] +requires = ["setuptools", "wheel"] +build-backend = "setuptools.build_meta" diff --git a/requirements.txt b/requirements.txt index 86290f8..4c59602 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1 +1,3 @@ -websockets ~= 10.0 +. +setuptools +build diff --git a/setup.cfg b/setup.cfg new file mode 100644 index 0000000..e6fa217 --- /dev/null +++ b/setup.cfg @@ -0,0 +1,13 @@ +[metadata] +name = ttun +version = 0.1.0 + +[options] +packages = ttun +install_requires = + websockets ~= 10.0 + aiohttp ~= 3.8 + +[options.entry_points] +console_scripts = + ttun = ttun.__main__:main diff --git a/ttun/__main__.py b/ttun/__main__.py index 5984286..f8a1978 100644 --- a/ttun/__main__.py +++ b/ttun/__main__.py @@ -1,88 +1,32 @@ import asyncio -import json -from dataclasses import dataclass -from functools import wraps -from typing import Coroutine, Callable, TypedDict +from argparse import ArgumentParser -import websockets -from websockets import WebSocketClientProtocol -from websockets.exceptions import ConnectionClosed +from ttun.client import Client -class Message(TypedDict): - type: str +def main(): + parser = ArgumentParser(prog='ttun') + parser.add_argument('port') + parser.add_argument('-s', '--subdomain', default=None) + args = parser.parse_args() + client = Client(port=args.port, subdomain=args.subdomain) -class Request(TypedDict): - path: str + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + loop.run_until_complete(client.connect()) -class Reply(TypedDict): - path: str - - -@dataclass -class Client: - connection = WebSocketClientProtocol - - async def send(self, data: dict): - await self.connection.send(json.dumps(data)) - - async def receive(self) -> dict: - return json.loads(await self.connection.recv()) - - @staticmethod - def loop(sleep: int = None): - async def wrapper(callback: Callable[[], Coroutine]): - while True: - try: - await callback() - - if sleep is not None: - await asyncio.sleep(sleep) - - except ConnectionClosed: - break - return wrapper - - async def connect(self) -> WebSocketClientProtocol: - self.connection = await websockets.connect('ws://localhost:8000/tunnel/') - - await self.send({ - 'type': 'config', - 'config': { - 'subdomain': 'test' - } - }) - - config = await self.receive() - print(config) - - if self.connection.open: - return self.connection - - async def handle_messages(self): - while True: - try: - message: Request = await self.receive() - print(message) - await self.connection.send(json.dumps({ - 'type': 'response', - 'response': { - 'test': 'Test' - } - })) - - except ConnectionClosed: - break - -client = Client() -loop = asyncio.get_event_loop() - -connection = loop.run_until_complete(client.connect()) - -loop.run_until_complete(asyncio.wait([ - asyncio.ensure_future(client.handle_messages()), - # asyncio.ensure_future(client.heartbeat()) -])) + print(client.config['url']) + try: + loop.run_until_complete(asyncio.wait([ + asyncio.ensure_future(client.handle_messages()), + ])) + except KeyboardInterrupt: + pass +if __name__ == '__main__': + main() diff --git a/ttun/client.py b/ttun/client.py new file mode 100644 index 0000000..defe023 --- /dev/null +++ b/ttun/client.py @@ -0,0 +1,76 @@ +import asyncio +import json +from base64 import b64encode +from dataclasses import dataclass +from typing import Optional, Callable, Coroutine + +import websockets +from aiohttp import ClientSession +from websockets import WebSocketClientProtocol +from websockets.exceptions import ConnectionClosed + +from ttun.types import Config, RequestData, ResponseData + + +@dataclass +class Client: + port: int + subdomain: Optional[str] + + url: Optional[Config] = None + connection = WebSocketClientProtocol + + async def send(self, data: dict): + await self.connection.send(json.dumps(data)) + + async def receive(self) -> dict: + return json.loads(await self.connection.recv()) + + @staticmethod + def loop(sleep: int = None): + async def wrapper(callback: Callable[[], Coroutine]): + while True: + try: + await callback() + + if sleep is not None: + await asyncio.sleep(sleep) + + except ConnectionClosed: + break + return wrapper + + async def connect(self) -> WebSocketClientProtocol: + self.connection = await websockets.connect('ws://localhost:8000/tunnel/') + + await self.send({ + 'subdomain': self.subdomain + }) + + self.config = await self.receive() + + if self.connection.open: + return self.connection + + async def handle_messages(self): + while True: + try: + request: RequestData = await self.receive() + + async with ClientSession() as session: + response = await session.request( + method=request['method'], + url=f'http://localhost:{self.port}{request["path"]}', + headers=request['headers'], + cookies=request['cookies'], + ) + + await self.send(ResponseData( + status=response.status, + headers=dict(response.headers), + cookies=dict(response.cookies), + body=b64encode(await response.read()).decode() + )) + + except ConnectionClosed: + break diff --git a/ttun/types.py b/ttun/types.py new file mode 100644 index 0000000..cf94b9e --- /dev/null +++ b/ttun/types.py @@ -0,0 +1,24 @@ +from typing import TypedDict, Optional + + +class Message(TypedDict): + type: str + + +class RequestData(TypedDict): + method: str + path: str + headers: dict + cookies: dict + body: Optional[str] + + +class ResponseData(TypedDict): + status: int + headers: dict + cookies: dict + body: Optional[str] + + +class Config(TypedDict): + url: str -- cgit v1.2.3