Merge pull request #159 from akquinet/token-map

fix(models): use a map to lookup token for environment

Author: rwxd

powerdns_api_proxy/config.py

@@ -33,18 +33,19 @@ def load_config(path: Optional[Path] = None) -> ProxyConfig:
     with open(path) as f:
         data = safe_load(f)
 
-    config = ProxyConfig.model_validate(data)
-    return config
+    return ProxyConfig(**data)
 
 
 def token_defined(config: ProxyConfig, token: str) -> bool:
     sha512 = hashlib.sha512()
     sha512.update(token.encode())
     token_digest = sha512.digest().hex()
-    for env in config.environments:
-        if token_digest == env.token_sha512:
-            logger.info(f'Authenticated environment "{env.name}"')
-            return True
+
+    if token_digest in config.token_env_map:
+        logger.info(
+            f'Authenticated environment "{config.token_env_map[token_digest].name}"'
+        )
+        return True
     return False
 
 
@@ -92,9 +93,10 @@ def get_environment_for_token(
     sha512 = hashlib.sha512()
     sha512.update(token.encode())
     token_digest = sha512.digest().hex()
-    for env in config.environments:
-        if token_digest == env.token_sha512:
-            return env
+
+    if token_digest in config.token_env_map:
+        return config.token_env_map[token_digest]
+
     raise ValueError("Could not find a environment for the given token")
 
 

powerdns_api_proxy/models.py

@@ -163,6 +163,8 @@ class ProxyConfig(BaseModel):
         </body>
     </html>
 """
+    # Dictionary for fast token lookups
+    token_env_map: dict[str, ProxyConfigEnvironment] = {}
 
     @field_validator("pdns_api_url")
     @classmethod
@@ -178,6 +180,13 @@ def api_token_defined(cls, v):
             raise ValueError("pdns_api_token must a non-empty string")
         return v
 
+    def __init__(self, **data):
+        super().__init__(**data)
+
+        # Automatically populate token_env_map during initialization
+        for env in self.environments:
+            self.token_env_map[env.token_sha512] = env
+
 
 class ResponseAllowed(BaseModel):
     zones: list[ProxyConfigZone]