From 06e71210f33d7b7701f93ff44b336b41e9e868ab Mon Sep 17 00:00:00 2001 From: Tom van der Lee Date: Thu, 30 Jul 2026 16:02:53 +0200 Subject: Auth --- ttun/__main__.py | 5 +- ttun/client.py | 145 +++++++++++++++++++++++++++++++++++++++++++++++++- ttun/settings.py | 5 +- ttun/token_storage.py | 42 +++++++++++++++ 4 files changed, 192 insertions(+), 5 deletions(-) create mode 100644 ttun/token_storage.py diff --git a/ttun/__main__.py b/ttun/__main__.py index 4b108f5..38b9d8d 100644 --- a/ttun/__main__.py +++ b/ttun/__main__.py @@ -34,9 +34,11 @@ def main(): parser.add_argument("port", help="The local port to expose") parser.add_argument( "--server", - default=f'{"wss" if SERVER_USING_SSL else "ws"}://{SERVER_HOSTNAME}', + # default=f'{"wss" if SERVER_USING_SSL else "ws"}://{SERVER_HOSTNAME}', + default=os.environ.get('TTUN_SERVER', SERVER_HOSTNAME), help="The hostname of the ttun server", ) + parser.add_argument('--server-using-ssl', default=SERVER_USING_SSL, action='store_true') parser.add_argument( "-s", "--subdomain", @@ -71,6 +73,7 @@ def main(): port=args.port, subdomains=args.subdomain, server=args.server, + using_ssl=args.server_using_ssl, to=args.to, https=args.https, headers=args.header, diff --git a/ttun/client.py b/ttun/client.py index 63bf258..1e45230 100644 --- a/ttun/client.py +++ b/ttun/client.py @@ -1,11 +1,16 @@ import asyncio +import hashlib import json import logging import os +import secrets import sys +import webbrowser +from asyncio import Future from asyncio import get_running_loop from base64 import b64decode from base64 import b64encode +from base64 import urlsafe_b64encode from datetime import datetime from time import perf_counter from typing import Awaitable @@ -14,6 +19,7 @@ from typing import Coroutine from typing import List from typing import Optional from typing import Tuple +from urllib.parse import urlencode from uuid import uuid4 import websockets @@ -21,10 +27,13 @@ from aiohttp import ClientConnectionError from aiohttp import ClientError from aiohttp import ClientSession from aiohttp import DummyCookieJar +from aiohttp import web from websockets.asyncio.client import ClientConnection from websockets.exceptions import ConnectionClosed +from websockets.exceptions import InvalidStatus from ttun import __version__ +from ttun import token_storage from ttun.pubsub import PubSub from ttun.types import Config from ttun.types import HttpMessage @@ -46,13 +55,15 @@ class Client: self, port: int, server: str, + using_ssl: bool = True, subdomains: List[str] = None, to: str = "127.0.0.1", https: bool = False, headers: List[Tuple[str, str]] = None, ): self.version = __version__ - self.server = server + self.http_origin = f'{"https://" if using_ssl else "http://"}{server}' + self.ws_origin= f'{"wss://" if using_ssl else "ws://"}{server}' self.subdomains = subdomains self._config: dict = None @@ -63,6 +74,9 @@ class Client: self.headers = [] if headers is None else headers + token = token_storage.load_token(self.http_origin) + self.token_type, self.token = token if token else (None, None) + self.websocket_connections = {} async def send(self, data: dict): @@ -72,6 +86,7 @@ class Client: data = json.loads(await self.connection.recv()) return data + @staticmethod def loop(sleep: int = None): async def wrapper(callback: Callable[[], Coroutine]): @@ -87,8 +102,134 @@ class Client: return wrapper + + async def authenticate(self) -> None: + async with ClientSession() as session: + async with session.get(f"{self.http_origin}/oauth/config/") as response: + config = await response.json() + + code_verifier = secrets.token_urlsafe(64) + code_challenge = ( + urlsafe_b64encode(hashlib.sha256(code_verifier.encode()).digest()) + .rstrip(b"=") + .decode() + ) + state = secrets.token_urlsafe(32) + + runner, port, code_future = await self._start_callback_server(state) + + try: + # redirect_uri must exactly match a redirect_uri registered on the + # server for this client_id; a loopback address will be rejected + # until the server accepts CLI-issued loopback redirect URIs + redirect_uri = f"http://127.0.0.1:{port}/callback" + authorize_url = f"{self.http_origin}/oauth/authorize/?" + urlencode( + { + "client_id": config["client_id"], + "redirect_uri": redirect_uri, + "response_type": "code", + "state": state, + "code_challenge": code_challenge, + "code_challenge_method": "S256", + } + ) + + print(f"Opening browser for authentication:\n{authorize_url}") + webbrowser.open(authorize_url) + + try: + code = await asyncio.wait_for(code_future, timeout=300) + except asyncio.TimeoutError: + print("Authentication timed out. Please try again.") + raise + except RuntimeError as e: + print(f"Authentication failed: {e}") + raise + finally: + await runner.cleanup() + + await self._exchange_code(code, redirect_uri, config["client_id"], code_verifier) + + async def _start_callback_server( + self, expected_state: str + ) -> Tuple[web.AppRunner, int, "Future[str]"]: + code_future: Future[str] = get_running_loop().create_future() + + async def callback(request: web.Request) -> web.Response: + if not code_future.done(): + self._resolve_callback(request, expected_state, code_future) + return web.Response( + text=( + "

Authentication complete. You may close this window.

" + ), + content_type="text/html", + ) + + app = web.Application() + app.add_routes([web.get("/callback", callback)]) + + runner = web.AppRunner(app) + await runner.setup() + site = web.TCPSite(runner, "127.0.0.1", 0) + await site.start() + port = runner.addresses[0][1] + + return runner, port, code_future + + @staticmethod + def _resolve_callback( + request: web.Request, expected_state: str, code_future: "Future[str]" + ) -> None: + if request.query.get("state") != expected_state: + code_future.set_exception(RuntimeError("OAuth state mismatch")) + elif "error" in request.query: + code_future.set_exception( + RuntimeError(f"Authorization denied: {request.query['error']}") + ) + elif "code" in request.query: + code_future.set_result(request.query["code"]) + else: + code_future.set_exception(RuntimeError("Callback missing code and error")) + + async def _exchange_code( + self, code: str, redirect_uri: str, client_id: str, code_verifier: str + ) -> None: + async with ClientSession() as session: + async with session.post( + f"{self.http_origin}/oauth/token/", + data={ + "grant_type": "authorization_code", + "code": code, + "redirect_uri": redirect_uri, + "client_id": client_id, + "code_verifier": code_verifier, + }, + ) as response: + token = await response.json() + + self.token_type = token["token_type"] + self.token = token["access_token"] + token_storage.save_token(self.http_origin, self.token_type, self.token) + + @staticmethod + def handle_authentication(func): + async def wrapper(self, *args, **kwargs): + try: + await func(self, *args, **kwargs) + except InvalidStatus: + await self.authenticate() + await func(self, *args, **kwargs) + return wrapper + + + @handle_authentication async def connect(self) -> ClientConnection: - self.connection = await websockets.connect(f"{self.server}/tunnel/") + additional_headers = ( + [("Authorization", f"{self.token_type} {self.token}")] if self.token else [] + ) + self.connection = await websockets.connect( + f"{self.ws_origin}/tunnel/", additional_headers=additional_headers + ) await self.send( { diff --git a/ttun/settings.py b/ttun/settings.py index 9c15f15..3db37d2 100644 --- a/ttun/settings.py +++ b/ttun/settings.py @@ -1,3 +1,4 @@ +import os import sys from configparser import ConfigParser from os import mkdir @@ -23,8 +24,8 @@ else: with open(config_file, "w") as f: config.write(f) -SERVER_HOSTNAME = config["server"].get("hostname") -SERVER_USING_SSL = config["server"].get("using_ssl").lower() in ["yes", "true"] +SERVER_HOSTNAME = os.environ.get('TTUN_SERVER_HOSTNAME', config["server"].get("hostname")) +SERVER_USING_SSL = os.environ.get('TTUN_SERVER_USING_SSL', config["server"].get("using_ssl").lower()) in ["yes", "true", "True", "1"] try: assert SERVER_HOSTNAME != "" diff --git a/ttun/token_storage.py b/ttun/token_storage.py new file mode 100644 index 0000000..cf5a9eb --- /dev/null +++ b/ttun/token_storage.py @@ -0,0 +1,42 @@ +import configparser +from pathlib import Path +from typing import Optional +from typing import Tuple + +import appdirs + +TOKEN_FILE = ( + Path(appdirs.user_config_dir("ttun-client", "ttun", roaming=True)) + / "authentication.ini" +) + + +def load_token(server: str) -> Optional[Tuple[str, str]]: + if not TOKEN_FILE.exists(): + return None + + config = configparser.ConfigParser() + config.read(TOKEN_FILE) + + if server not in config: + return None + + token_type = config[server].get("token_type") + access_token = config[server].get("access_token") + + if not token_type or not access_token: + return None + + return token_type, access_token + + +def save_token(server: str, token_type: str, access_token: str) -> None: + config = configparser.ConfigParser() + if TOKEN_FILE.exists(): + config.read(TOKEN_FILE) + + config[server] = {"token_type": token_type, "access_token": access_token} + + TOKEN_FILE.parent.mkdir(parents=True, exist_ok=True) + with open(TOKEN_FILE, "w") as f: + config.write(f) -- cgit v1.2.3