From 8fc177d87986d53eeaea80e76e309101e66298c1 Mon Sep 17 00:00:00 2001
From: Facundo Tuesca <facundo.tuesca@trailofbits.com>
Date: Thu, 2 May 2024 21:01:47 +0200
Subject: [PATCH] Add AttestationPayload model

---
 pyproject.toml                          |  3 ++-
 src/pypi_attestation_models/__init__.py |  2 ++
 src/pypi_attestation_models/_impl.py    | 35 +++++++++++++++++++++++--
 test/test_impl.py                       | 12 +++++++++
 4 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index ce45190..ecca645 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -73,10 +73,11 @@ target-version = "py39"
 
 [tool.ruff.lint]
 select = ["ALL"]
+# ANN102 is deprecated
 # D203 and D213 are incompatible with D211 and D212 respectively.
 # COM812 and ISC001 can cause conflicts when using ruff as a formatter.
 # See https://docs.astral.sh/ruff/formatter/#conflicting-lint-rules.
-ignore = ["D203", "D213", "COM812", "ISC001"]
+ignore = ["ANN102", "D203", "D213", "COM812", "ISC001"]
 
 [tool.ruff.lint.per-file-ignores]
 
diff --git a/src/pypi_attestation_models/__init__.py b/src/pypi_attestation_models/__init__.py
index b3ba743..cd34895 100644
--- a/src/pypi_attestation_models/__init__.py
+++ b/src/pypi_attestation_models/__init__.py
@@ -4,6 +4,7 @@
 
 from ._impl import (
     Attestation,
+    AttestationPayload,
     ConversionError,
     InvalidAttestationError,
     VerificationMaterial,
@@ -13,6 +14,7 @@
 
 __all__ = [
     "Attestation",
+    "AttestationPayload",
     "ConversionError",
     "InvalidAttestationError",
     "VerificationMaterial",
diff --git a/src/pypi_attestation_models/_impl.py b/src/pypi_attestation_models/_impl.py
index 75f8dd8..1f03787 100644
--- a/src/pypi_attestation_models/_impl.py
+++ b/src/pypi_attestation_models/_impl.py
@@ -7,14 +7,19 @@
 
 import binascii
 from base64 import b64decode, b64encode
-from typing import Annotated, Any, Literal, NewType
+from hashlib import sha256
+from typing import TYPE_CHECKING, Annotated, Any, Literal, NewType
 
+import rfc8785
 from annotated_types import MinLen  # noqa: TCH002
 from cryptography import x509
 from cryptography.hazmat.primitives import serialization
 from pydantic import BaseModel
 from sigstore.models import Bundle, LogEntry
 
+if TYPE_CHECKING:
+    from pathlib import Path  # pragma: no cover
+
 
 class ConversionError(ValueError):
     """The base error for all errors during conversion."""
@@ -62,8 +67,34 @@ class Attestation(BaseModel):
     message_signature: str
     """
     The attestation's signature, as `base64(raw-sig)`, where `raw-sig`
-    is the raw bytes of the signing operation.
+    is the raw bytes of the signing operation over the attestation payload.
+    """
+
+
+class AttestationPayload(BaseModel):
+    """Attestation Payload object as defined in PEP 740."""
+
+    distribution: str
     """
+    The file name of the Python package distribution.
+    """
+
+    digest: str
+    """
+    The SHA-256 digest of the distribution's contents, as a hexadecimal string.
+    """
+
+    @classmethod
+    def from_dist(cls, dist: Path) -> AttestationPayload:
+        """Create an `AttestationPayload` from a distribution file."""
+        return AttestationPayload(
+            distribution=dist.name,
+            digest=sha256(dist.read_bytes()).hexdigest(),
+        )
+
+    def __bytes__(self: AttestationPayload) -> bytes:
+        """Convert to bytes using a canonicalized JSON representation (from RFC8785)."""
+        return rfc8785.dumps(self.model_dump())
 
 
 def sigstore_to_pypi(sigstore_bundle: Bundle) -> Attestation:
diff --git a/test/test_impl.py b/test/test_impl.py
index 1516a0d..bd45c44 100644
--- a/test/test_impl.py
+++ b/test/test_impl.py
@@ -1,5 +1,6 @@
 """Internal implementation tests."""
 
+import hashlib
 import json
 from pathlib import Path
 
@@ -80,3 +81,14 @@ def test_verification_roundtrip() -> None:
                 identity="facundo.tuesca@trailofbits.com", issuer="https://accounts.google.com"
             ),
         )
+
+
+def test_attestation_payload() -> None:
+    payload = impl.AttestationPayload.from_dist(artifact_path)
+
+    assert payload.digest == hashlib.sha256(artifact_path.read_bytes()).hexdigest()
+    assert payload.distribution == artifact_path.name
+
+    expected = f'{{"digest":"{payload.digest}","distribution":"{payload.distribution}"}}'
+
+    assert bytes(payload) == bytes(expected, "utf-8")