[monitoringlib/auth] Add support for JWTs as client authentification (RFC 7523 section 2.2)

monitoring/mock_uss/app.py

@@ -154,13 +154,15 @@ def require_config_value(config_key: str) -> None:
     "ClientIdClientSecret": 4,
     "Keycloak": 3,
     "FlightPassport": 4,
+    "PrivateKeyJWT": 5,
 }
 
 _SECRETS = {
     "UsernamePassword": (2, "password"),
     "ClientIdClientSecret": (2, "client_secret"),
     "Keycloak": (2, "client_secret"),
     "FlightPassport": (2, "client_secret"),
+    "PrivateKeyJWT": (3, "key"),
 }
 
 

monitoring/mock_uss/app_test.py

@@ -110,6 +110,24 @@ def test_client_secret_mixed_with_trailing_positional():
     )
 
 
+def test_jwt_bearer_key_args():
+    assert (
+        sanitize_secrets(
+            K, "PrivateKeyJWT(http://host/token,cli,key=eyJrdHkiOiJSU0EifQ,key_id=k1)"
+        )
+        == "PrivateKeyJWT(http://host/token, cli, /auth/uss1.key, ***)"
+    )
+
+
+def test_jwt_bearer_key_kwarg():
+    assert (
+        sanitize_secrets(
+            K, "PrivateKeyJWT(http://host/token,cli,key=eyJrdHkiOiJSU0EifQ,key_id=k1)"
+        )
+        == "PrivateKeyJWT(http://host/token, cli, key=***, key_id=k1)"
+    )
+
+
 def test_unknown_adapter_hides_everything():
     assert (
         sanitize_secrets(K, "MysteryAuth(http://host/token,cli,whatever,foo=bar)")

monitoring/monitorlib/README.md

@@ -30,6 +30,8 @@ AuthAdapter's `__init__` constructor.  Both ordinal (e.g.,
 * `SignedRequest(https://example.interuss.org/oauth/token, client_id=uss1.com,
    key_path=/auth/uss1.key, cert_url=https://uss1.com/uss1.der)`
 * `ClientIdClientSecret(https://example.interuss.org/token, uss1, dXNzMQ==)`
+* `PrivateKeyJWT(https://example.interuss.org/token, uss1, key_path=/auth/uss1.key)`
+* `PrivateKeyJWT(https://example.interuss.org/token, uss1, key=LS0tLS1CRUdJTi<...>, key_id=uss1)`
 
 ### Testing
 

monitoring/monitorlib/auth.py

@@ -594,6 +594,115 @@ def __init__(
         self._send_request_as_data = send_request_as_data
 
 
+class PrivateKeyJWT(AuthAdapter):
+    """Auth adapter that authenticates with a private key JWT client assertion.
+
+    See RFC 7523 section 2.2.
+    """
+
+    def __init__(
+        self,
+        token_endpoint: str,
+        client_id: str,
+        key_path: str | None = None,
+        key: str | None = None,
+        key_id: str | None = None,
+    ):
+        """Create an AuthAdapter that retrieves tokens with a client assertion.
+
+        Args:
+          token_endpoint: URL of the authorization server's token endpoint.
+          client_id: ID of client for which the token is being requested.
+          key_path: Path to a PEM or JWK file containing the private key with which
+            to sign the client assertion.  Mutually exclusive with key.
+          key: base64 encoding of the PEM or JWK content of the private key with
+            which to sign the client assertion.  Mutually exclusive with key_path.
+          key_id: If specified, the specific ID to supply in the JWS header.  If not
+            specified, no key ID is supplied.
+        """
+
+        super().__init__()
+
+        self._oauth_token_endpoint = token_endpoint
+        self._client_id = client_id
+
+        if key_path and key:
+            raise ValueError("Specify only one of key_path or key")
+
+        if key_path:
+            with open(key_path) as f:
+                key_content = f.read()
+        elif key:
+            key_content = base64.b64decode(key + "=" * (-len(key) % 4)).decode("utf-8")
+        else:
+            raise ValueError("Either key_path or key must be specified")
+
+        if key_content.lstrip().startswith("{"):
+            self._key = jwcrypto.jwk.JWK.from_json(key_content)
+        else:
+            self._key = jwcrypto.jwk.JWK.from_pem(key_content.encode("utf-8"))
+
+        kty = self._key["kty"]
+
+        if kty == "RSA":
+            self._alg = "RS256"
+        elif kty == "EC":
+            crv_algs = {"P-256": "ES256", "P-384": "ES384", "P-521": "ES512"}
+            crv = self._key["crv"]
+
+            if crv not in crv_algs:
+                raise ValueError(f"Unsupported EC curve `{crv}`")
+
+            self._alg = crv_algs[crv]
+        elif kty == "OKP":
+            self._alg = "EdDSA"
+        else:
+            raise ValueError(f"Unsupported key type `{kty}`")
+
+        self._kid = key_id
+
+    def issue_token(self, intended_audience: str, scopes: list[str]) -> str:
+
+        timestamp = int(
+            (datetime.datetime.now(datetime.UTC) - _UNIX_EPOCH).total_seconds()
+        )
+
+        header = {"typ": "JWT", "alg": self._alg}
+        if self._kid:
+            header["kid"] = self._kid
+
+        assertion = jwcrypto.jwt.JWT(
+            header=header,
+            claims={
+                "iss": self._client_id,
+                "sub": self._client_id,
+                "aud": self._oauth_token_endpoint,
+                "iat": timestamp,
+                "exp": timestamp + 300,
+                "jti": str(uuid.uuid4()),
+            },
+        )
+
+        assertion.make_signed_token(self._key)
+
+        response = requests.post(
+            self._oauth_token_endpoint,
+            data={
+                "grant_type": "client_credentials",
+                "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer",
+                "client_assertion": assertion.serialize(),
+                "audience": intended_audience,
+                "scope": " ".join(scopes),
+            },
+        )
+
+        if response.status_code != 200:
+            raise AccessTokenError(
+                "Unable to retrieve access token:\n" + response.content.decode("utf-8")
+            )
+        return response.json()["access_token"]
+
+
 class AccessTokenError(RuntimeError):
     def __init__(self, msg):
         super().__init__(msg)

monitoring/monitorlib/auth_test.py

@@ -0,0 +1,91 @@
+import base64
+import json
+import os
+from unittest import mock
+
+import jwcrypto.jwk
+import jwcrypto.jwt
+import pytest
+
+from monitoring.monitorlib import auth
+
+_KEY_PATH = os.path.join(
+    os.path.dirname(__file__), "..", "..", "build", "test-certs", "auth2.key"
+)
+with open(_KEY_PATH, "rb") as f:
+    _KEY = jwcrypto.jwk.JWK.from_pem(f.read())
+
+
+class _Response:
+    def __init__(self, status_code: int, content: bytes):
+        self.status_code = status_code
+        self.content = content
+
+    def json(self):
+        return json.loads(self.content.decode("utf-8"))
+
+
+def test_jwt_bearer_key_from_pem_file():
+    adapter = auth.PrivateKeyJWT("http://host/token", "cli", key_path=_KEY_PATH)
+    assert adapter._alg == "RS256"
+    assert adapter._kid is None
+
+
+def test_jwt_bearer_key_from_jwk():
+    adapter = auth.PrivateKeyJWT(
+        "http://host/token",
+        "cli",
+        key=base64.b64encode(_KEY.export(private_key=True).encode("utf-8")).decode(
+            "utf-8"
+        ),
+        key_id="k1",
+    )
+    assert adapter._alg == "RS256"
+    assert adapter._kid == "k1"
+
+
+def test_jwt_bearer_missing_or_duplicate_key():
+    with pytest.raises(ValueError):
+        auth.PrivateKeyJWT("http://host/token", "cli")
+    with pytest.raises(ValueError):
+        auth.PrivateKeyJWT("http://host/token", "cli", key_path=_KEY_PATH, key="{}")
+
+
+def test_jwt_bearer_issue_token():
+    adapter = auth.PrivateKeyJWT("http://host/token", "cli", key_path=_KEY_PATH)
+    with mock.patch.object(
+        auth.requests, "post", return_value=_Response(200, b'{"access_token": "t0ken"}')
+    ) as post:
+        assert (
+            adapter.issue_token("https://uss.example.com", ["scope1", "scope2"])
+            == "t0ken"
+        )
+
+    assert post.call_args[0][0] == "http://host/token"
+    payload = post.call_args[1]["data"]
+    assert payload["grant_type"] == "client_credentials"
+    assert payload["audience"] == "https://uss.example.com"
+    assert payload["scope"] == "scope1 scope2"
+
+    header = jwcrypto.jwt.JWT(
+        jwt=payload["client_assertion"], key=_KEY
+    ).token.jose_header
+    assert header["alg"] == "RS256"
+    assert "kid" not in header
+
+    claims = json.loads(
+        jwcrypto.jwt.JWT(jwt=payload["client_assertion"], key=_KEY).claims
+    )
+    assert claims["iss"] == "cli"
+    assert claims["sub"] == "cli"
+    assert claims["aud"] == "http://host/token"
+    assert claims["exp"] == claims["iat"] + 300
+
+
+def test_jwt_bearer_issue_token_error():
+    adapter = auth.PrivateKeyJWT("http://host/token", "cli", key_path=_KEY_PATH)
+    with mock.patch.object(
+        auth.requests, "post", return_value=_Response(401, b"invalid_grant")
+    ):
+        with pytest.raises(auth.AccessTokenError):
+            adapter.issue_token("https://uss.example.com", ["scope1"])