fix: replace python-jose library with PyJWK. Add signature verification errors. Add class ProtectedHeader(TypedDict) to signing.py.

sokoliva · sokoliva · commit 701eb3902922 · 2025-12-10T13:39:49.000Z
diff --git a/pyproject.toml b/pyproject.toml
@@ -36,7 +36,7 @@ telemetry = ["opentelemetry-api>=1.33.0", "opentelemetry-sdk>=1.33.0"]
 postgresql = ["sqlalchemy[asyncio,postgresql-asyncpg]>=2.0.0"]
 mysql = ["sqlalchemy[asyncio,aiomysql]>=2.0.0"]
 sqlite = ["sqlalchemy[asyncio,aiosqlite]>=2.0.0"]
-signing = ["python-jose>=3.0.0"]
+signing = ["PyJWT>=2.0.0"]
 
 sql = ["a2a-sdk[postgresql,mysql,sqlite]"]
 
diff --git a/src/a2a/utils/signing.py b/src/a2a/utils/signing.py
@@ -1,101 +1,111 @@
 import json
 
 from collections.abc import Callable
-from typing import Any
+from typing import Any, TypedDict
 
 
 try:
-    from jose import jws
-    from jose.backends.base import Key
-    from jose.exceptions import JOSEError
-    from jose.utils import base64url_decode, base64url_encode
+    import jwt
+
+    from jwt.api_jwk import PyJWK
+    from jwt.exceptions import PyJWTError
+    from jwt.utils import base64url_decode, base64url_encode
 except ImportError as e:
     raise ImportError(
-        'A2A Signing requires python-jose to be installed. '
+        'A2A Signing requires PyJWT to be installed. '
         'Install with: '
         "'pip install a2a-sdk[signing]'"
     ) from e
 
 from a2a.types import AgentCard, AgentCardSignature
 
 
+class SignatureVerificationError(Exception):
+    """Base exception for signature verification errors."""
+
+
+class NoSignatureError(SignatureVerificationError):
+    """Exception raised when no signature is found on an AgentCard."""
+
+
+class InvalidSignaturesError(SignatureVerificationError):
+    """Exception raised when all signatures are invalid."""
+
+
+class ProtectedHeader(TypedDict):
+    """Protected header parameters for JWS (JSON Web Signature)."""
+
+    kid: str
+    """ Key identifier. """
+    alg: str | None
+    """ Algorithm used for signing. """
+    jku: str | None
+    """ JSON Web Key Set URL. """
+    typ: str | None
+    """ Token type.
+
+    Best practice: SHOULD be "JOSE" for JWS tokens.
+    """
+
+
 def clean_empty(d: Any) -> Any:
-    """Recursively remove empty strings, lists, dicts, and None values from a dictionary."""
+    """Recursively remove empty strings, lists and dicts from a dictionary."""
     if isinstance(d, dict):
         cleaned_dict: dict[Any, Any] = {k: clean_empty(v) for k, v in d.items()}
-        return {
-            k: v
-            for k, v in cleaned_dict.items()
-            if v is not None and (isinstance(v, (bool, int, float)) or v)
-        }
+        return {k: v for k, v in cleaned_dict.items() if v}
     if isinstance(d, list):
         cleaned_list: list[Any] = [clean_empty(v) for v in d]
-        return [
-            v
-            for v in cleaned_list
-            if v is not None and (isinstance(v, (bool, int, float)) or v)
-        ]
-    return d if d not in ['', [], {}, None] else None
+        return [v for v in cleaned_list if v]
+    return d if d not in ['', [], {}] else None
 
 
 def canonicalize_agent_card(agent_card: AgentCard) -> str:
     """Canonicalizes the Agent Card JSON according to RFC 8785 (JCS)."""
     card_dict = agent_card.model_dump(
         exclude={'signatures'},
         exclude_defaults=True,
+        exclude_none=True,
         by_alias=True,
     )
-    # Ensure 'protocol_version' is always included
-    protocol_version_alias = (
-        AgentCard.model_fields['protocol_version'].alias or 'protocol_version'
-    )
-    if protocol_version_alias not in card_dict:
-        card_dict[protocol_version_alias] = agent_card.protocol_version
-
-    # Recursively remove empty/None values
+    # Recursively remove empty values
     cleaned_dict = clean_empty(card_dict)
-
     return json.dumps(cleaned_dict, separators=(',', ':'), sort_keys=True)
 
 
 def create_agent_card_signer(
-    signing_key: str | bytes | dict[str, Any] | Key,
-    kid: str,
-    alg: str = 'HS256',
-    jku: str | None = None,
+    signing_key: PyJWK | str | bytes,
+    protected_header: ProtectedHeader,
+    header: dict[str, Any] | None = None,
 ) -> Callable[[AgentCard], AgentCard]:
     """Creates a function that signs an AgentCard and adds the signature.
 
     Args:
         signing_key: The private key for signing.
-        kid: Key ID for the signing key.
-        alg: The algorithm to use (e.g., "ES256", "RS256").
-        jku: Optional URL to the JSON Web Keys.
+        protected_header: The protected header parameters.
+        header: Unprotected header parameters.
 
     Returns:
         A callable that takes an AgentCard and returns the modified AgentCard with a signature.
     """
 
     def agent_card_signer(agent_card: AgentCard) -> AgentCard:
-        """The actual card_modifier function."""
+        """Signs agent card."""
         canonical_payload = canonicalize_agent_card(agent_card)
+        payload_dict = json.loads(canonical_payload)
 
-        headers = {'kid': kid, 'typ': 'JOSE'}
-        if jku:
-            headers['jku'] = jku
-
-        jws_string = jws.sign(
-            payload=canonical_payload.encode('utf-8'),
+        jws_string = jwt.encode(
+            payload=payload_dict,
             key=signing_key,
-            headers=headers,
-            algorithm=alg,
+            algorithm=protected_header.get('alg', 'HS256'),
+            headers=protected_header,
         )
 
-        # The result of jws.sign is a compact serialization: HEADER.PAYLOAD.SIGNATURE
-        protected_header, _, signature = jws_string.split('.')
+        # The result of jwt.encode is a compact serialization: HEADER.PAYLOAD.SIGNATURE
+        protected, _, signature = jws_string.split('.')
 
         agent_card_signature = AgentCardSignature(
-            protected=protected_header,
+            header=header,
+            protected=protected,
             signature=signature,
         )
 
@@ -108,9 +118,7 @@ def agent_card_signer(agent_card: AgentCard) -> AgentCard:
 
 
 def create_signature_verifier(
-    key_provider: Callable[
-        [str | None, str | None], str | bytes | dict[str, Any] | Key
-    ],
+    key_provider: Callable[[str | None, str | None], PyJWK | str | bytes],
     algorithms: list[str],
 ) -> Callable[[AgentCard], None]:
     """Creates a function that verifies AgentCard signatures.
@@ -126,14 +134,17 @@ def create_signature_verifier(
     def signature_verifier(
         agent_card: AgentCard,
     ) -> None:
-        """The actual signature_verifier function."""
+        """Verifies agent card signatures.
+
+        Checks if at least one signature matches the key, otherwise raises an error.
+        """
         if not agent_card.signatures:
-            raise JOSEError('No signatures found on AgentCard')
+            raise NoSignatureError('AgentCard has no signatures to verify.')
 
         last_error = None
         for agent_card_signature in agent_card.signatures:
             try:
-                # fetch kid and jku from protected header
+                # get verification key
                 protected_header_json = base64url_decode(
                     agent_card_signature.protected.encode('utf-8')
                 ).decode('utf-8')
@@ -146,20 +157,22 @@ def signature_verifier(
                 encoded_payload = base64url_encode(
                     canonical_payload.encode('utf-8')
                 ).decode('utf-8')
-                token = f'{agent_card_signature.protected}.{encoded_payload}.{agent_card_signature.signature}'
 
-                jws.verify(
-                    token=token,
+                token = f'{agent_card_signature.protected}.{encoded_payload}.{agent_card_signature.signature}'
+                jwt.decode(
+                    jwt=token,
                     key=verification_key,
                     algorithms=algorithms,
                 )
                 # Found a valid signature, exit the loop and function
                 break
-            except JOSEError as e:
+            except PyJWTError as e:
                 last_error = e
                 continue
         else:
             # This block runs only if the loop completes without a break
-            raise JOSEError('No valid signature found') from last_error
+            raise InvalidSignaturesError(
+                'No valid signature found'
+            ) from last_error
 
     return signature_verifier
diff --git a/tests/integration/test_client_server_integration.py b/tests/integration/test_client_server_integration.py
@@ -8,8 +8,8 @@
 import pytest
 import pytest_asyncio
 from grpc.aio import Channel
-from jose.backends.base import Key
 
+from jwt.api_jwk import PyJWK
 from a2a.client import ClientConfig
 from a2a.client.base_client import BaseClient
 from a2a.client.transports import JsonRpcTransport, RestTransport
@@ -89,7 +89,7 @@
 )
 
 
-def create_key_provider(verification_key: str | bytes | dict[str, Any] | Key):
+def create_key_provider(verification_key: PyJWK | str | bytes):
     """Creates a key provider function for testing."""
 
     def key_provider(kid: str | None, jku: str | None):
@@ -754,6 +754,7 @@ async def test_http_transport_get_authenticated_card(
     transport = RestTransport(httpx_client=httpx_client, agent_card=agent_card)
     result = await transport.get_card()
     assert result.name == extended_agent_card.name
+    assert transport.agent_card is not None
     assert transport.agent_card.name == extended_agent_card.name
     assert transport._needs_extended_card is False
 
@@ -776,6 +777,7 @@ def channel_factory(address: str) -> Channel:
     transport = GrpcTransport(channel=channel, agent_card=agent_card)
 
     # The transport starts with a minimal card, get_card() fetches the full one
+    assert transport.agent_card is not None
     transport.agent_card.supports_authenticated_extended_card = True
     result = await transport.get_card()
 
@@ -845,7 +847,7 @@ async def test_json_transport_base_client_send_message_with_extensions(
 
 
 @pytest.mark.asyncio
-async def test_json_transport_get_signed_base_card_no_initial(
+async def test_json_transport_get_signed_base_card(
     jsonrpc_setup: TransportSetup, agent_card: AgentCard
 ) -> None:
     """Tests fetching and verifying a symmetrically signed AgentCard via JSON-RPC.
@@ -860,7 +862,13 @@ async def test_json_transport_get_signed_base_card_no_initial(
     # Setup signing on the server side
     key = 'key12345'
     signer = create_agent_card_signer(
-        signing_key=key, alg='HS384', kid='testkey'
+        signing_key=key,
+        protected_header={
+            'alg': 'HS384',
+            'kid': 'testkey',
+            'jku': None,
+            'typ': 'JOSE',
+        },
     )
 
     app_builder = A2AFastAPIApplication(
@@ -885,6 +893,7 @@ async def test_json_transport_get_signed_base_card_no_initial(
     assert result.name == agent_card.name
     assert result.signatures is not None
     assert len(result.signatures) == 1
+    assert transport.agent_card is not None
     assert transport.agent_card.name == agent_card.name
     assert transport._needs_extended_card is False
 
@@ -911,7 +920,13 @@ async def test_json_transport_get_signed_extended_card(
     private_key = asymmetric.ec.generate_private_key(asymmetric.ec.SECP256R1())
     public_key = private_key.public_key()
     signer = create_agent_card_signer(
-        signing_key=private_key, alg='ES256', kid='testkey'
+        signing_key=private_key,
+        protected_header={
+            'alg': 'ES256',
+            'kid': 'testkey',
+            'jku': None,
+            'typ': 'JOSE',
+        },
     )
 
     app_builder = A2AFastAPIApplication(
@@ -937,6 +952,7 @@ async def test_json_transport_get_signed_extended_card(
     assert result.name == extended_agent_card.name
     assert result.signatures is not None
     assert len(result.signatures) == 1
+    assert transport.agent_card is not None
     assert transport.agent_card.name == extended_agent_card.name
     assert transport._needs_extended_card is False
 
@@ -964,7 +980,13 @@ async def test_json_transport_get_signed_base_and_extended_cards(
     private_key = asymmetric.ec.generate_private_key(asymmetric.ec.SECP256R1())
     public_key = private_key.public_key()
     signer = create_agent_card_signer(
-        signing_key=private_key, alg='ES256', kid='testkey'
+        signing_key=private_key,
+        protected_header={
+            'alg': 'ES256',
+            'kid': 'testkey',
+            'jku': None,
+            'typ': 'JOSE',
+        },
     )
 
     app_builder = A2AFastAPIApplication(
@@ -993,6 +1015,7 @@ async def test_json_transport_get_signed_base_and_extended_cards(
     assert result.name == extended_agent_card.name
     assert result.signatures is not None
     assert len(result.signatures) == 1
+    assert transport.agent_card is not None
     assert transport.agent_card.name == extended_agent_card.name
     assert transport._needs_extended_card is False
 
@@ -1019,7 +1042,13 @@ async def test_rest_transport_get_signed_card(
     private_key = asymmetric.ec.generate_private_key(asymmetric.ec.SECP256R1())
     public_key = private_key.public_key()
     signer = create_agent_card_signer(
-        signing_key=private_key, alg='ES256', kid='testkey'
+        signing_key=private_key,
+        protected_header={
+            'alg': 'ES256',
+            'kid': 'testkey',
+            'jku': None,
+            'typ': 'JOSE',
+        },
     )
 
     app_builder = A2ARESTFastAPIApplication(
@@ -1048,6 +1077,7 @@ async def test_rest_transport_get_signed_card(
     assert result.name == extended_agent_card.name
     assert result.signatures is not None
     assert len(result.signatures) == 1
+    assert transport.agent_card is not None
     assert transport.agent_card.name == extended_agent_card.name
     assert transport._needs_extended_card is False
 
@@ -1066,7 +1096,13 @@ async def test_grpc_transport_get_signed_card(
     private_key = asymmetric.ec.generate_private_key(asymmetric.ec.SECP256R1())
     public_key = private_key.public_key()
     signer = create_agent_card_signer(
-        signing_key=private_key, alg='ES256', kid='testkey'
+        signing_key=private_key,
+        protected_header={
+            'alg': 'ES256',
+            'kid': 'testkey',
+            'jku': None,
+            'typ': 'JOSE',
+        },
     )
 
     server = grpc.aio.server()
diff --git a/tests/utils/test_signing.py b/tests/utils/test_signing.py
