2 files changed, 173 insertions, 10 deletions

diff --git a/.gitignore b/.gitignore
index 4906f4c..a6d8f21 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,6 +1,98 @@
 
-# Created by https://www.toptal.com/developers/gitignore/api/python
-# Edit at https://www.toptal.com/developers/gitignore?templates=python
+# Created by https://www.toptal.com/developers/gitignore/api/python,pycharm+all
+# Edit at https://www.toptal.com/developers/gitignore?templates=python,pycharm+all
+
+### PyCharm+all ###
+# Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio, WebStorm and Rider
+# Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839
+
+# User-specific stuff
+.idea/**/workspace.xml
+.idea/**/tasks.xml
+.idea/**/usage.statistics.xml
+.idea/**/dictionaries
+.idea/**/shelf
+
+# AWS User-specific
+.idea/**/aws.xml
+
+# Generated files
+.idea/**/contentModel.xml
+
+# Sensitive or high-churn files
+.idea/**/dataSources/
+.idea/**/dataSources.ids
+.idea/**/dataSources.local.xml
+.idea/**/sqlDataSources.xml
+.idea/**/dynamic.xml
+.idea/**/uiDesigner.xml
+.idea/**/dbnavigator.xml
+
+# Gradle
+.idea/**/gradle.xml
+.idea/**/libraries
+
+# Gradle and Maven with auto-import
+# When using Gradle or Maven with auto-import, you should exclude module files,
+# since they will be recreated, and may cause churn.  Uncomment if using
+# auto-import.
+# .idea/artifacts
+# .idea/compiler.xml
+# .idea/jarRepositories.xml
+# .idea/modules.xml
+# .idea/*.iml
+# .idea/modules
+# *.iml
+# *.ipr
+
+# CMake
+cmake-build-*/
+
+# Mongo Explorer plugin
+.idea/**/mongoSettings.xml
+
+# File-based project format
+*.iws
+
+# IntelliJ
+out/
+
+# mpeltonen/sbt-idea plugin
+.idea_modules/
+
+# JIRA plugin
+atlassian-ide-plugin.xml
+
+# Cursive Clojure plugin
+.idea/replstate.xml
+
+# Crashlytics plugin (for Android Studio and IntelliJ)
+com_crashlytics_export_strings.xml
+crashlytics.properties
+crashlytics-build.properties
+fabric.properties
+
+# Editor-based Rest Client
+.idea/httpRequests
+
+# Android studio 3.1+ serialized cache file
+.idea/caches/build_file_checksums.ser
+
+### PyCharm+all Patch ###
+# Ignores the whole .idea folder and all .iml files
+# See https://github.com/joeblau/gitignore.io/issues/186 and https://github.com/joeblau/gitignore.io/issues/360
+
+.idea/
+
+# Reason: https://github.com/joeblau/gitignore.io/issues/186#issuecomment-249601023
+
+*.iml
+modules.xml
+.idea/misc.xml
+*.ipr
+
+# Sonarlint plugin
+.idea/sonarlint
 
 ### Python ###
 # Byte-compiled / optimized / DLL files
@@ -142,4 +234,4 @@ dmypy.json
 # Cython debug symbols
 cython_debug/
 
-# End of https://www.toptal.com/developers/gitignore/api/python
+# End of https://www.toptal.com/developers/gitignore/api/python,pycharm+all
diff --git a/ttun/__main__.py b/ttun/__main__.py
index 36e074e..5984286 100644
--- a/ttun/__main__.py
+++ b/ttun/__main__.py
@@ -1,17 +1,88 @@
 import asyncio
 import json
+from dataclasses import dataclass
+from functools import wraps
+from typing import Coroutine, Callable, TypedDict
 
 import websockets
 from websockets import WebSocketClientProtocol
+from websockets.exceptions import ConnectionClosed
 
-async def test():
-    async with websockets.connect('ws://localhost:8000/tunnel/') as websocket:
-        ws: WebSocketClientProtocol = websocket
 
-        await ws.send(json.dumps({ 'Lol':  'Test' }))
+class Message(TypedDict):
+    type: str
 
+
+class Request(TypedDict):
+    path: str
+
+
+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:
-            data = await ws.recv()
-            print(data)
+            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())
+]))
 
-asyncio.run(test())