PSP レコードの復号

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140import base64
import hashlib
import hmac
import json

import rfc8785
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import ec, rsa
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from pyhpke import AEADId, CipherSuite, KDFId, KEMId, KEMKey, KEMKeyPair


def recover_record(bundle_json, capsule, original_key_id, original_public_key_pem):
    """Return plaintext from exact inner bundle bytes and a decrypted capsule.

    Use the documented schema and trusted original key ID/RSA public PEM.
    No cloud calls occur here. The RSA private key stays in the cloud.
    """
    record, access, runtime, nonce, ciphertext = validate_bundle(
        bundle_json, original_key_id, original_public_key_pem
    )
    e = recover_scalar(record, access, capsule)

    ephemeral = ec.derive_private_key(e, ec.SECP256R1())
    pair = KEMKeyPair(
        KEMKey.from_pyca_cryptography_key(ephemeral),
        KEMKey.from_pyca_cryptography_key(ephemeral.public_key()),
    )
    suite = CipherSuite.new(
        KEMId.DHKEM_P256_HKDF_SHA256, KDFId.HKDF_SHA256, AEADId.EXPORT_ONLY
    )
    enc, sender = suite.create_sender_context(
        KEMKey.from_pyca_cryptography_key(runtime),
        info=b"OpenAI-PSP-v4/HPKE-export",
        eks=pair,
    )
    require(enc == decode64(record["encrypted_dek"]), "Wrong recovered scalar")
    aad = rfc8785.dumps({k: v for k, v in record.items() if k != "ciphertext"})
    key = sender.export(b"OpenAI-PSP-v4/payload-key" + hashlib.sha256(aad).digest(), 32)
    return AESGCM(key).decrypt(nonce, ciphertext, aad).decode("utf-8")


def validate_bundle(bundle_json, original_key_id, original_public_key_pem):
    bundle = json.loads(bundle_json.decode("utf-8"), object_pairs_hook=unique_object)
    require(
        type(bundle["bundle_version"]) is int and bundle["bundle_version"] == 1,
        "Unsupported bundle version",
    )
    record, access = bundle["protected_record"], bundle["customer_access"]
    require(access is not None, "Customer recovery is unavailable")
    require(
        record["version"] == "v4"
        and record["algorithm"]
        == "HPKE-v1-DHKEM_P256_HKDF_SHA256-EXPORT_ONLY-AES256_GCM",
        "Unsupported record",
    )
    require(
        type(access["version"]) is int
        and access["version"] == 1
        and access["algorithm"] == "RSA-OAEP-256",
        "Unsupported recovery capsule",
    )
    require(access["customer_key_id"] == original_key_id, "Wrong key registration")
    require(len(decode64(access["recovery_secret"])) == 384, "Wrong RSA length")

    public_key = serialization.load_pem_public_key(original_public_key_pem)
    require(isinstance(public_key, rsa.RSAPublicKey), "Expected an RSA public key")
    require(public_key.key_size == 3072, "Expected RSA-3072")
    spki = public_key.public_bytes(
        serialization.Encoding.DER, serialization.PublicFormat.SubjectPublicKeyInfo
    )
    fingerprint = decode64(access["customer_public_key_fingerprint"])
    require(
        hmac.compare_digest(fingerprint, hashlib.sha256(spki).digest()), "Wrong key"
    )
    runtime = point(record["safety_review_runtime_public_key"])
    coordinates = runtime.public_numbers()
    jwk = {
        "crv": "P-256",
        "kty": "EC",
        "x": b64url(coordinates.x.to_bytes(32, "big")),
        "y": b64url(coordinates.y.to_bytes(32, "big")),
    }
    thumbprint = b64url(hashlib.sha256(rfc8785.dumps(jwk)).digest())
    require(record["key_thumbprint"] == thumbprint, "Wrong runtime key")
    nonce, ciphertext = decode64(record["nonce"]), decode64(record["ciphertext"])
    require(len(nonce) == 12 and len(ciphertext) >= 16, "Invalid AES-GCM encoding")
    return record, access, runtime, nonce, ciphertext


def recover_scalar(record, access, capsule):
    header = {k: v for k, v in access.items() if k != "recovery_secret"}
    binding_input = {
        "purpose": "OpenAI-PSP-v4/customer-recovery-binding",
        "bundle_version": 1,
        "record": record,
        "recovery_header": header,
    }
    binding = hashlib.sha256(rfc8785.dumps(binding_input)).digest()
    require(
        isinstance(capsule, bytes)
        and len(capsule) == 72
        and capsule[:8] == b"PSPv4\x00\x00\x01"
        and hmac.compare_digest(capsule[40:], binding),
        "Invalid recovery capsule binding",
    )
    e = int.from_bytes(capsule[8:40], "big")
    order = 0xFFFFFFFF00000000FFFFFFFFFFFFFFFFBCE6FAADA7179E84F3B9CAC2FC632551
    require(1 <= e < order, "Invalid P-256 private scalar")
    return e


def require(condition, message):
    if not condition:
        raise ValueError(message)


def unique_object(pairs):
    require(len(pairs) == len(dict(pairs)), "Duplicate JSON field")
    return dict(pairs)


def b64url(data):
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")


def decode64(value):
    # Restore omitted "=" padding: -len(value) % 4 is how many characters
    # are needed to reach a multiple of four (e.g. length 10 needs "==").
    data = base64.b64decode(
        value + "=" * (-len(value) % 4), altchars=b"-_", validate=True
    )
    require(b64url(data) == value, "Noncanonical base64url")
    return data


def point(value):
    encoded = decode64(value)
    require(len(encoded) == 65 and encoded[0] == 4, "Invalid P-256 encoding")
    return ec.EllipticCurvePublicKey.from_encoded_point(ec.SECP256R1(), encoded)