Jona/support entra auth

.gitignore

@@ -6,6 +6,7 @@ docker-compose.override.y*ml
 sld/config.json
 __pycache__
 .env
+.envrc
 venv
 .vscode
 index.json

compose.yml

@@ -8,10 +8,10 @@ services:
     ports:
       - "5403:5432"
     environment:
-      - POSTGRES_PASSWORD=insecure
+      - POSTGRES_PASSWORD="${POSTGRES_PASSWORD}"
       - POSTGRES_DB=postgres
 
-    command: ["postgres", "-c", "log_statement=all"]
+    command: [ "postgres", "-c", "log_statement=all" ]
     volumes:
       - "${POSTGRES_DATA_DIR:-/var/lib/postgresql/data}:/var/lib/postgresql/data"
 
@@ -22,12 +22,15 @@ services:
     ports:
       - 9191:8000
     environment:
-      PROXY_URL: http://map
+      PRIVATE_PROXY_URL: http://private:8080
+      PUBLIC_PROXY_URL: http://public:8080
       JWKS_PATH: /app/test_jwk.json
       LOG_LEVEL: DEBUG
-      #JWKS_URL: "https://iam.amsterdam.nl/auth/realms/datapunt-ad-acc/protocol/openid-connect/certs"
-
-  # Bare bones 'webserver' to behave as public loadbalancer
+      OAUTH_KEYCLOAK_CLAIMS: "${OAUTH_KEYCLOAK_CLAIMS}"
+      OAUTH_KEYCLOAK_URL: "${OAUTH_KEYCLOAK_URL}"
+      OAUTH_ENTRA_CLAIMS: "${OAUTH_ENTRA_CLAIMS}"
+      OAUTH_ENTRA_URL: "${OAUTH_ENTRA_URL}"
+      # Bare bones 'webserver' to behave as public loadbalancer
   publicdummy:
     image: python:3.13-slim-bullseye
     ports:
@@ -63,7 +66,7 @@ services:
       args:
         BASE_IMAGE: mapserver/public
     depends_on:
-    - public
+      - public
     ports:
       - "8384:8080"
     environment:
@@ -90,7 +93,7 @@ services:
       args:
         BASE_IMAGE: mapserver/public
     depends_on:
-    - public
+      - public
     ports:
       - "8385:8080"
     environment:

jwtproxy/src/server.py

@@ -9,7 +9,8 @@
 from os import environ as env
 
 from aiohttp import ClientSession, ClientTimeout, DummyCookieJar, TCPConnector, web
-from jwcrypto.jwt import JWT, JWKSet, JWTExpired
+from jwcrypto.common import JWException
+from jwcrypto.jwt import JWT, JWKSet, JWTExpired, JWTMissingKey
 from yarl import URL
 
 LOG_LEVEL = env.get("LOG_LEVEL", "INFO")
@@ -121,26 +122,45 @@ async def fetch_token(req: web.Request):
 
 
 async def audit(req: web.Request, token: str):
-    # Note that this function only supports token formats issued by
-    # the iam KeyCloak server.
-
-    j = JWT()
-    jwks = await get_jwk()
-
-    try:
-        j.deserialize(token, key=jwks)
-    except ValueError as e:
-        app_logger.warning(e)
-        raise web.HTTPUnauthorized(reason="Invalid token") from e
-    except JWTExpired as e:
-        app_logger.warning(e)
-        raise web.HTTPUnauthorized(reason="Token expired") from e
-
-    claims = json.loads(j.claims)
-    try:
-        if "fp_mdw" not in claims["realm_access"]["roles"]:
-            raise ValueError("Required scope not present in claims")
-    except (KeyError, ValueError):
+    # This function now supports tokens issued by both Keycloak and Entra ID
+
+    # First check if the access token is valid
+    claims = None
+    for provider in await get_jwks():
+        try:
+            j = JWT(
+                jwt=token,
+                key=provider["jwks"],
+                check_claims=provider["check_claims"],
+            )
+        except JWTMissingKey:
+            continue
+        except JWTExpired as e:
+            app_logger.warning(e)
+            raise web.HTTPUnauthorized(reason="Token expired") from e
+        except (JWException, ValueError) as e:
+            app_logger.warning(e)
+            continue
+
+        candidate_claims = json.loads(j.claims)
+        if not isinstance(candidate_claims, dict):
+            raise web.HTTPUnauthorized(reason="Invalid token")
+        claims = candidate_claims
+        break
+
+    if claims is None:
+        raise web.HTTPUnauthorized(reason="Invalid token")
+
+    # Then check is the validated access token contains the correct scope
+    realm_access = claims.get("realm_access")
+    keycloak_roles = realm_access.get("roles") if isinstance(realm_access, dict) else None
+    entra_roles = claims.get("roles")
+    has_access = (
+        isinstance(keycloak_roles, list) and "fp_mdw" in keycloak_roles
+    ) or (isinstance(entra_roles, list) and "FP/MDW" in entra_roles)
+
+    # Deny requests if no_access:
+    if not has_access:
         audit_logger.info(
             "Access to %s (%s) denied for subject %s (username: %s)",
             req.url,
@@ -191,31 +211,59 @@ async def on_shutdown(app):
 
 
 lock = asyncio.Lock()
-cached_jwk = None
-
+cached_jwks = None
 
-async def get_jwk():
-    """If we can get the pubkey from a URL, get it from there.
-    Otherwise, default to the filesystem."""
-    # https://login.microsoftonline.com/{tenant}/discovery/v2.0/keys
+async def get_jwks():
+    """Load and cache the configured issuer/JWKS pairs."""
     async with lock:
-        global cached_jwk
-        if cached_jwk is not None:
-            return cached_jwk
+        global cached_jwks
+        if cached_jwks is not None:
+            return cached_jwks
+
+        providers = (
+            (env.get("OAUTH_KEYCLOAK_CLAIMS"), env.get("OAUTH_KEYCLOAK_URL")),
+            (env.get("OAUTH_ENTRA_CLAIMS"), env.get("OAUTH_ENTRA_URL")),
+        )
+        keysets = []
+        for check_claims, url in providers:
+            if check_claims is None and url is None:
+                continue
+            if not check_claims or not url:
+                raise RuntimeError("Each provider must have a JWKS_URL and CHECK_CLAIMS var")
+
+            check_claims = dict(item.split("=", 1) for item in check_claims.split(","))
+            if not check_claims.get("iss"):
+                raise RuntimeError("check_claims must include an issuer (iss)")
+
+            # 
+            check_claims["exp"] = None
 
-        url = env.get("JWKS_URL")
-        if url:
             async with SessionManager.session().get(
                 url,
                 timeout=ClientTimeout(sock_connect=3, sock_read=3),
                 skip_auto_headers=["Accept-Encoding"],  # we dont need a compressed jwks
             ) as resp:
-                payload = await resp.text()
-        else:
-            with open(env.get("JWKS_PATH")) as file:
-                payload = file.read()
-        cached_jwk = JWKSet.from_json(payload)
-        return cached_jwk
+                resp.raise_for_status()
+                provider = {
+                    "check_claims": check_claims,
+                    "jwks": JWKSet.from_json(await resp.text()),
+                }
+                keysets.append(provider)
+
+        if not keysets:
+            jwks_path = env.get("JWKS_PATH")
+            if not jwks_path:
+                raise RuntimeError("Configure issuer/JWKS URL pairs or JWKS_PATH")
+            with open(jwks_path) as file:
+                keysets.append(
+                    {
+                        "check_claims": {"exp": None},
+                        "jwks": JWKSet.from_json(file.read()),
+                    }
+                )
+
+        cached_jwks = keysets
+        return cached_jwks
 
 
 async def check_protection(path, method):

jwtproxy/src/tests/conftest.py

@@ -26,6 +26,13 @@ def set_env():
     os.environ["PRIVATE_PROXY_URL"] = PRIVATE_PROXY_URL
     os.environ["PUBLIC_PROXY_URL"] = PUBLIC_PROXY_URL
     os.environ["JWKS_PATH"] = "test_jwk.json"
+    for name in (
+        "OAUTH_KEYCLOAK_CLAIMS",
+        "OAUTH_KEYCLOAK_URL",
+        "OAUTH_ENTRA_CLAIMS",
+        "OAUTH_ENTRA_URL",
+    ):
+        os.environ.pop(name, None)
 
 
 @pytest.fixture