4 files changed, 192 insertions, 5 deletions

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=(
+                    "<p>Authentication complete. You may close this window.</p>"
+                ),
+                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