add passphrase logic to jwt auth

Author: pnilan

airbyte_cdk/sources/declarative/auth/jwt.py

@@ -9,6 +9,9 @@
 from typing import Any, Mapping, Optional, Union
 
 import jwt
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric.types import PrivateKeyTypes
 
 from airbyte_cdk.sources.declarative.auth.declarative_authenticator import DeclarativeAuthenticator
 from airbyte_cdk.sources.declarative.interpolation.interpolated_boolean import InterpolatedBoolean
@@ -74,6 +77,7 @@ class JwtAuthenticator(DeclarativeAuthenticator):
     aud: Optional[Union[InterpolatedString, str]] = None
     additional_jwt_headers: Optional[Mapping[str, Any]] = None
     additional_jwt_payload: Optional[Mapping[str, Any]] = None
+    passphrase: Optional[Union[InterpolatedString, str]] = None
 
     def __post_init__(self, parameters: Mapping[str, Any]) -> None:
         self._secret_key = InterpolatedString.create(self.secret_key, parameters=parameters)
@@ -103,6 +107,11 @@ def __post_init__(self, parameters: Mapping[str, Any]) -> None:
         self._additional_jwt_payload = InterpolatedMapping(
             self.additional_jwt_payload or {}, parameters=parameters
         )
+        self._passphrase = (
+            InterpolatedString.create(self.passphrase, parameters=parameters)
+            if self.passphrase
+            else None
+        )
 
     def _get_jwt_headers(self) -> dict[str, Any]:
         """
@@ -149,16 +158,24 @@ def _get_jwt_payload(self) -> dict[str, Any]:
         payload["nbf"] = nbf
         return payload
 
-    def _get_secret_key(self) -> str:
+    def _get_secret_key(self) -> PrivateKeyTypes | str | bytes:
         """
         Returns the secret key used to sign the JWT.
         """
         secret_key: str = self._secret_key.eval(self.config, json_loads=json.loads)
-        return (
-            base64.b64encode(secret_key.encode()).decode()
-            if self._base64_encode_secret_key
-            else secret_key
-        )
+
+        if self._passphrase:
+            return serialization.load_pem_private_key(
+                secret_key.encode(),
+                password=self._passphrase.eval(self.config, json_loads=json.loads).encode(),
+                backend=default_backend(),
+            )
+        else:
+            return (
+                base64.b64encode(secret_key.encode()).decode()
+                if self._base64_encode_secret_key
+                else secret_key
+            )
 
     def _get_signed_token(self) -> Union[str, Any]:
         """

airbyte_cdk/sources/declarative/declarative_component_schema.yaml

@@ -1270,6 +1270,12 @@ definitions:
         title: Additional JWT Payload Properties
         description: Additional properties to be added to the JWT payload.
         additionalProperties: true
+      passphrase:
+        title: Passphrase
+        description: A passphrase/password used to encrypt the private key. Only provide a passphrase if required by the API for JWT authentication. The API will typically provide the passphrase when generating the public/private key pair.
+        type: string
+        examples:
+          - "{{ config['passphrase'] }}"
       $parameters:
         type: object
         additionalProperties: true

airbyte_cdk/sources/declarative/interpolation/macros.py

@@ -6,6 +6,7 @@
 import datetime
 import re
 import typing
+import uuid
 from typing import Optional, Union
 from urllib.parse import quote_plus
 
@@ -207,6 +208,13 @@ def camel_case_to_snake_case(value: str) -> str:
     return re.sub(r"(?<!^)(?=[A-Z])", "_", value).lower()
 
 
+def random_uuid() -> str:
+    """
+    Generates a UUID4
+    """
+    return str(uuid.uuid4())
+
+
 _macros_list = [
     now_utc,
     today_utc,

airbyte_cdk/sources/declarative/models/declarative_component_schema.py

@@ -448,6 +448,12 @@ class JwtAuthenticator(BaseModel):
         description="Additional properties to be added to the JWT payload.",
         title="Additional JWT Payload Properties",
     )
+    passphrase: Optional[str] = Field(
+        None,
+        description="A passphrase/password used to encrypt the private key. Only provide a passphrase if required by the API for JWT authentication. The API will typically provide the passphrase when generating the public/private key pair.",
+        examples=["{{ config['passphrase'] }}"],
+        title="Passphrase",
+    )
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
 
 

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -2705,6 +2705,7 @@ def create_jwt_authenticator(
             aud=jwt_payload.aud,
             additional_jwt_headers=model.additional_jwt_headers,
             additional_jwt_payload=model.additional_jwt_payload,
+            passphrase=model.passphrase,
         )
 
     def create_list_partition_router(

unit_tests/sources/declarative/auth/test_jwt.py

@@ -8,6 +8,9 @@
 import freezegun
 import jwt
 import pytest
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
 
 from airbyte_cdk.sources.declarative.auth.jwt import JwtAuthenticator
 
@@ -185,3 +188,100 @@ def test_get_header_prefix(self, header_prefix, expected):
             header_prefix=header_prefix,
         )
         assert authenticator._get_header_prefix() == expected
+
+    def test_get_secret_key_with_passphrase(self):
+        """Test _get_secret_key method with encrypted private key and passphrase."""
+        # Generate a test RSA private key
+        private_key = rsa.generate_private_key(
+            public_exponent=65537, key_size=2048, backend=default_backend()
+        )
+
+        passphrase = b"test_passphrase"
+        encrypted_pem = private_key.private_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PrivateFormat.PKCS8,
+            encryption_algorithm=serialization.BestAvailableEncryption(passphrase),
+        )
+
+        authenticator = JwtAuthenticator(
+            config={},
+            parameters={},
+            secret_key=encrypted_pem.decode(),
+            algorithm="RS256",
+            token_duration=1200,
+            passphrase="test_passphrase",
+        )
+
+        result_key = authenticator._get_secret_key()
+
+        assert isinstance(result_key, rsa.RSAPrivateKey)
+
+        original_public_key = private_key.public_key()
+        result_public_key = result_key.public_key()
+
+        original_public_numbers = original_public_key.public_numbers()
+        result_public_numbers = result_public_key.public_numbers()
+
+        assert original_public_numbers.n == result_public_numbers.n
+        assert original_public_numbers.e == result_public_numbers.e
+
+    def test_get_secret_key_with_wrong_passphrase_raises_error(self):
+        """Test that _get_secret_key raises error with wrong passphrase."""
+        private_key = rsa.generate_private_key(
+            public_exponent=65537, key_size=2048, backend=default_backend()
+        )
+
+        passphrase = b"correct_passphrase"
+        encrypted_pem = private_key.private_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PrivateFormat.PKCS8,
+            encryption_algorithm=serialization.BestAvailableEncryption(passphrase),
+        )
+
+        authenticator = JwtAuthenticator(
+            config={},
+            parameters={},
+            secret_key=encrypted_pem.decode(),
+            algorithm="RS256",
+            token_duration=1200,
+            passphrase="wrong_passphrase",
+        )
+
+        with pytest.raises(Exception):
+            authenticator._get_secret_key()
+
+    def test_get_signed_token_with_passphrase_protected_key(self):
+        """Test that JWT signing works with passphrase-protected RSA private key."""
+        private_key = rsa.generate_private_key(
+            public_exponent=65537, key_size=2048, backend=default_backend()
+        )
+
+        passphrase = b"test_passphrase"
+        encrypted_pem = private_key.private_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PrivateFormat.PKCS8,
+            encryption_algorithm=serialization.BestAvailableEncryption(passphrase),
+        )
+
+        authenticator = JwtAuthenticator(
+            config={},
+            parameters={},
+            secret_key=encrypted_pem.decode(),
+            algorithm="RS256",
+            token_duration=1000,
+            passphrase="test_passphrase",
+            typ="JWT",
+            iss="test_issuer",
+        )
+
+        signed_token = authenticator._get_signed_token()
+
+        assert isinstance(signed_token, str)
+        assert len(signed_token.split(".")) == 3
+
+        public_key = private_key.public_key()
+        decoded_payload = jwt.decode(signed_token, public_key, algorithms=["RS256"])
+
+        assert decoded_payload["iss"] == "test_issuer"
+        assert "iat" in decoded_payload
+        assert "exp" in decoded_payload