From 42462d595d8a1f86760d0d6878de5ac29a392fa0 Mon Sep 17 00:00:00 2001
From: Facundo Tuesca <facundo.tuesca@trailofbits.com>
Date: Tue, 19 Nov 2024 23:34:12 +0100
Subject: [PATCH] attestations: add support for upload from GitLab CI/CD

Signed-off-by: Facundo Tuesca <facundo.tuesca@trailofbits.com>
---
 tests/unit/attestations/test_services.py | 19 +++++++++++++------
 tests/unit/oidc/models/test_gitlab.py    | 19 +++++++++++++++++++
 warehouse/attestations/services.py       |  1 -
 warehouse/oidc/models/gitlab.py          |  9 +++++++++
 4 files changed, 41 insertions(+), 7 deletions(-)

diff --git a/tests/unit/attestations/test_services.py b/tests/unit/attestations/test_services.py
index 793e6a16ea02..4457c98b9d8d 100644
--- a/tests/unit/attestations/test_services.py
+++ b/tests/unit/attestations/test_services.py
@@ -40,8 +40,12 @@ class TestNullIntegrityService:
     def test_interface_matches(self):
         assert verifyClass(IIntegrityService, services.NullIntegrityService)
 
-    def test_build_provenance(self, db_request, dummy_attestation):
-        db_request.oidc_publisher = GitHubPublisherFactory.create()
+    @pytest.mark.parametrize(
+        "publisher_factory",
+        [GitHubPublisherFactory, GitLabPublisherFactory],
+    )
+    def test_build_provenance(self, db_request, dummy_attestation, publisher_factory):
+        db_request.oidc_publisher = publisher_factory.create()
 
         file = FileFactory.create()
         service = services.NullIntegrityService.create_service(None, db_request)
@@ -89,7 +93,6 @@ def test_parse_attestations_fails_no_publisher(self, db_request):
     @pytest.mark.parametrize(
         "publisher_factory",
         [
-            GitLabPublisherFactory,
             GooglePublisherFactory,
             ActiveStatePublisherFactory,
         ],
@@ -267,7 +270,7 @@ def test_parse_attestations_succeeds(
 
     @pytest.mark.parametrize(
         "publisher_factory",
-        [GitHubPublisherFactory],
+        [GitHubPublisherFactory, GitLabPublisherFactory],
     )
     def test_build_provenance_succeeds(
         self, metrics, db_request, publisher_factory, dummy_attestation
@@ -295,8 +298,12 @@ def test_build_provenance_succeeds(
         ]
 
 
-def test_extract_attestations_from_request_empty_list(db_request):
-    db_request.oidc_publisher = GitHubPublisherFactory.create()
+@pytest.mark.parametrize(
+    "publisher_factory",
+    [GitHubPublisherFactory, GitLabPublisherFactory],
+)
+def test_extract_attestations_from_request_empty_list(db_request, publisher_factory):
+    db_request.oidc_publisher = publisher_factory.create()
     db_request.POST = {"attestations": json.dumps([])}
 
     with pytest.raises(
diff --git a/tests/unit/oidc/models/test_gitlab.py b/tests/unit/oidc/models/test_gitlab.py
index 00a247d12bed..ca427789c4e9 100644
--- a/tests/unit/oidc/models/test_gitlab.py
+++ b/tests/unit/oidc/models/test_gitlab.py
@@ -710,6 +710,25 @@ def test_gitlab_publisher_verify_url(
         )
         assert publisher.verify_url(url) == expected
 
+    @pytest.mark.parametrize("environment", ["", "some-env"])
+    def test_gitlab_publisher_attestation_identity(self, environment):
+        publisher = gitlab.GitLabPublisher(
+            project="project",
+            namespace="group/subgroup",
+            workflow_filepath="workflow_filename.yml",
+            environment=environment,
+        )
+
+        identity = publisher.attestation_identity
+        assert identity is not None
+        assert identity.repository == publisher.project_path
+        assert identity.workflow_filepath == publisher.workflow_filepath
+
+        if not environment:
+            assert identity.environment is None
+        else:
+            assert identity.environment == publisher.environment
+
 
 class TestPendingGitLabPublisher:
     def test_reify_does_not_exist_yet(self, db_request):
diff --git a/warehouse/attestations/services.py b/warehouse/attestations/services.py
index ae8b8399a313..10e9f00378dd 100644
--- a/warehouse/attestations/services.py
+++ b/warehouse/attestations/services.py
@@ -151,7 +151,6 @@ def parse_attestations(
         artifact. Attestations are only allowed when uploading via a Trusted
         Publisher, because a Trusted Publisher provides the identity that will be
         used to verify the attestations.
-        Only GitHub Actions Trusted Publishers are supported.
         """
 
         attestations = _extract_attestations_from_request(request)
diff --git a/warehouse/oidc/models/gitlab.py b/warehouse/oidc/models/gitlab.py
index 2bc0bce8523a..22b4e58fe676 100644
--- a/warehouse/oidc/models/gitlab.py
+++ b/warehouse/oidc/models/gitlab.py
@@ -14,6 +14,7 @@
 
 from typing import Any
 
+from pypi_attestations import GitLabPublisher as GitLabIdentity, Publisher
 from sqlalchemy import ForeignKey, String, UniqueConstraint
 from sqlalchemy.dialects.postgresql import UUID
 from sqlalchemy.orm import Query, mapped_column
@@ -258,6 +259,14 @@ def publisher_url(self, claims=None):
         base = self.publisher_base_url
         return f"{base}/commit/{claims['sha']}" if claims else base
 
+    @property
+    def attestation_identity(self) -> Publisher | None:
+        return GitLabIdentity(
+            repository=self.project_path,
+            workflow_filepath=self.workflow_filepath,
+            environment=self.environment if self.environment else None,
+        )
+
     def stored_claims(self, claims=None):
         claims = claims if claims else {}
         return {"ref_path": claims.get("ref_path"), "sha": claims.get("sha")}