Pass headers to the serverv1.6.0

2 files changed, 13 insertions, 1 deletions

diff --git a/ttun/__main__.py b/ttun/__main__.py
index 49d8b07..b14d0e7 100644
--- a/ttun/__main__.py
+++ b/ttun/__main__.py
@@ -21,7 +21,8 @@ header_regex = re.compile("(?P<header>\w+)")
 
 
 def header(v) -> Tuple[str, str]:
-    return v
+    name, *value = v.split(":")
+    return name.strip(), ":".join(value).strip()
 
 
 def main():
@@ -67,6 +68,7 @@ def main():
         server=args.server,
         to=args.to,
         https=args.https,
+        headers=args.header,
     )
 
     try:
diff --git a/ttun/client.py b/ttun/client.py
index 5434f0e..1218283 100644
--- a/ttun/client.py
+++ b/ttun/client.py
@@ -7,7 +7,9 @@ from time import perf_counter
 from typing import Awaitable
 from typing import Callable
 from typing import Coroutine
+from typing import List
 from typing import Optional
+from typing import Tuple
 from uuid import uuid4
 
 import websockets
@@ -32,6 +34,7 @@ class Client:
         subdomain: str = None,
         to: str = "127.0.0.1",
         https: bool = False,
+        headers: List[Tuple[str, str]] = None,
     ):
         self.server = server
         self.subdomain = subdomain
@@ -41,6 +44,8 @@ class Client:
 
         self.proxy_origin = f'{"https" if https else "http"}://{to}:{port}'
 
+        self.headers = [] if headers is None else headers
+
     async def send(self, data: dict):
         await self.connection.send(json.dumps(data))
 
@@ -80,6 +85,11 @@ class Client:
             while True:
                 try:
                     request: RequestData = await self.receive()
+
+                    request["headers"] = [
+                        *request["headers"],
+                        *self.headers,
+                    ]
                     await self.proxy_request(
                         session=session,
                         request=request,