6 files changed, 142 insertions, 80 deletions

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