"""Pure-Python bridge prover for Apple's trusted-device HSA2 flow.""" from __future__ import annotations import base64 import hashlib import hmac import secrets from dataclasses import dataclass from typing import Optional from cryptography.exceptions import InvalidTag from cryptography.hazmat.primitives.ciphers.aead import AESGCM _SCRYPT_PARAMS = { "n": 16384, "r": 8, "p": 1, "dklen": 64, } _CLIENT_IDENTITY = b"com.apple.security.webprover" _SERVER_IDENTITY = b"com.apple.security.webverifier" _SPAKE2_CONTEXT = b"SPAKE2Web" _KEY_LENGTH = 32 _VERIFIER_KEY_INFO = b"webVerifier" _PROVER_KEY_INFO = b"webProver" _P256_P = int("FFFFFFFF00000001000000000000000000000000FFFFFFFFFFFFFFFFFFFFFFFF", 16) _P256_A = (_P256_P - 3) % _P256_P _P256_B = int("5AC635D8AA3A93E7B3EBBD55769886BC651D06B0CC53B0F63BCE3C3E27D2604B", 16) _P256_ORDER = int( "FFFFFFFF00000000FFFFFFFFFFFFFFFFBCE6FAADA7179E84F3B9CAC2FC632551", 16 ) _P256_GX = int("6B17D1F2E12C4247F8BCE6E563A440F277037D812DEB33A0F4A13945D898C296", 16) _P256_GY = int("4FE342E2FE1A7F9B8EE7EB4A7C0F9E162BCE33576B315ECECBB6406837BF51F5", 16) _SPAKE2_M = "02886e2f97ace46e55ba9dd7242579f2993b64e16ef3dcab95afd497333d8fa12f" _SPAKE2_N = "03d8bbd6c639c62937b04d997f38c3770719c629d7014d49a24b4f98baa1292b49" _AES_GCM_LAYOUTS = {0: (12, 16)} @dataclass(frozen=True) class _Point: """Affine P-256 point used by the bridge SPAKE2 math helpers.""" x: Optional[int] y: Optional[int] @property def is_infinity(self) -> bool: """Return whether this point is the point at infinity.""" return self.x is None or self.y is None _INFINITY = _Point(None, None) _GENERATOR = _Point(_P256_GX, _P256_GY) def _int_to_bytes(value: int, length: Optional[int] = None) -> bytes: """Encode an integer using big-endian bytes.""" if length is None: length = max(1, (value.bit_length() + 7) // 8) return value.to_bytes(length, "big") def _b64_to_bytes(value: str) -> bytes: """Decode a base64 string into raw bytes.""" return base64.b64decode(value.encode("ascii")) def _bytes_to_b64(value: bytes) -> str: """Encode raw bytes as an ASCII base64 string.""" return base64.b64encode(value).decode("ascii") def _encode_point(point: _Point) -> str: """Encode a P-256 point using SEC1 uncompressed point format.""" if point.is_infinity: raise ValueError("Cannot encode the point at infinity.") return "04" + _int_to_bytes(point.x, 32).hex() + _int_to_bytes(point.y, 32).hex() def _decode_point(value: str) -> _Point: """Decode a compressed or uncompressed SEC1 point into affine coordinates.""" raw = bytes.fromhex(value) if len(raw) == 65 and raw[0] == 0x04: point = _Point( int.from_bytes(raw[1:33], "big"), int.from_bytes(raw[33:65], "big"), ) elif len(raw) == 33 and raw[0] in (0x02, 0x03): x_coord = int.from_bytes(raw[1:], "big") rhs = (pow(x_coord, 3, _P256_P) + _P256_A * x_coord + _P256_B) % _P256_P y_coord = pow(rhs, (_P256_P + 1) // 4, _P256_P) if y_coord & 1 != raw[0] & 1: y_coord = (-y_coord) % _P256_P point = _Point(x_coord, y_coord) else: raise ValueError("Unsupported P-256 point encoding.") if not _is_on_curve(point): raise ValueError("Invalid P-256 point.") return point def _is_on_curve(point: _Point) -> bool: """Return whether a point lies on the configured P-256 curve.""" if point.is_infinity: return False assert point.x is not None and point.y is not None return ( pow(point.y, 2, _P256_P) - (pow(point.x, 3, _P256_P) + _P256_A * point.x + _P256_B) ) % _P256_P == 0 def _negate(point: _Point) -> _Point: """Return the additive inverse of a P-256 point.""" if point.is_infinity: return point assert point.x is not None and point.y is not None return _Point(point.x, (-point.y) % _P256_P) def _add_points(left: _Point, right: _Point) -> _Point: """Add two affine P-256 points.""" if left.is_infinity: return right if right.is_infinity: return left assert left.x is not None and left.y is not None assert right.x is not None and right.y is not None if left.x == right.x and (left.y + right.y) % _P256_P == 0: return _INFINITY if left.x == right.x and left.y == right.y: if left.y == 0: return _INFINITY slope = ( (3 * left.x * left.x + _P256_A) * pow(2 * left.y, -1, _P256_P) ) % _P256_P else: slope = ((right.y - left.y) * pow(right.x - left.x, -1, _P256_P)) % _P256_P x_coord = (slope * slope - left.x - right.x) % _P256_P y_coord = (slope * (left.x - x_coord) - left.y) % _P256_P return _Point(x_coord, y_coord) def _multiply_point(point: _Point, scalar: int) -> _Point: """Multiply a P-256 point by a scalar using double-and-add.""" scalar %= _P256_ORDER result = _INFINITY addend = point while scalar: if scalar & 1: result = _add_points(result, addend) addend = _add_points(addend, addend) scalar >>= 1 return result def _concat_length_prefixed(*parts: bytes) -> bytes: """Concatenate transcript parts using the bridge's length-prefixed format.""" output = bytearray() for part in parts: output.extend(len(part).to_bytes(8, "little")) output.extend(part) return bytes(output) def _hkdf_like(ikm: bytes, salt: bytes, info: bytes, length: int) -> bytes: """Derive key material using the bridge worker's HKDF-like expansion.""" hash_len = hashlib.sha256().digest_size if not salt: salt = b"\x00" * hash_len prk = hmac.new(salt, ikm, hashlib.sha256).digest() blocks = bytearray() previous = b"" counter = 1 while len(blocks) < length: previous = hmac.new( prk, previous + info + bytes([counter]), hashlib.sha256, ).digest() blocks.extend(previous) counter += 1 return bytes(blocks[:length]) def _confirmation_key_length(info: bytes, requested_length: int) -> int: """Return the bridge-specific output length for a given HKDF info label.""" if b"ConfirmationKeys" in info: return 64 return requested_length def _derive_key(ikm: bytes, info: bytes, length: int = 64) -> bytes: """Derive one bridge sub-key from raw shared-secret material.""" return _hkdf_like( ikm=ikm, salt=b"", info=info, length=_confirmation_key_length(info, length), ) def _derive_prover_and_verifier_keys(raw_key_hex: str) -> tuple[str, str]: """Split the raw bridge key into prover and verifier AES/HMAC keys.""" raw_key = bytes.fromhex(raw_key_hex) verifier_key = _derive_key(raw_key, _VERIFIER_KEY_INFO, _KEY_LENGTH) prover_key = _derive_key(raw_key, _PROVER_KEY_INFO, _KEY_LENGTH) return verifier_key.hex(), prover_key.hex() @dataclass(frozen=True) class _ClientSharedSecret: """Client-side shared-secret transcript and derived confirmation keys.""" transcript: bytes share_p: str share_v: str def __post_init__(self) -> None: """Derive confirmation keys and the final shared key from the transcript.""" digest = hashlib.sha256(self.transcript).digest() object.__setattr__(self, "_hash_transcript", digest) confirmations = _derive_key(digest, b"ConfirmationKeys", 64) object.__setattr__(self, "_confirm_client", confirmations[:32]) object.__setattr__(self, "_confirm_server", confirmations[32:]) shared_key = _derive_key(digest, b"SharedKey", _KEY_LENGTH) object.__setattr__(self, "_shared_key", shared_key) def get_confirmation(self) -> str: """Return the prover's HMAC confirmation message.""" return hmac.new( self._confirm_client, bytes.fromhex(self.share_v), hashlib.sha256, ).hexdigest() def verify(self, message_hex: str) -> bytes: """Verify the server confirmation and return the shared key bytes.""" expected = hmac.new( self._confirm_server, bytes.fromhex(self.share_p), hashlib.sha256, ).hexdigest() if expected != message_hex: raise ValueError("invalid confirmation from server") return self._shared_key @dataclass(frozen=True) class _ServerSharedSecret: """Server-side shared-secret transcript and derived confirmation keys.""" transcript: bytes share_p: str share_v: str def __post_init__(self) -> None: """Derive confirmation keys and the final shared key from the transcript.""" digest = hashlib.sha256(self.transcript).digest() confirmations = _derive_key(digest, b"ConfirmationKeys", 64) object.__setattr__(self, "_confirm_client", confirmations[:32]) object.__setattr__(self, "_confirm_server", confirmations[32:]) object.__setattr__( self, "_shared_key", _derive_key(digest, b"SharedKey", _KEY_LENGTH), ) def get_confirmation(self) -> str: """Return the verifier's HMAC confirmation message.""" return hmac.new( self._confirm_server, bytes.fromhex(self.share_p), hashlib.sha256, ).hexdigest() def verify(self, message_hex: str) -> bytes: """Verify the prover confirmation and return the shared key bytes.""" expected = hmac.new( self._confirm_client, bytes.fromhex(self.share_v), hashlib.sha256, ).hexdigest() if expected != message_hex: raise ValueError("invalid confirmation from client") return self._shared_key class _ClientHandshake: """Client-side SPAKE2 handshake state for Apple's bridge prover.""" def __init__( self, *, x_scalar: int, w0: int, w1: int, ) -> None: """Initialize the prover handshake with the derived SPAKE2 scalars.""" self._x = x_scalar self._w0 = w0 self._w1 = w1 self._message1_point: Optional[_Point] = None self.share_p: Optional[str] = None def get_message(self) -> str: """Return the prover's first SPAKE2 message.""" point = _add_points( _multiply_point(_GENERATOR, self._x), _multiply_point(_decode_point(_SPAKE2_M), self._w0), ) self._message1_point = point self.share_p = _encode_point(point) return self.share_p def finish(self, server_message_hex: str) -> _ClientSharedSecret: """Finish the handshake using the verifier's first message.""" if self._message1_point is None or self.share_p is None: raise ValueError("get_message must be called before finish") server_point = _decode_point(server_message_hex) if server_point.is_infinity: raise ValueError("invalid curve point") adjusted = _add_points( server_point, _negate(_multiply_point(_decode_point(_SPAKE2_N), self._w0)), ) y_point = _multiply_point(adjusted, self._x) v_point = _multiply_point(adjusted, self._w1) transcript = _concat_length_prefixed( _SPAKE2_CONTEXT, _CLIENT_IDENTITY, _SERVER_IDENTITY, bytes.fromhex(_encode_point(_decode_point(_SPAKE2_M))), bytes.fromhex(_encode_point(_decode_point(_SPAKE2_N))), bytes.fromhex(_encode_point(self._message1_point)), bytes.fromhex(_encode_point(server_point)), bytes.fromhex(_encode_point(y_point)), bytes.fromhex(_encode_point(v_point)), _int_to_bytes(self._w0), ) return _ClientSharedSecret( transcript=transcript, share_p=self.share_p, share_v=server_message_hex, ) class _ServerHandshake: """Server-side SPAKE2 handshake state used by the local test helper.""" def __init__( self, *, y_scalar: int, w0: int, verifier_point: _Point, ) -> None: """Initialize the verifier handshake with its scalar and verifier point.""" self._y = y_scalar self._w0 = w0 self._verifier_point = verifier_point self._message1_point: Optional[_Point] = None self.share_v: Optional[str] = None def get_message(self) -> str: """Return the verifier's first SPAKE2 message.""" point = _add_points( _multiply_point(_GENERATOR, self._y), _multiply_point(_decode_point(_SPAKE2_N), self._w0), ) self._message1_point = point self.share_v = _encode_point(point) return self.share_v def finish(self, client_message_hex: str) -> _ServerSharedSecret: """Finish the verifier handshake using the prover's first message.""" if self._message1_point is None or self.share_v is None: raise ValueError("get_message must be called before finish") client_point = _decode_point(client_message_hex) if client_point.is_infinity: raise ValueError("invalid curve point") adjusted = _add_points( client_point, _negate(_multiply_point(_decode_point(_SPAKE2_M), self._w0)), ) y_point = _multiply_point(adjusted, self._y) verifier_share = _multiply_point(self._verifier_point, self._y) transcript = _concat_length_prefixed( _SPAKE2_CONTEXT, _CLIENT_IDENTITY, _SERVER_IDENTITY, bytes.fromhex(_encode_point(_decode_point(_SPAKE2_M))), bytes.fromhex(_encode_point(_decode_point(_SPAKE2_N))), bytes.fromhex(_encode_point(client_point)), bytes.fromhex(_encode_point(self._message1_point)), bytes.fromhex(_encode_point(y_point)), bytes.fromhex(_encode_point(verifier_share)), _int_to_bytes(self._w0), ) return _ServerSharedSecret( transcript=transcript, share_p=client_message_hex, share_v=self.share_v, ) def _compute_w0_w1(password: str, salt_b64: str) -> tuple[int, int]: """Derive the SPAKE2 scalars from the user code and bridge salt.""" derived = hashlib.scrypt( password.encode("utf-8"), salt=_b64_to_bytes(salt_b64), **_SCRYPT_PARAMS, ) midpoint = len(derived) // 2 return ( int.from_bytes(derived[:midpoint], "big"), int.from_bytes(derived[midpoint:], "big"), ) def _random_nonzero_scalar() -> int: """Return a random scalar in the non-zero P-256 subgroup range.""" scalar = 0 while scalar == 0: scalar = secrets.randbelow(_P256_ORDER) return scalar class TrustedDeviceBridgeProver: """Client-side prover mirroring Apple's prover worker.""" def __init__(self) -> None: """Initialize empty prover state for one bridge verification attempt.""" self._client: Optional[_ClientHandshake] = None self._shared_secret: Optional[_ClientSharedSecret] = None self._raw_key: Optional[str] = None self._verified = False self._verifier_key: Optional[str] = None self._prover_key: Optional[str] = None def init_with_salt(self, salt_b64: str, code: str) -> None: """Initialize the prover with Apple's salt and the user-entered code.""" w0, w1 = _compute_w0_w1(code, salt_b64) self._client = _ClientHandshake( x_scalar=_random_nonzero_scalar(), w0=w0, w1=w1, ) self._shared_secret = None self._raw_key = None self._verified = False self._verifier_key = None self._prover_key = None def get_message1(self) -> str: """Return the prover's first bridge message.""" if self._client is None: raise ValueError("init_with_salt must be called before get_message1") return self._client.get_message() def process_message1(self, message_hex: str) -> str: """Process Apple's first bridge message and return the prover confirmation.""" if self._client is None: raise ValueError("init_with_salt must be called before process_message1") self._shared_secret = self._client.finish(message_hex) return self.get_message2() def get_message2(self) -> str: """Return the prover confirmation generated from the shared transcript.""" if self._shared_secret is None: raise ValueError("process_message1 must be called before get_message2") return self._shared_secret.get_confirmation() def process_message2(self, message_hex: str) -> dict[str, object]: """Verify Apple's confirmation and persist the derived bridge keys.""" if self._shared_secret is None: raise ValueError("process_message1 must be called before process_message2") raw_key = self._shared_secret.verify(message_hex).hex() self._raw_key = raw_key self._verifier_key, self._prover_key = _derive_prover_and_verifier_keys(raw_key) self._verified = True return {"isVerified": True, "key": raw_key} def is_verified(self) -> bool: """Return whether the bridge confirmation exchange has completed.""" return self._verified def get_key(self) -> str: """Return the raw shared bridge key as hexadecimal.""" if self._raw_key is None: raise ValueError("No bridge key is available yet.") return self._raw_key def decrypt_message(self, ciphertext_b64: str) -> str: """Decrypt Apple's final encrypted validation code.""" if self._verifier_key is None: raise ValueError("Bridge verifier key is not available.") try: payload = _b64_to_bytes(ciphertext_b64) version = payload[0] iv_length, tag_length = _AES_GCM_LAYOUTS[version] iv = payload[1 : 1 + iv_length] tag = payload[1 + iv_length : 1 + iv_length + tag_length] ciphertext = payload[1 + iv_length + tag_length :] plaintext = AESGCM(bytes.fromhex(self._verifier_key)).decrypt( iv, ciphertext + tag, bytes([version]), ) return plaintext.decode("utf-8") except (IndexError, KeyError, InvalidTag, UnicodeDecodeError) as exc: raise ValueError("Malformed bridge payload") from exc class _TrustedDeviceBridgeServerProver: """Internal test helper mirroring Apple's server-side bridge flow.""" def __init__(self, *, password: str, salt_b64: str) -> None: """Initialize the local verifier helper with the same password and salt.""" w0, w1 = _compute_w0_w1(password, salt_b64) verifier_point = _multiply_point(_GENERATOR, w1) self._server = _ServerHandshake( y_scalar=_random_nonzero_scalar(), w0=w0, verifier_point=verifier_point, ) self._shared_secret: Optional[_ServerSharedSecret] = None self._raw_key: Optional[str] = None self._verifier_key: Optional[str] = None self._prover_key: Optional[str] = None def get_message1(self) -> str: """Return the verifier's first bridge message.""" return self._server.get_message() def process_message1(self, client_message_hex: str) -> str: """Process the prover message and return the verifier confirmation.""" self._shared_secret = self._server.finish(client_message_hex) return self.get_message2() def get_message2(self) -> str: """Return the verifier confirmation generated from the shared transcript.""" if self._shared_secret is None: raise ValueError("process_message1 must be called before get_message2") return self._shared_secret.get_confirmation() def verify_message2(self, message_hex: str) -> str: """Verify the prover confirmation and persist the derived bridge keys.""" if self._shared_secret is None: raise ValueError("process_message1 must be called before verify_message2") raw_key = self._shared_secret.verify(message_hex).hex() self._raw_key = raw_key self._verifier_key, self._prover_key = _derive_prover_and_verifier_keys(raw_key) return raw_key def encrypt_message(self, plaintext: str) -> str: """Encrypt a plaintext validation code using Apple's AES-GCM payload layout.""" if self._verifier_key is None: raise ValueError("Bridge verifier key is not available.") version = 0 iv_length, tag_length = _AES_GCM_LAYOUTS[version] iv = secrets.token_bytes(iv_length) encrypted = AESGCM(bytes.fromhex(self._verifier_key)).encrypt( iv, plaintext.encode("utf-8"), bytes([version]), ) ciphertext = encrypted[:-tag_length] tag = encrypted[-tag_length:] payload = bytes([version]) + iv + tag + ciphertext return _bytes_to_b64(payload)