Add token validation

dev-tools/entrypoint-server.py

@@ -1,13 +1,22 @@
 import os
 
+import importlib_metadata
+
 os.chdir(os.path.dirname(os.path.abspath(__file__)))
 os.environ["DUETECTOR_LOG_LEVEL"] = "DEBUG"
 
 import re
 import sys
 from pathlib import Path
 
-from pkg_resources import load_entry_point
+
+def load_entry_point(distribution, group, name):
+    dist_obj = importlib_metadata.distribution(distribution)
+    eps = [ep for ep in dist_obj.entry_points if ep.group == group and ep.name == name]
+    if not eps:
+        raise ImportError("Entry point %r not found" % ((group, name),))
+    return eps[0].load()
+
 
 db_file = Path("./duetector-dbcollector.sqlite3")
 config_file = Path("./config.toml")

dev-tools/entrypoint.py

@@ -1,13 +1,22 @@
 import os
 
+import importlib_metadata
+
 os.chdir(os.path.dirname(os.path.abspath(__file__)))
 os.environ["DUETECTOR_LOG_LEVEL"] = "DEBUG"
 
 import re
 import sys
 from pathlib import Path
 
-from pkg_resources import load_entry_point
+
+def load_entry_point(distribution, group, name):
+    dist_obj = importlib_metadata.distribution(distribution)
+    eps = [ep for ep in dist_obj.entry_points if ep.group == group and ep.name == name]
+    if not eps:
+        raise ImportError("Entry point %r not found" % ((group, name),))
+    return eps[0].load()
+
 
 db_file = Path("./duetector-dbcollector.sqlite3")
 config_file = Path("./config.toml")

duetector/service/app.py

@@ -1,13 +1,26 @@
-from fastapi import FastAPI
+from typing import Optional
+
+from fastapi import Depends, FastAPI, HTTPException, Query
 
 from duetector.__init__ import __version__
+from duetector.service.config import Config, get_server_config
 from duetector.service.control.routes import r as cr
 from duetector.service.query.routes import r as qr
 
+
+async def verify_token(
+    server_config: Config = Depends(get_server_config),
+    token: Optional[str] = Query(default=""),
+):
+    if server_config.token and token != server_config.token:
+        raise HTTPException(403, "Invalid token")
+
+
 app = FastAPI(
     title="Duetector",
     description="Data Usage Extensible Detector for data usage observability",
     version=__version__,
+    dependencies=[Depends(verify_token)],
 )
 app.include_router(qr)
 app.include_router(cr)

duetector/service/base.py

@@ -0,0 +1,22 @@
+from typing import Any, Dict, Optional
+
+from fastapi import Depends
+
+from duetector.config import Configuable
+from duetector.service.config import get_config
+
+
+class Controller(Configuable):
+    config_scope = None
+
+    default_config = {}
+
+    def __init__(self, config: Optional[Dict[str, Any]] = None, *args, **kwargs):
+        super().__init__(config, *args, **kwargs)
+
+
+def get_controller(controller_type: type):
+    def _(config: dict = Depends(get_config)) -> Controller:
+        return controller_type(config)
+
+    return _

duetector/service/config.py

@@ -1,12 +1,14 @@
 import os
 from typing import Any, Dict
 
+from fastapi import Depends
+
 try:
     from functools import cache
 except ImportError:
     from functools import lru_cache as cache
 
-from duetector.config import ConfigLoader
+from duetector.config import Config, ConfigLoader, Configuable
 
 CONFIG_PATH_ENV = "DUETECTOR_SERVER_CONFIG_PATH"
 
@@ -24,3 +26,14 @@ def get_config() -> Dict[str, Any]:
         load_env=False,
         dump_when_load=False,
     ).load_config()
+
+
+class ServerConfig(Configuable):
+    config_scope = "server"
+    default_config = {"token": ""}
+
+
+def get_server_config(
+    config: Dict[str, Any] = Depends(get_config),
+) -> Config:
+    return ServerConfig(config).config

duetector/service/control/controller.py

@@ -0,0 +1,5 @@
+from duetector.service.base import Controller
+
+
+class DaemonControler(Controller):
+    pass

duetector/service/control/routes.py

@@ -9,4 +9,4 @@
 
 @r.get("/")
 async def root(config: dict = Depends(get_config)):
-    return config
+    pass

duetector/service/query/controller.py

@@ -0,0 +1,17 @@
+from typing import Any, Dict, Optional
+
+from duetector.analyzer.base import Analyzer
+from duetector.analyzer.db import DBAnalyzer
+from duetector.service.base import Controller
+
+
+class AnalyzerController(Controller):
+    def __init__(self, config: Optional[Dict[str, Any]] = None, *args, **kwargs):
+        super().__init__(config, *args, **kwargs)
+
+        # TODO: Make this configurable
+        self.analyzer: Analyzer = self._init_analyzer(DBAnalyzer)
+
+    def _init_analyzer(self, analyzer: type):
+        analyzer_config = getattr(self.config, analyzer.config_scope)._config_dict
+        return analyzer(analyzer_config)

duetector/service/query/routes.py

@@ -1,6 +1,8 @@
 from fastapi import APIRouter, Depends
 
+from duetector.service.base import get_controller
 from duetector.service.config import get_config
+from duetector.service.query.controller import AnalyzerController
 
 r = APIRouter(
     prefix="/query",
@@ -10,3 +12,8 @@
 @r.get("/")
 async def root(config: dict = Depends(get_config)):
     return config
+
+
+@r.get("/who")
+async def who(controller: AnalyzerController = Depends(get_controller(AnalyzerController))):
+    return str(controller.analyzer)

duetector/static/config.toml

@@ -104,3 +104,6 @@ table_prefix = "duetector_tracking"
 
 [db_analyzer.db.engine]
 url = "sqlite:///duetector-dbcollector.sqlite3"
+
+[server]
+token = ""

duetector/tools/config_generator.py

@@ -9,6 +9,7 @@
 from duetector.log import logger
 from duetector.managers import CollectorManager, FilterManager, TracerManager
 from duetector.monitors import BccMonitor, ShMonitor
+from duetector.service.config import ServerConfig
 
 
 def _recursive_load(config_scope: str, config_dict: dict, default_config: dict):
@@ -62,6 +63,8 @@ class ConfigGenerator:
     All analyzers to inspect.
     """
 
+    others = [ServerConfig]
+
     def __init__(
         self,
         load: bool = True,
@@ -87,6 +90,9 @@ def __init__(
 
         for a in self.analyzer:
             _recursive_load(a.config_scope, self.dynamic_config, a.default_config)
+
+        for o in self.others:
+            _recursive_load(o.config_scope, self.dynamic_config, o.default_config)
         # This will generate default config file if not exists
         if load:
             self.loaded_config = ConfigLoader(path, load_env, dump_when_load=False).load_config()