Merge branch 'feat_ta_support' of https://github.com/GT-BAITA/iam-proxy-italia into GT-BAITA-feat_ta_support

Author: peppelinux

iam-proxy-italia-project/backends/cieoidc/cieoidc.py

@@ -1,19 +1,82 @@
 import logging
 import inspect
-
+import time
+from datetime import datetime, timezone
+from types import SimpleNamespace
+from typing import Optional, List
 
 from satosa.backends.base import BackendModule
 from satosa.backends.oauth import get_metadata_desc_for_oauth_backend
 
 from .utils.endpoints_loader import EndpointsLoader
+from .storage.db_engine import OidcDbEngine
+from .models.trust_chain_cache import TrustChainCache
 
 from pyeudiw.federation.trust_chain_builder import TrustChainBuilder
 from pyeudiw.federation.statements import EntityStatement, get_entity_configurations
 
+from .utils.exceptions import TrustChainNotFoundError
+
 
 logger = logging.getLogger(__name__)
 
 
+def _trust_chain_from_cache(cached: TrustChainCache):
+    """
+    Build a minimal trust-chain-like object from TrustChainCache.
+    Has .subject and .subject_configuration.payload as required by authorization endpoint.
+    """
+    wrapper = SimpleNamespace()
+    wrapper.subject = cached.provider_url
+    wrapper.subject_configuration = SimpleNamespace(payload=cached.payload)
+    return wrapper
+
+
+def _is_cache_expired(cached: TrustChainCache, now=None) -> bool:
+    """Return True if the cached payload is expired (exp in the past)."""
+    exp = cached.exp or cached.payload.get("exp")
+    if exp is None:
+        return False
+    t = now if now is not None else time.time()
+    return t >= exp
+
+
+class TrustChainResolver:
+    """
+    Resolves trust chains from cache or builds them on-demand via discovery.
+    When a provider is requested but not in the cache (e.g. startup failed),
+    discovery is performed and the resulting trust chain is stored for reuse.
+    """
+
+    def __init__(self, trust_chains: dict, build_callback):
+        """
+        :param trust_chains: Dict of provider_url -> TrustChainBuilder (mutated when new chains are built)
+        :param build_callback: Callable(provider_url) -> TrustChainBuilder; raises TrustChainNotFoundError on failure
+        """
+        self._chains = trust_chains
+        self._build = build_callback
+
+    def __contains__(self, key):
+        return key in self._chains
+
+    def __getitem__(self, key):
+        return self._chains[key]
+
+    def keys(self):
+        return self._chains.keys()
+
+    def get_or_build(self, provider_url: str) -> TrustChainBuilder:
+        """Get trust chain from cache, or discover and store it on-demand."""
+        for key in (
+            provider_url,
+            provider_url.rstrip("/"),
+            provider_url + "/" if not provider_url.endswith("/") else None,
+        ):
+            if key and key in self._chains:
+                return self._chains[key]
+        return self._build(provider_url)
+
+
 class CieOidcBackend(BackendModule):
 
     def __init__(self, callback, internal_attributes, module_config, base_url, name):
@@ -23,7 +86,12 @@ def __init__(self, callback, internal_attributes, module_config, base_url, name)
         super().__init__(callback, internal_attributes, base_url, name)
         self.config = module_config
         self.endpoints = {}
+        self._validated_trust_anchors: List[EntityStatement] = []
         self.trust_chain = self._generate_trust_chains()
+        self._trust_chain_resolver = TrustChainResolver(
+            self.trust_chain,
+            self.get_or_build_trust_chain,
+        )
         metadata = self.config.get("metadata", {}).get("openid_relying_party", {})
         self._client_id = metadata.get("client_id") or f"{base_url}/{name}"
 
@@ -58,7 +126,7 @@ def register_endpoints(self):
             self.name,
             self.auth_callback_func,
             self.converter,
-            self.trust_chain,
+            self._trust_chain_resolver,
         )
 
         url_map = []
@@ -81,42 +149,116 @@ def get_metadata_desc(self):
         meta = get_metadata_desc_for_oauth_backend(self._client_id, self.config)
         return meta
 
-    def _generate_trust_chains(self) -> dict:
-        '''
-        private method _generate_trust_chains:
-        This method generate a list of trust-chain. After create a entity statement
-        for Trust Anchor, validate itself, and call the generate_trust_chain
-        for all providers into configuration.
-        Add all providers into dictionary.
-        '''
-        logger.debug(
-            f"Entering method: {inspect.getframeinfo(inspect.currentframe()).function}. "
-        )
+    def _get_storage(self) -> Optional[OidcDbEngine]:
+        """Create and return storage engine; connect if needed. Returns None if no storage configured."""
+        if getattr(self, "_storage_engine", None) is not None:
+            return self._storage_engine
+        storage_config = self.config.get("storage") or {}
+        if not storage_config:
+            return None
+        try:
+            engine = OidcDbEngine(storage_config)
+            engine.connect()
+            self._storage_engine = engine
+            return engine
+        except Exception as e:
+            logger.warning("Could not initialize storage for trust chain persistence: %s", e)
+            return None
+
+    def _store_trust_chain(self, chain, provider_url: str) -> None:
+        """Persist trust chain to database if storage is available."""
+        engine = self._get_storage()
+        if engine is None:
+            return
+        try:
+            payload = chain.subject_configuration.payload
+            exp = payload.get("exp")
+            variants = {
+                provider_url.rstrip("/"),
+                provider_url.rstrip("/") + "/"
+            }
+
+            for url in variants:
+                cached = TrustChainCache(
+                    provider_url=url,
+                    payload=payload,
+                    exp=exp,
+                    created=datetime.now(timezone.utc),
+                )
+                engine.add_or_update_trust_chain(cached)
+        except Exception as e:
+            logger.warning("Could not persist trust chain for %s: %s", provider_url, e)
 
+    def _generate_trust_chains(self) -> dict:
+        """try load from DB, or can try discovery with TA's list."""
         httpc_params = self.config["trust_chain"]["config"]["httpc_params"]
-
-        ta_url = self.config["trust_chain"]["config"]["trust_anchor"][0]
-        jwt = get_entity_configurations(ta_url, httpc_params=httpc_params)[0]
-
-        trust_anchor_ec = EntityStatement(jwt, httpc_params=httpc_params)
-
-        trust_anchor_ec.validate_by_itself()
-
         providers = self.config["providers"]
-
         trust_chains = dict()
 
         for provider_url in providers:
+            # try load from DB
+            engine = self._get_storage()
+            if engine:
+                cached = engine.get_trust_chain_by_provider(provider_url)
+                if cached and not _is_cache_expired(cached):
+                    chain = _trust_chain_from_cache(cached)
+                    self._add_to_dict(trust_chains, provider_url, chain)
+                    continue
+
+            # Build via discovery, tryng each TA
             try:
-                trust_chains[provider_url] = CieOidcBackend.generate_trust_chain(
-                    trust_anchor_ec, provider_url, httpc_params)
-            except Exception as exception:
+                tas = self._ensure_trust_anchors()
+                chain_built = False
+                for ta_ec in tas:
+                    try:
+                        chain = self.generate_trust_chain(
+                            ta_ec, provider_url, httpc_params
+                        )
+                        self._add_to_dict(trust_chains, provider_url, chain)
+                        self._store_trust_chain(chain, provider_url)
+                        logger.info(
+                            "Provider %s linked to TA %s", provider_url, ta_ec.sub
+                        )
+                        chain_built = True
+                        break
+                    except Exception as e:
+                        logger.warning(
+                            "Failed to build trust chain for provider %s with TA %s: %s",
+                            provider_url,
+                            getattr(ta_ec, "sub", "<unknown>"),
+                            e,
+                        )
+                if not chain_built:
+                    logger.error(
+                        "Could not build trust chain for provider %s with any configured trust anchor",
+                        provider_url,
+                    )
+            except Exception as e:
                 logger.error(
-                    f"Exception {exception} generated from this provider {provider_url}"
+                    "Could not resolve trust chain for %s: %s", provider_url, e
                 )
 
         return trust_chains
 
+    def _add_to_dict(self, d, url, chain):
+        """Helper to add a normalized URL in a dict."""
+        # Always store the exact URL key.
+        d[url] = chain
+        # Also store the normalized variant (with/without trailing slash),
+        # but avoid silently overwriting an existing normalized entry.
+        norm = url.rstrip("/") if url.endswith("/") else url + "/"
+        if norm != url:
+            if norm in d:
+                logger.warning(
+                    "Duplicate provider URL variants configured: %s and %s; "
+                    "keeping existing trust chain for %s",
+                    url,
+                    norm,
+                    norm,
+                )
+            else:
+                d[norm] = chain
+
     @staticmethod
     def generate_trust_chain(
         trust_anchor_ec: EntityStatement, provider_endpoint: str, httpc_params
@@ -141,3 +283,58 @@ def generate_trust_chain(
         trust_chain.start()
         trust_chain.apply_metadata_policy()
         return trust_chain
+
+    def _ensure_trust_anchors(self) -> List[EntityStatement]:
+        """Return a list of valid TAs."""
+        if not self._validated_trust_anchors:
+            httpc_params = self.config["trust_chain"]["config"]["httpc_params"]
+            ta_urls = self.config["trust_chain"]["config"]["trust_anchor"]
+
+            for ta_url in ta_urls:
+                try:
+                    jwt = get_entity_configurations(ta_url, httpc_params=httpc_params)[0]
+                    ta_ec = EntityStatement(jwt, httpc_params=httpc_params)
+                    ta_ec.validate_by_itself()
+                    self._validated_trust_anchors.append(ta_ec)
+                except Exception as e:
+                    logger.error(f"Failed to validate TA {ta_url}: {e}")
+
+            if not self._validated_trust_anchors:
+                raise ValueError("No valid Trust Anchors could be loaded.")
+
+        return self._validated_trust_anchors
+
+    def get_or_build_trust_chain(self, provider_url: str) -> TrustChainBuilder:
+        """
+        Get trust chain from cache, or from DB, or discover and build it on-demand.
+        Newly built chains are stored in memory and in the database.
+        """
+        providers = self.config.get("providers", [])
+        provider_variants = [provider_url, provider_url.rstrip("/")]
+        if not provider_url.endswith("/"):
+            provider_variants.append(provider_url + "/")
+        if not any(p in providers for p in provider_variants if p):
+            raise TrustChainNotFoundError(f"Provider {provider_url} not in allowed list.")
+
+        # Try load from DB (in-memory cache already checked by TrustChainResolver)
+        engine = self._get_storage()
+        if engine:
+            cached = engine.get_trust_chain_by_provider(provider_url)
+            if cached and not _is_cache_expired(cached):
+                chain = _trust_chain_from_cache(cached)
+                self._add_to_dict(self.trust_chain, provider_url, chain)
+                return chain
+
+        httpc_params = self.config["trust_chain"]["config"]["httpc_params"]
+        tas = self._ensure_trust_anchors()
+
+        for ta_ec in tas:
+            try:
+                chain = self.generate_trust_chain(ta_ec, provider_url, httpc_params)
+                self._add_to_dict(self.trust_chain, provider_url, chain)
+                self._store_trust_chain(chain, provider_url)
+                return chain
+            except Exception:
+                continue
+
+        raise TrustChainNotFoundError(f"Failed to build trust chain for {provider_url} with any TA.")