codekingpro/portable-devtools

1from authlib.common.encoding import (2    to_bytes,3    to_unicode,4    urlsafe_b64encode,5    json_b64encode,6)7from authlib.jose.util import (8    extract_header,9    extract_segment, ensure_dict,10)11from authlib.jose.errors import (12    DecodeError,13    MissingAlgorithmError,14    UnsupportedAlgorithmError,15    BadSignatureError,16    InvalidHeaderParameterNameError,17)18from .models import JWSHeader, JWSObject19 20 21class JsonWebSignature:22 23    #: Registered Header Parameter Names defined by Section 4.124    REGISTERED_HEADER_PARAMETER_NAMES = frozenset([25        'alg', 'jku', 'jwk', 'kid',26        'x5u', 'x5c', 'x5t', 'x5t#S256',27        'typ', 'cty', 'crit'28    ])29 30    #: Defined available JWS algorithms in the registry31    ALGORITHMS_REGISTRY = {}32 33    def __init__(self, algorithms=None, private_headers=None):34        self._private_headers = private_headers35        self._algorithms = algorithms36 37    @classmethod38    def register_algorithm(cls, algorithm):39        if not algorithm or algorithm.algorithm_type != 'JWS':40            raise ValueError(41                f'Invalid algorithm for JWS, {algorithm!r}')42        cls.ALGORITHMS_REGISTRY[algorithm.name] = algorithm43 44    def serialize_compact(self, protected, payload, key):45        """Generate a JWS Compact Serialization. The JWS Compact Serialization46        represents digitally signed or MACed content as a compact, URL-safe47        string, per `Section 7.1`_.48 49        .. code-block:: text50 51            BASE64URL(UTF8(JWS Protected Header)) || '.' ||52            BASE64URL(JWS Payload) || '.' ||53            BASE64URL(JWS Signature)54 55        :param protected: A dict of protected header56        :param payload: A bytes/string of payload57        :param key: Private key used to generate signature58        :return: byte59        """60        jws_header = JWSHeader(protected, None)61        self._validate_private_headers(protected)62        algorithm, key = self._prepare_algorithm_key(protected, payload, key)63 64        protected_segment = json_b64encode(jws_header.protected)65        payload_segment = urlsafe_b64encode(to_bytes(payload))66 67        # calculate signature68        signing_input = b'.'.join([protected_segment, payload_segment])69        signature = urlsafe_b64encode(algorithm.sign(signing_input, key))70        return b'.'.join([protected_segment, payload_segment, signature])71 72    def deserialize_compact(self, s, key, decode=None):73        """Exact JWS Compact Serialization, and validate with the given key.74        If key is not provided, the returned dict will contain the signature,75        and signing input values. Via `Section 7.1`_.76 77        :param s: text of JWS Compact Serialization78        :param key: key used to verify the signature79        :param decode: a function to decode payload data80        :return: JWSObject81        :raise: BadSignatureError82 83        .. _`Section 7.1`: https://tools.ietf.org/html/rfc7515#section-7.184        """85        try:86            s = to_bytes(s)87            signing_input, signature_segment = s.rsplit(b'.', 1)88            protected_segment, payload_segment = signing_input.split(b'.', 1)89        except ValueError:90            raise DecodeError('Not enough segments')91 92        protected = _extract_header(protected_segment)93        jws_header = JWSHeader(protected, None)94 95        payload = _extract_payload(payload_segment)96        if decode:97            payload = decode(payload)98 99        signature = _extract_signature(signature_segment)100        rv = JWSObject(jws_header, payload, 'compact')101        algorithm, key = self._prepare_algorithm_key(jws_header, payload, key)102        if algorithm.verify(signing_input, signature, key):103            return rv104        raise BadSignatureError(rv)105 106    def serialize_json(self, header_obj, payload, key):107        """Generate a JWS JSON Serialization. The JWS JSON Serialization108        represents digitally signed or MACed content as a JSON object,109        per `Section 7.2`_.110 111        :param header_obj: A dict/list of header112        :param payload: A string/dict of payload113        :param key: Private key used to generate signature114        :return: JWSObject115 116        Example ``header_obj`` of JWS JSON Serialization::117 118            {119                "protected: {"alg": "HS256"},120                "header": {"kid": "jose"}121            }122 123        Pass a dict to generate flattened JSON Serialization, pass a list of124        header dict to generate standard JSON Serialization.125        """126        payload_segment = json_b64encode(payload)127 128        def _sign(jws_header):129            self._validate_private_headers(jws_header)130            _alg, _key = self._prepare_algorithm_key(jws_header, payload, key)131 132            protected_segment = json_b64encode(jws_header.protected)133            signing_input = b'.'.join([protected_segment, payload_segment])134            signature = urlsafe_b64encode(_alg.sign(signing_input, _key))135 136            rv = {137                'protected': to_unicode(protected_segment),138                'signature': to_unicode(signature)139            }140            if jws_header.header is not None:141                rv['header'] = jws_header.header142            return rv143 144        if isinstance(header_obj, dict):145            data = _sign(JWSHeader.from_dict(header_obj))146            data['payload'] = to_unicode(payload_segment)147            return data148 149        signatures = [_sign(JWSHeader.from_dict(h)) for h in header_obj]150        return {151            'payload': to_unicode(payload_segment),152            'signatures': signatures153        }154 155    def deserialize_json(self, obj, key, decode=None):156        """Exact JWS JSON Serialization, and validate with the given key.157        If key is not provided, it will return a dict without signature158        verification. Header will still be validated. Via `Section 7.2`_.159 160        :param obj: text of JWS JSON Serialization161        :param key: key used to verify the signature162        :param decode: a function to decode payload data163        :return: JWSObject164        :raise: BadSignatureError165 166        .. _`Section 7.2`: https://tools.ietf.org/html/rfc7515#section-7.2167        """168        obj = ensure_dict(obj, 'JWS')169 170        payload_segment = obj.get('payload')171        if payload_segment is None:172            raise DecodeError('Missing "payload" value')173 174        payload_segment = to_bytes(payload_segment)175        payload = _extract_payload(payload_segment)176        if decode:177            payload = decode(payload)178 179        if 'signatures' not in obj:180            # flattened JSON JWS181            jws_header, valid = self._validate_json_jws(182                payload_segment, payload, obj, key)183 184            rv = JWSObject(jws_header, payload, 'flat')185            if valid:186                return rv187            raise BadSignatureError(rv)188 189        headers = []190        is_valid = True191        for header_obj in obj['signatures']:192            jws_header, valid = self._validate_json_jws(193                payload_segment, payload, header_obj, key)194            headers.append(jws_header)195            if not valid:196                is_valid = False197 198        rv = JWSObject(headers, payload, 'json')199        if is_valid:200            return rv201        raise BadSignatureError(rv)202 203    def serialize(self, header, payload, key):204        """Generate a JWS Serialization. It will automatically generate a205        Compact or JSON Serialization depending on the given header. If a206        header is in a JSON header format, it will call207        :meth:`serialize_json`, otherwise it will call208        :meth:`serialize_compact`.209 210        :param header: A dict/list of header211        :param payload: A string/dict of payload212        :param key: Private key used to generate signature213        :return: byte/dict214        """215        if isinstance(header, (list, tuple)):216            return self.serialize_json(header, payload, key)217        if 'protected' in header:218            return self.serialize_json(header, payload, key)219        return self.serialize_compact(header, payload, key)220 221    def deserialize(self, s, key, decode=None):222        """Deserialize JWS Serialization, both compact and JSON format.223        It will automatically deserialize depending on the given JWS.224 225        :param s: text of JWS Compact/JSON Serialization226        :param key: key used to verify the signature227        :param decode: a function to decode payload data228        :return: dict229        :raise: BadSignatureError230 231        If key is not provided, it will still deserialize the serialization232        without verification.233        """234        if isinstance(s, dict):235            return self.deserialize_json(s, key, decode)236 237        s = to_bytes(s)238        if s.startswith(b'{') and s.endswith(b'}'):239            return self.deserialize_json(s, key, decode)240        return self.deserialize_compact(s, key, decode)241 242    def _prepare_algorithm_key(self, header, payload, key):243        if 'alg' not in header:244            raise MissingAlgorithmError()245 246        alg = header['alg']247        if self._algorithms is not None and alg not in self._algorithms:248            raise UnsupportedAlgorithmError()249        if alg not in self.ALGORITHMS_REGISTRY:250            raise UnsupportedAlgorithmError()251 252        algorithm = self.ALGORITHMS_REGISTRY[alg]253        if callable(key):254            key = key(header, payload)255        elif key is None and 'jwk' in header:256            key = header['jwk']257        key = algorithm.prepare_key(key)258        return algorithm, key259 260    def _validate_private_headers(self, header):261        # only validate private headers when developers set262        # private headers explicitly263        if self._private_headers is not None:264            names = self.REGISTERED_HEADER_PARAMETER_NAMES.copy()265            names = names.union(self._private_headers)266 267            for k in header:268                if k not in names:269                    raise InvalidHeaderParameterNameError(k)270 271    def _validate_json_jws(self, payload_segment, payload, header_obj, key):272        protected_segment = header_obj.get('protected')273        if not protected_segment:274            raise DecodeError('Missing "protected" value')275 276        signature_segment = header_obj.get('signature')277        if not signature_segment:278            raise DecodeError('Missing "signature" value')279 280        protected_segment = to_bytes(protected_segment)281        protected = _extract_header(protected_segment)282        header = header_obj.get('header')283        if header and not isinstance(header, dict):284            raise DecodeError('Invalid "header" value')285 286        jws_header = JWSHeader(protected, header)287        algorithm, key = self._prepare_algorithm_key(jws_header, payload, key)288        signing_input = b'.'.join([protected_segment, payload_segment])289        signature = _extract_signature(to_bytes(signature_segment))290        if algorithm.verify(signing_input, signature, key):291            return jws_header, True292        return jws_header, False293 294 295def _extract_header(header_segment):296    return extract_header(header_segment, DecodeError)297 298 299def _extract_signature(signature_segment):300    return extract_segment(signature_segment, DecodeError, 'signature')301 302 303def _extract_payload(payload_segment):304    return extract_segment(payload_segment, DecodeError, 'payload')305