summaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
-rw-r--r--ttun/__main__.py4
-rw-r--r--ttun/client.py10
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+)")
21 21
22 22
23def header(v) -> Tuple[str, str]: 23def header(v) -> Tuple[str, str]:
24 return v 24 name, *value = v.split(":")
25 return name.strip(), ":".join(value).strip()
25 26
26 27
27def main(): 28def main():
@@ -67,6 +68,7 @@ def main():
67 server=args.server, 68 server=args.server,
68 to=args.to, 69 to=args.to,
69 https=args.https, 70 https=args.https,
71 headers=args.header,
70 ) 72 )
71 73
72 try: 74 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
7from typing import Awaitable 7from typing import Awaitable
8from typing import Callable 8from typing import Callable
9from typing import Coroutine 9from typing import Coroutine
10from typing import List
10from typing import Optional 11from typing import Optional
12from typing import Tuple
11from uuid import uuid4 13from uuid import uuid4
12 14
13import websockets 15import websockets
@@ -32,6 +34,7 @@ class Client:
32 subdomain: str = None, 34 subdomain: str = None,
33 to: str = "127.0.0.1", 35 to: str = "127.0.0.1",
34 https: bool = False, 36 https: bool = False,
37 headers: List[Tuple[str, str]] = None,
35 ): 38 ):
36 self.server = server 39 self.server = server
37 self.subdomain = subdomain 40 self.subdomain = subdomain
@@ -41,6 +44,8 @@ class Client:
41 44
42 self.proxy_origin = f'{"https" if https else "http"}://{to}:{port}' 45 self.proxy_origin = f'{"https" if https else "http"}://{to}:{port}'
43 46
47 self.headers = [] if headers is None else headers
48
44 async def send(self, data: dict): 49 async def send(self, data: dict):
45 await self.connection.send(json.dumps(data)) 50 await self.connection.send(json.dumps(data))
46 51
@@ -80,6 +85,11 @@ class Client:
80 while True: 85 while True:
81 try: 86 try:
82 request: RequestData = await self.receive() 87 request: RequestData = await self.receive()
88
89 request["headers"] = [
90 *request["headers"],
91 *self.headers,
92 ]
83 await self.proxy_request( 93 await self.proxy_request(
84 session=session, 94 session=session,
85 request=request, 95 request=request,