first commit

2026-09-14 23:22:10 +09:00
commit 9f9c8fb69b
4207 changed files with 1333187 additions and 0 deletions
@@ -0,0 +1,142 @@
+from typing import Any
+from .models import (
+    BaseJSONEncryption,
+    GeneralJSONEncryption,
+    FlattenedJSONEncryption,
+    Recipient,
+)
+from .registry import JWERegistry
+from .types import (
+    JSONRecipientDict,
+    GeneralJSONSerialization,
+    FlattenedJSONSerialization,
+)
+from ..util import (
+    to_bytes,
+    to_str,
+    json_b64encode,
+    json_b64decode,
+    urlsafe_b64encode,
+    urlsafe_b64decode,
+)
+from .._keys import Key
+
+
+__all__ = [
+    "represent_general_json",
+    "represent_flattened_json",
+    "extract_general_json",
+    "extract_flattened_json",
+]
+
+
+def represent_general_json(obj: GeneralJSONEncryption) -> GeneralJSONSerialization:
+    data: GeneralJSONSerialization = __represent_json_serialization(obj)
+    recipients = []
+    for recipient in obj.recipients:
+        item: JSONRecipientDict = {}
+        assert recipient.header is not None
+        assert recipient.encrypted_key is not None
+        item["header"] = recipient.header
+        item["encrypted_key"] = to_str(urlsafe_b64encode(recipient.encrypted_key))
+        recipients.append(item)
+    data["recipients"] = recipients
+    return data
+
+
+def represent_flattened_json(obj: FlattenedJSONEncryption) -> FlattenedJSONSerialization:
+    data: FlattenedJSONSerialization = __represent_json_serialization(obj)
+    recipient = obj.recipients[0]
+    assert recipient is not None
+    assert recipient.header is not None
+    assert recipient.encrypted_key is not None
+    data["header"] = recipient.header
+    data["encrypted_key"] = to_str(urlsafe_b64encode(recipient.encrypted_key))
+    return data
+
+
+def __represent_json_serialization(obj: BaseJSONEncryption) -> Any:
+    data: dict[str, Any] = {
+        "protected": to_str(json_b64encode(obj.protected)),
+        "iv": to_str(obj.base64_segments["iv"]),
+        "ciphertext": to_str(obj.base64_segments["ciphertext"]),
+        "tag": to_str(obj.base64_segments["tag"]),
+    }
+    if obj.aad:
+        data["aad"] = to_str(urlsafe_b64encode(obj.aad))
+
+    if obj.unprotected:
+        data["unprotected"] = obj.unprotected
+    return data
+
+
+def extract_general_json(data: GeneralJSONSerialization, registry: JWERegistry) -> GeneralJSONEncryption:
+    protected_segment = to_bytes(data["protected"])
+    registry.validate_protected_header_size(protected_segment)
+    protected = json_b64decode(protected_segment)
+
+    unprotected = data.get("unprotected")
+    recipients = data["recipients"]
+    registry.validate_recipient_count(recipients)
+    base64_segments, bytes_segments, aad = __extract_segments(data, registry)
+
+    obj = GeneralJSONEncryption(protected, None, unprotected, aad)
+    obj.base64_segments = base64_segments
+    obj.bytes_segments = bytes_segments
+    for item in recipients:
+        recipient = __extract_recipient(obj, item, registry)
+        obj.recipients.append(recipient)
+    return obj
+
+
+def extract_flattened_json(data: FlattenedJSONSerialization, registry: JWERegistry) -> FlattenedJSONEncryption:
+    protected_segment = to_bytes(data["protected"])
+    registry.validate_protected_header_size(protected_segment)
+    protected = json_b64decode(protected_segment)
+    unprotected = data.get("unprotected")
+    base64_segments, bytes_segments, aad = __extract_segments(data, registry)
+    obj = FlattenedJSONEncryption(protected, None, unprotected, aad)
+    obj.base64_segments = base64_segments
+    obj.bytes_segments = bytes_segments
+    recipient = __extract_recipient(obj, data, registry)
+    obj.recipients.append(recipient)
+    return obj
+
+
+def __extract_segments(
+    data: GeneralJSONSerialization | FlattenedJSONSerialization,
+    registry: JWERegistry,
+) -> tuple[dict[str, bytes], dict[str, bytes], bytes | None]:
+    base64_segments: dict[str, bytes] = {
+        "iv": to_bytes(data["iv"]),
+        "ciphertext": to_bytes(data["ciphertext"]),
+        "tag": to_bytes(data["tag"]),
+    }
+    registry.validate_initialization_vector_size(base64_segments["iv"])
+    registry.validate_ciphertext_size(base64_segments["ciphertext"])
+    registry.validate_auth_tag_size(base64_segments["tag"])
+    bytes_segments: dict[str, bytes] = {
+        "iv": urlsafe_b64decode(base64_segments["iv"]),
+        "ciphertext": urlsafe_b64decode(base64_segments["ciphertext"]),
+        "tag": urlsafe_b64decode(base64_segments["tag"]),
+    }
+    if "aad" in data:
+        base64_segments["aad"] = to_bytes(data["aad"])
+        aad = urlsafe_b64decode(base64_segments["aad"])
+        bytes_segments["aad"] = aad
+    else:
+        aad = None
+    return base64_segments, bytes_segments, aad
+
+
+def __extract_recipient(
+    obj: FlattenedJSONEncryption | GeneralJSONEncryption,
+    data: FlattenedJSONSerialization | JSONRecipientDict,
+    registry: JWERegistry,
+) -> Recipient[Key]:
+    recipient: Recipient[Key] = Recipient(obj, data.get("header"))
+    if "encrypted_key" in data:
+        ek_segment = to_bytes(data["encrypted_key"])
+        registry.validate_encrypted_key_size(ek_segment)
+        recipient.encrypted_key = urlsafe_b64decode(ek_segment)
+    return recipient