Skip to content

Commit

Permalink
feat(low-code): added json.loads to jwt authenticator (#301)
Browse files Browse the repository at this point in the history
Co-authored-by: octavia-squidington-iii <[email protected]>
  • Loading branch information
lazebnyi and octavia-squidington-iii authored Jan 30, 2025
1 parent a6d55be commit dea2cc9
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 11 deletions.
28 changes: 17 additions & 11 deletions airbyte_cdk/sources/declarative/auth/jwt.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
#

import base64
import json
from dataclasses import InitVar, dataclass
from datetime import datetime
from typing import Any, Mapping, Optional, Union
Expand Down Expand Up @@ -104,21 +105,21 @@ def __post_init__(self, parameters: Mapping[str, Any]) -> None:
)

def _get_jwt_headers(self) -> dict[str, Any]:
""" "
"""
Builds and returns the headers used when signing the JWT.
"""
headers = self._additional_jwt_headers.eval(self.config)
headers = self._additional_jwt_headers.eval(self.config, json_loads=json.loads)
if any(prop in headers for prop in ["kid", "alg", "typ", "cty"]):
raise ValueError(
"'kid', 'alg', 'typ', 'cty' are reserved headers and should not be set as part of 'additional_jwt_headers'"
)

if self._kid:
headers["kid"] = self._kid.eval(self.config)
headers["kid"] = self._kid.eval(self.config, json_loads=json.loads)
if self._typ:
headers["typ"] = self._typ.eval(self.config)
headers["typ"] = self._typ.eval(self.config, json_loads=json.loads)
if self._cty:
headers["cty"] = self._cty.eval(self.config)
headers["cty"] = self._cty.eval(self.config, json_loads=json.loads)
headers["alg"] = self._algorithm
return headers

Expand All @@ -130,18 +131,19 @@ def _get_jwt_payload(self) -> dict[str, Any]:
exp = now + self._token_duration if isinstance(self._token_duration, int) else now
nbf = now

payload = self._additional_jwt_payload.eval(self.config)
payload = self._additional_jwt_payload.eval(self.config, json_loads=json.loads)
if any(prop in payload for prop in ["iss", "sub", "aud", "iat", "exp", "nbf"]):
raise ValueError(
"'iss', 'sub', 'aud', 'iat', 'exp', 'nbf' are reserved properties and should not be set as part of 'additional_jwt_payload'"
)

if self._iss:
payload["iss"] = self._iss.eval(self.config)
payload["iss"] = self._iss.eval(self.config, json_loads=json.loads)
if self._sub:
payload["sub"] = self._sub.eval(self.config)
payload["sub"] = self._sub.eval(self.config, json_loads=json.loads)
if self._aud:
payload["aud"] = self._aud.eval(self.config)
payload["aud"] = self._aud.eval(self.config, json_loads=json.loads)

payload["iat"] = now
payload["exp"] = exp
payload["nbf"] = nbf
Expand All @@ -151,7 +153,7 @@ def _get_secret_key(self) -> str:
"""
Returns the secret key used to sign the JWT.
"""
secret_key: str = self._secret_key.eval(self.config)
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
Expand All @@ -176,7 +178,11 @@ def _get_header_prefix(self) -> Union[str, None]:
"""
Returns the header prefix to be used when attaching the token to the request.
"""
return self._header_prefix.eval(self.config) if self._header_prefix else None
return (
self._header_prefix.eval(self.config, json_loads=json.loads)
if self._header_prefix
else None
)

@property
def auth_header(self) -> str:
Expand Down
14 changes: 14 additions & 0 deletions unit_tests/sources/declarative/auth/test_jwt.py
Original file line number Diff line number Diff line change
Expand Up @@ -126,6 +126,20 @@ def test_get_secret_key(self, base64_encode_secret_key, secret_key, expected):
)
assert authenticator._get_secret_key() == expected

def test_get_secret_key_from_config(
self,
):
authenticator = JwtAuthenticator(
config={"secrets": '{"secret_key": "test"}'},
parameters={},
secret_key="{{ json_loads(config['secrets'])['secret_key'] }}",
algorithm="test_algo",
token_duration=1200,
base64_encode_secret_key=False,
)
expected = "test"
assert authenticator._get_secret_key() == expected

def test_get_signed_token(self):
authenticator = JwtAuthenticator(
config={},
Expand Down

0 comments on commit dea2cc9

Please sign in to comment.