codekingpro/portable-devtools

1from __future__ import annotations2 3import json4import time5from typing import Any6 7from .algorithms import get_default_algorithms, has_crypto, requires_cryptography8from .exceptions import InvalidKeyError, PyJWKError, PyJWKSetError, PyJWTError9from .types import JWKDict10 11 12class PyJWK:13    def __init__(self, jwk_data: JWKDict, algorithm: str | None = None) -> None:14        self._algorithms = get_default_algorithms()15        self._jwk_data = jwk_data16 17        kty = self._jwk_data.get("kty", None)18        if not kty:19            raise InvalidKeyError(f"kty is not found: {self._jwk_data}")20 21        if not algorithm and isinstance(self._jwk_data, dict):22            algorithm = self._jwk_data.get("alg", None)23 24        if not algorithm:25            # Determine alg with kty (and crv).26            crv = self._jwk_data.get("crv", None)27            if kty == "EC":28                if crv == "P-256" or not crv:29                    algorithm = "ES256"30                elif crv == "P-384":31                    algorithm = "ES384"32                elif crv == "P-521":33                    algorithm = "ES512"34                elif crv == "secp256k1":35                    algorithm = "ES256K"36                else:37                    raise InvalidKeyError(f"Unsupported crv: {crv}")38            elif kty == "RSA":39                algorithm = "RS256"40            elif kty == "oct":41                algorithm = "HS256"42            elif kty == "OKP":43                if not crv:44                    raise InvalidKeyError(f"crv is not found: {self._jwk_data}")45                if crv == "Ed25519":46                    algorithm = "EdDSA"47                else:48                    raise InvalidKeyError(f"Unsupported crv: {crv}")49            else:50                raise InvalidKeyError(f"Unsupported kty: {kty}")51 52        if not has_crypto and algorithm in requires_cryptography:53            raise PyJWKError(f"{algorithm} requires 'cryptography' to be installed.")54 55        self.Algorithm = self._algorithms.get(algorithm)56 57        if not self.Algorithm:58            raise PyJWKError(f"Unable to find an algorithm for key: {self._jwk_data}")59 60        self.key = self.Algorithm.from_jwk(self._jwk_data)61 62    @staticmethod63    def from_dict(obj: JWKDict, algorithm: str | None = None) -> "PyJWK":64        return PyJWK(obj, algorithm)65 66    @staticmethod67    def from_json(data: str, algorithm: None = None) -> "PyJWK":68        obj = json.loads(data)69        return PyJWK.from_dict(obj, algorithm)70 71    @property72    def key_type(self) -> str | None:73        return self._jwk_data.get("kty", None)74 75    @property76    def key_id(self) -> str | None:77        return self._jwk_data.get("kid", None)78 79    @property80    def public_key_use(self) -> str | None:81        return self._jwk_data.get("use", None)82 83 84class PyJWKSet:85    def __init__(self, keys: list[JWKDict]) -> None:86        self.keys = []87 88        if not keys:89            raise PyJWKSetError("The JWK Set did not contain any keys")90 91        if not isinstance(keys, list):92            raise PyJWKSetError("Invalid JWK Set value")93 94        for key in keys:95            try:96                self.keys.append(PyJWK(key))97            except PyJWTError:98                # skip unusable keys99                continue100 101        if len(self.keys) == 0:102            raise PyJWKSetError(103                "The JWK Set did not contain any usable keys. Perhaps 'cryptography' is not installed?"104            )105 106    @staticmethod107    def from_dict(obj: dict[str, Any]) -> "PyJWKSet":108        keys = obj.get("keys", [])109        return PyJWKSet(keys)110 111    @staticmethod112    def from_json(data: str) -> "PyJWKSet":113        obj = json.loads(data)114        return PyJWKSet.from_dict(obj)115 116    def __getitem__(self, kid: str) -> "PyJWK":117        for key in self.keys:118            if key.key_id == kid:119                return key120        raise KeyError(f"keyset has no key for kid: {kid}")121 122 123class PyJWTSetWithTimestamp:124    def __init__(self, jwk_set: PyJWKSet):125        self.jwk_set = jwk_set126        self.timestamp = time.monotonic()127 128    def get_jwk_set(self) -> PyJWKSet:129        return self.jwk_set130 131    def get_timestamp(self) -> float:132        return self.timestamp133