1# Copyright 2021 Google LLC 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14"""A JWT HMAC key manager.""" 15 16from typing import Optional, Type 17 18from tink.proto import jwt_hmac_pb2 19from tink.proto import tink_pb2 20from tink import core 21from tink.cc.pybind import tink_bindings 22from tink.jwt import _json_util 23from tink.jwt import _jwt_error 24from tink.jwt import _jwt_format 25from tink.jwt import _jwt_mac 26from tink.jwt import _jwt_validator 27from tink.jwt import _raw_jwt 28from tink.jwt import _verified_jwt 29 30_JWT_HMAC_KEY_TYPE = 'type.googleapis.com/google.crypto.tink.JwtHmacKey' 31 32_ALGORITHM_STRING = { 33 jwt_hmac_pb2.HS256: 'HS256', 34 jwt_hmac_pb2.HS384: 'HS384', 35 jwt_hmac_pb2.HS512: 'HS512' 36} 37 38 39class _JwtHmac(_jwt_mac.JwtMacInternal): 40 """Interface for authenticating and verifying JWT with JWS MAC.""" 41 42 def __init__(self, cc_mac: tink_bindings.Mac, algorithm: str, 43 custom_kid: Optional[str]): 44 self._cc_mac = cc_mac 45 self._algorithm = algorithm 46 self._custom_kid = custom_kid 47 48 @core.use_tink_errors 49 def _compute_mac(self, data: bytes) -> bytes: 50 return self._cc_mac.compute_mac(data) 51 52 @core.use_tink_errors 53 def _verify_mac(self, mac_value: bytes, data: bytes) -> None: 54 self._cc_mac.verify_mac(mac_value, data) 55 56 def compute_mac_and_encode_with_kid(self, raw_jwt: _raw_jwt.RawJwt, 57 kid: Optional[str]) -> str: 58 """Computes a MAC and encodes the token. 59 60 Args: 61 raw_jwt: The RawJwt token to be MACed and encoded. 62 kid: Optional "kid" header value. It is set by the wrapper for keys with 63 output prefix TINK, and it is None for output prefix RAW. 64 65 Returns: 66 The MACed token encoded in the JWS compact serialization format. 67 Raises: 68 tink.TinkError if the operation fails. 69 """ 70 if self._custom_kid is not None: 71 if kid is not None: 72 raise _jwt_error.JwtInvalidError( 73 'custom_kid must not be set for keys with output prefix type TINK') 74 kid = self._custom_kid 75 unsigned = _jwt_format.create_unsigned_compact(self._algorithm, kid, 76 raw_jwt) 77 return _jwt_format.create_signed_compact(unsigned, 78 self._compute_mac(unsigned)) 79 80 def verify_mac_and_decode_with_kid( 81 self, compact: str, validator: _jwt_validator.JwtValidator, 82 kid: Optional[str]) -> _verified_jwt.VerifiedJwt: 83 """Verifies, validates and decodes a MACed compact JWT token.""" 84 parts = _jwt_format.split_signed_compact(compact) 85 unsigned_compact, json_header, json_payload, mac = parts 86 self._verify_mac(mac, unsigned_compact) 87 header = _json_util.json_loads(json_header) 88 _jwt_format.validate_header( 89 header=header, 90 algorithm=self._algorithm, 91 tink_kid=kid, 92 custom_kid=self._custom_kid) 93 raw_jwt = _raw_jwt.raw_jwt_from_json( 94 _jwt_format.get_type_header(header), json_payload) 95 _jwt_validator.validate(validator, raw_jwt) 96 return _verified_jwt.VerifiedJwt._create(raw_jwt) # pylint: disable=protected-access 97 98 99class MacCcToPyJwtMacKeyManager(core.KeyManager[_jwt_mac.JwtMacInternal]): 100 """Transforms C++ KeyManager into a Python KeyManager.""" 101 102 def __init__(self): 103 self._cc_key_manager = tink_bindings.MacKeyManager.from_cc_registry( 104 'type.googleapis.com/google.crypto.tink.JwtHmacKey') 105 106 def primitive_class(self) -> Type[_jwt_mac.JwtMacInternal]: 107 return _jwt_mac.JwtMacInternal 108 109 @core.use_tink_errors 110 def primitive(self, key_data: tink_pb2.KeyData) -> _jwt_mac.JwtMacInternal: 111 if key_data.type_url != _JWT_HMAC_KEY_TYPE: 112 raise _jwt_error.JwtInvalidError('Invalid key data key type') 113 jwt_hmac_key = jwt_hmac_pb2.JwtHmacKey.FromString(key_data.value) 114 algorithm = _ALGORITHM_STRING[jwt_hmac_key.algorithm] 115 cc_mac = self._cc_key_manager.primitive(key_data.SerializeToString()) 116 if jwt_hmac_key.HasField('custom_kid'): 117 custom_kid = jwt_hmac_key.custom_kid.value 118 else: 119 custom_kid = None 120 return _JwtHmac(cc_mac, algorithm, custom_kid) 121 122 def key_type(self) -> str: 123 return self._cc_key_manager.key_type() 124 125 @core.use_tink_errors 126 def new_key_data(self, 127 key_template: tink_pb2.KeyTemplate) -> tink_pb2.KeyData: 128 data = self._cc_key_manager.new_key_data(key_template.SerializeToString()) 129 return tink_pb2.KeyData.FromString(data) 130 131 132def register(): 133 tink_bindings.register_jwt() 134 core.Registry.register_key_manager( 135 MacCcToPyJwtMacKeyManager(), new_key_allowed=True) 136