auth: Refactor get_id_token()

Use token source classes to hide the details from users. There are two
sources: from file and from command.

Author: holesch

not_my_board/_agent.py

@@ -13,8 +13,6 @@
 from dataclasses import dataclass, field
 from typing import List, Tuple
 
-import not_my_board._auth as auth
-import not_my_board._http as http
 import not_my_board._jsonrpc as jsonrpc
 import not_my_board._models as models
 import not_my_board._usbip as usbip
@@ -25,18 +23,10 @@
 Address = Tuple[str, int]
 
 
-async def agent(hub_url, ca_files, token_store_path, token_cmd):
-    io = _AgentIO(hub_url, http.Client(ca_files), token_store_path, token_cmd)
-    async with Agent(hub_url, io) as agent_:
-        await agent_.serve_forever()
-
-
-class _AgentIO:
-    def __init__(self, hub_url, http_client, token_store_path, token_cmd):
+class AgentIO:
+    def __init__(self, hub_url, http_client):
         self._hub_url = hub_url
         self._http = http_client
-        self._token_store_path = token_store_path
-        self._token_cmd = token_cmd
 
     @contextlib.asynccontextmanager
     async def hub_rpc(self):
@@ -111,33 +101,19 @@ async def _handle_port_forward_client(self, proxy, target, client_r, client_w):
             await client_w.drain()
             await util.relay_streams(client_r, client_w, remote_r, remote_w)
 
-    async def get_id_token(self):
-        if self._token_cmd:
-            logger.debug("Executing token command: %s", self._token_cmd)
-            proc = await asyncio.create_subprocess_shell(
-                self._token_cmd, stdout=asyncio.subprocess.PIPE
-            )
-            stdout, _ = await proc.communicate()
-            if proc.returncode:
-                raise RuntimeError(f"{self._token_cmd!r} exited with {proc.returncode}")
-            return stdout.decode("utf-8").rstrip()
-
-        return await auth.get_id_token(
-            self._token_store_path, self._hub_url, self._http
-        )
-
 
 class Agent(util.ContextStack):
-    def __init__(self, hub_url, io):
+    def __init__(self, hub_url, io, token_src):
         url = urllib.parse.urlsplit(hub_url)
         self._hub_host = url.netloc.split(":")[0]
         self._io = io
         self._locks = weakref.WeakValueDictionary()
         self._reservations = {}
+        self._token_src = token_src
 
     async def _context_stack(self, stack):
         self._hub = await stack.enter_async_context(self._io.hub_rpc())
-        self._hub.set_api_object(_WebsocketInterface(self._io))
+        self._hub.set_api_object(self._token_src)
         stack.push_async_callback(self._cleanup)
         self._unix_server = await stack.enter_async_context(self._io.unix_server(self))
 
@@ -246,14 +222,6 @@ async def status(self):
         ]
 
 
-class _WebsocketInterface:
-    def __init__(self, io):
-        self._io = io
-
-    async def get_id_token(self):
-        return await self._io.get_id_token()
-
-
 def _filter_places(import_description, places):
     candidates = {}
 

not_my_board/_auth/__init__.py

@@ -1,2 +1,2 @@
-from ._login import LoginFlow, get_id_token
+from ._login import IdTokenFromCmd, IdTokenFromFile, LoginFlow
 from ._openid import Validator

not_my_board/_auth/_login.py

@@ -71,16 +71,37 @@ async def oidc_callback_registered(self):
         self._ready_event.set()
 
 
-async def get_id_token(token_store_path, hub_url, http_client):
-    token_store = _TokenStore(token_store_path)
-    async with token_store:
-        id_token, refresh_token = token_store.get_tokens(hub_url)
-        id_token, refresh_token = await ensure_fresh(
-            id_token, refresh_token, http_client
-        )
-        token_store.save_tokens(hub_url, id_token, refresh_token)
+class IdTokenFromFile:
+    def __init__(self, hub_url, http_client, token_store_path):
+        self._hub_url = hub_url
+        self._http = http_client
+        self._token_store = _TokenStore(token_store_path)
+
+    async def get_id_token(self):
+        async with self._token_store:
+            id_token, refresh_token = self._token_store.get_tokens(self._hub_url)
+            id_token, refresh_token = await ensure_fresh(
+                id_token, refresh_token, self._http
+            )
+            self._token_store.save_tokens(self._hub_url, id_token, refresh_token)
+
+        return id_token
 
-    return id_token
+
+class IdTokenFromCmd:
+    def __init__(self, hub_url, http_client, cmd):
+        self._hub_url = hub_url
+        self._http = http_client
+        self._cmd = cmd
+
+    async def get_id_token(self):
+        proc = await asyncio.create_subprocess_shell(
+            self._cmd, stdout=asyncio.subprocess.PIPE
+        )
+        stdout, _ = await proc.communicate()
+        if proc.returncode:
+            raise RuntimeError(f"{self._cmd!r} exited with {proc.returncode}")
+        return stdout.decode("utf-8").rstrip()
 
 
 class _TokenStore(util.ContextStack):

not_my_board/_export.py

@@ -8,7 +8,6 @@
 
 import h11
 
-import not_my_board._auth as auth
 import not_my_board._http as http
 import not_my_board._jsonrpc as jsonrpc
 import not_my_board._models as models
@@ -18,21 +17,14 @@
 logger = logging.getLogger(__name__)
 
 
-async def export(hub_url, place, ca_files, token_store_path):
-    http_client = http.Client(ca_files)
-    async with Exporter(hub_url, place, http_client, token_store_path) as exporter:
-        await exporter.register_place()
-        await exporter.serve_forever()
-
-
 class Exporter(util.ContextStack):
-    def __init__(self, hub_url, export_desc_path, http_client, token_store_path):
+    def __init__(self, hub_url, export_desc_path, http_client, token_src):
         self._hub_url = hub_url
         self._ip_to_tasks_map = {}
         export_desc_content = export_desc_path.read_text()
         self._place = models.ExportDesc(**util.toml_loads(export_desc_content))
         self._http = http_client
-        self._token_store_path = token_store_path
+        self._token_src = token_src
 
         tcp_targets = {
             f"{tcp.host}:{tcp.port}".encode()
@@ -127,9 +119,7 @@ async def _tunnel(self, client_r, client_w, target, trailing_data):
                 await util.relay_streams(client_r, client_w, remote_r, remote_w)
 
     async def get_id_token(self):
-        return await auth.get_id_token(
-            self._token_store_path, self._hub_url, self._http
-        )
+        return await self._token_src.get_id_token()
 
 
 def format_date_time(dt=None):

not_my_board/cli/__init__.py

@@ -6,12 +6,12 @@
 import pathlib
 import sys
 
+import not_my_board._agent as agent
 import not_my_board._auth as auth
 import not_my_board._client as client
+import not_my_board._export as export
 import not_my_board._http as http
 import not_my_board._util as util
-from not_my_board._agent import agent
-from not_my_board._export import export
 from not_my_board._hub import run_hub
 
 try:
@@ -160,11 +160,26 @@ def _hub_command(_):
 
 
 async def _export_command(args):
-    await export(args.hub_url, args.export_description, args.cacert, TOKEN_STORE_PATH)
+    http_client = http.Client(args.cacert)
+    token_src = auth.IdTokenFromFile(args.hub_url, http_client, TOKEN_STORE_PATH)
+    async with export.Exporter(
+        args.hub_url, args.export_description, http_client, token_src
+    ) as exporter:
+        await exporter.register_place()
+        await exporter.serve_forever()
 
 
 async def _agent_command(args):
-    await agent(args.hub_url, args.cacert, TOKEN_STORE_PATH, args.token_cmd)
+    http_client = http.Client(args.cacert)
+    io = agent.AgentIO(args.hub_url, http_client)
+
+    if args.token_cmd:
+        token_src = auth.IdTokenFromCmd(args.hub_url, http_client, args.token_cmd)
+    else:
+        token_src = auth.IdTokenFromFile(args.hub_url, http_client, TOKEN_STORE_PATH)
+
+    async with agent.Agent(args.hub_url, io, token_src) as agent_:
+        await agent_.serve_forever()
 
 
 async def _reserve_command(args):

tests/test_agent.py

@@ -212,7 +212,7 @@ async def port_forward(self, ready_event, proxy, target, local_port):
 @pytest.fixture()
 async def agent_io():
     io = FakeAgentIO()
-    async with agentmodule.Agent(HUB_URL, io) as agent:
+    async with agentmodule.Agent(HUB_URL, io, None) as agent:
         async with util.background_task(agent.serve_forever()):
             yield io
 

tests/test_auth.py

@@ -210,12 +210,11 @@ def token_store_path():
 
 
 class FakeExporter:
-    def __init__(self, rpc, token_store_path_, http, id_token=None):
+    def __init__(self, rpc, token_src, http):
         rpc.set_api_object(self)
         self._rpc = rpc
         self._http = http
-        self._token_store_path = token_store_path_
-        self._id_token = id_token
+        self._token_src = token_src
 
     async def communicate_forever(self):
         await self._rpc.communicate_forever()
@@ -228,9 +227,15 @@ async def register_place(self):
         await self._rpc.register_place(place)
 
     async def get_id_token(self):
-        if self._id_token is not None:
-            return self._id_token
-        return await auth.get_id_token(self._token_store_path, HUB_URL, self._http)
+        return await self._token_src.get_id_token()
+
+
+class FakeTokenSource:
+    def __init__(self, id_token):
+        self._id_token = id_token
+
+    async def get_id_token(self):
+        return self._id_token
 
 
 async def test_login_success(hub, http_client, token_store_path):
@@ -247,7 +252,8 @@ async def test_login_success(hub, http_client, token_store_path):
         assert claims["aud"] == CLIENT_ID
 
     rpc1, rpc2 = fake_rpc_pair()
-    fake_exporter = FakeExporter(rpc1, token_store_path, http_client)
+    token_src = auth.IdTokenFromFile(HUB_URL, http_client, token_store_path)
+    fake_exporter = FakeExporter(rpc1, token_src, http_client)
     hub_coro = hub.communicate("3.1.1.1", rpc2)
     exporter_coro = fake_exporter.communicate_forever()
     async with util.background_task(hub_coro):
@@ -299,7 +305,8 @@ async def test_permissions(claims, is_allowed):
     tokens = http_client_.issue_tokens(full_claims)
 
     rpc1, rpc2 = fake_rpc_pair()
-    fake_exporter = FakeExporter(rpc1, "", http_client_, tokens["id"])
+    token_src = FakeTokenSource(tokens["id"])
+    fake_exporter = FakeExporter(rpc1, token_src, http_client_)
     hub_coro = hub_.communicate("3.1.1.1", rpc2)
     exporter_coro = fake_exporter.communicate_forever()
     async with util.background_task(hub_coro):
@@ -327,7 +334,8 @@ async def test_permission_lost(hub, http_client, token_store_path, fake_time):
     token_store_path.write_text(json.dumps(token_store_content))
 
     rpc1, rpc2 = fake_rpc_pair()
-    fake_exporter = FakeExporter(rpc1, token_store_path, http_client)
+    token_src = auth.IdTokenFromFile(HUB_URL, http_client, token_store_path)
+    fake_exporter = FakeExporter(rpc1, token_src, http_client)
     hub_coro = hub.communicate("3.1.1.1", rpc2)
     exporter_coro = fake_exporter.communicate_forever()