feat(relocation): Allow relocation artifacts to be downloaded (#69444)

azaslavsky · web-flow · commit 7323263194ea · 2024-04-23T20:27:00.000Z
If the relocation artifact is a `.tar` file, it is first decrypted.

This ability is only available to Superusers/staff who have the
`relocation.admin` permission.
diff --git a/src/sentry/api/endpoints/relocations/artifacts/details.py b/src/sentry/api/endpoints/relocations/artifacts/details.py
@@ -0,0 +1,92 @@
+import logging
+
+from cryptography.fernet import Fernet
+from rest_framework.exceptions import PermissionDenied
+from rest_framework.request import Request
+from rest_framework.response import Response
+
+from sentry.api.api_owners import ApiOwner
+from sentry.api.api_publish_status import ApiPublishStatus
+from sentry.api.base import Endpoint, region_silo_endpoint
+from sentry.api.exceptions import ResourceDoesNotExist, StaffRequired, SuperuserRequired
+from sentry.api.permissions import SuperuserOrStaffFeatureFlaggedPermission
+from sentry.auth.elevated_mode import has_elevated_mode
+from sentry.auth.staff import has_staff_option
+from sentry.backup.crypto import (
+    GCPKMSDecryptor,
+    get_default_crypto_key_version,
+    unwrap_encrypted_export_tarball,
+)
+from sentry.models.files.utils import get_relocation_storage
+from sentry.models.relocation import Relocation
+from sentry.utils import json
+
+ERR_NEED_RELOCATION_ADMIN = (
+    "Cannot view relocation artifacts, as you do not have the appropriate permissions."
+)
+
+logger = logging.getLogger(__name__)
+
+
+@region_silo_endpoint
+class RelocationArtifactDetailsEndpoint(Endpoint):
+    owner = ApiOwner.OPEN_SOURCE
+    publish_status = {
+        # TODO(getsentry/team-ospo#214): Stabilize before GA.
+        "GET": ApiPublishStatus.EXPERIMENTAL,
+    }
+    permission_classes = (SuperuserOrStaffFeatureFlaggedPermission,)
+
+    def get(
+        self, request: Request, relocation_uuid: str, artifact_kind: str, file_name: str
+    ) -> Response:
+        """
+        Get a single relocation artifact.
+        ``````````````````````````````````````````````````
+
+        :pparam string relocation_uuid: a UUID identifying the relocation.
+        :pparam string artifact_kind: one of `conf` | `in` | `out` | `findings`.
+        :pparam string file_name: The name of the file itself.
+
+        :auth: required
+        """
+
+        logger.info("relocations.artifact.details.get.start", extra={"caller": request.user.id})
+
+        # TODO(schew2381): Remove the superuser reference below after feature flag is removed.
+        # Must be superuser/staff AND have a `UserPermission` of `relocation.admin` to see access!
+        if not has_elevated_mode(request):
+            if has_staff_option(request.user):
+                raise StaffRequired
+            raise SuperuserRequired
+
+        if not request.access.has_permission("relocation.admin"):
+            raise PermissionDenied(
+                "Cannot view relocation artifacts, as you do not have the appropriate permissions."
+            )
+
+        try:
+            relocation: Relocation = Relocation.objects.get(uuid=relocation_uuid)
+        except Relocation.DoesNotExist:
+            raise ResourceDoesNotExist
+
+        file_path = f"runs/{relocation.uuid}/{artifact_kind}/{file_name}"
+        relocation_storage = get_relocation_storage()
+        if not relocation_storage.exists(file_path):
+            raise ResourceDoesNotExist
+
+        # TODO(azaslavsky): We can probably get all clever and stream these files, but it's not
+        # necessary for now.
+        with relocation_storage.open(file_path) as fp:
+            if not file_name.endswith(".tar"):
+                return self.respond({"contents": fp.read()})
+
+            unwrapped = unwrap_encrypted_export_tarball(fp)
+            decryptor = GCPKMSDecryptor.from_bytes(
+                json.dumps(get_default_crypto_key_version()).encode("utf-8")
+            )
+            plaintext_data_encryption_key = decryptor.decrypt_data_encryption_key(unwrapped)
+            fernet = Fernet(plaintext_data_encryption_key)
+            return self.respond(
+                {"contents": fernet.decrypt(unwrapped.encrypted_json_blob).decode("utf-8")}
+            )
diff --git a/src/sentry/api/urls.py b/src/sentry/api/urls.py
@@ -41,6 +41,7 @@
     ReleaseThresholdStatusIndexEndpoint,
 )
 from sentry.api.endpoints.relocations.abort import RelocationAbortEndpoint
+from sentry.api.endpoints.relocations.artifacts.details import RelocationArtifactDetailsEndpoint
 from sentry.api.endpoints.relocations.artifacts.index import RelocationArtifactIndexEndpoint
 from sentry.api.endpoints.relocations.cancel import RelocationCancelEndpoint
 from sentry.api.endpoints.relocations.details import RelocationDetailsEndpoint
@@ -870,6 +871,11 @@ def create_group_urls(name_prefix: str) -> list[URLPattern | URLResolver]:
         RelocationArtifactIndexEndpoint.as_view(),
         name="sentry-api-0-relocations-artifacts-index",
     ),
+    re_path(
+        r"^(?P<relocation_uuid>[^\/]+)/artifacts/(?P<artifact_kind>[^\/]+)/(?P<file_name>[^\/]+)$",
+        RelocationArtifactDetailsEndpoint.as_view(),
+        name="sentry-api-0-relocations-artifacts-details",
+    ),
 ]
 
 RELAY_URLS = [
diff --git a/tests/sentry/api/endpoints/relocations/artifacts/test_details.py b/tests/sentry/api/endpoints/relocations/artifacts/test_details.py
@@ -0,0 +1,243 @@
+from datetime import datetime, timezone
+from io import BytesIO, StringIO
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from types import SimpleNamespace
+from unittest.mock import patch
+from uuid import uuid4
+
+from google_crc32c import value as crc32c
+
+from sentry.api.endpoints.relocations.artifacts.index import ERR_NEED_RELOCATION_ADMIN
+from sentry.backup.crypto import (
+    LocalFileDecryptor,
+    LocalFileEncryptor,
+    create_encrypted_export_tarball,
+    unwrap_encrypted_export_tarball,
+)
+from sentry.models.files.utils import get_relocation_storage
+from sentry.models.relocation import Relocation
+from sentry.testutils.cases import APITestCase
+from sentry.testutils.helpers.backups import FakeKeyManagementServiceClient, generate_rsa_key_pair
+from sentry.testutils.helpers.options import override_options
+from sentry.utils.relocation import OrderedTask
+
+TEST_DATE_ADDED = datetime(2023, 1, 23, 1, 23, 45, tzinfo=timezone.utc)
+RELOCATION_ADMIN_PERMISSION = "relocation.admin"
+
+
+class GetRelocationArtifactDetailsTest(APITestCase):
+    endpoint = "sentry-api-0-relocations-artifacts-details"
+    method = "GET"
+
+    def setUp(self):
+        super().setUp()
+        self.owner = self.create_user(email="owner@example.com", is_superuser=False, is_staff=False)
+        self.superuser = self.create_user(is_superuser=True)
+        self.staff_user = self.create_user(is_staff=True)
+        self.relocation: Relocation = Relocation.objects.create(
+            date_added=TEST_DATE_ADDED,
+            creator_id=self.owner.id,
+            owner_id=self.owner.id,
+            status=Relocation.Status.PAUSE.value,
+            step=Relocation.Step.PREPROCESSING.value,
+            want_org_slugs=["foo"],
+            want_usernames=["alice", "bob"],
+            latest_notified=Relocation.EmailKind.STARTED.value,
+            latest_task=OrderedTask.PREPROCESSING_SCAN.name,
+            latest_task_attempts=1,
+        )
+
+
+class GetRelocationArtifactDetailsGoodTest(GetRelocationArtifactDetailsTest):
+    def setUp(self):
+        super().setUp()
+        dir = f"runs/{self.relocation.uuid}"
+        self.relocation_storage = get_relocation_storage()
+
+        # These files are unencrypted, so just save the file name as the content for testing
+        # purposes.
+        self.relocation_storage.save(
+            f"{dir}/somedir/file.json", StringIO(f'"{dir}/somedir/file.json"')
+        )
+
+        # `.tar` files should be encrypted.
+        with TemporaryDirectory() as tmp_dir:
+            (priv_key_pem, pub_key_pem) = generate_rsa_key_pair()
+            tmp_priv_key_path = Path(tmp_dir).joinpath("key")
+            self.priv_key_pem = priv_key_pem
+            with open(tmp_priv_key_path, "wb") as f:
+                f.write(priv_key_pem)
+
+            tmp_pub_key_path = Path(tmp_dir).joinpath("key.pub")
+            self.pub_key_pem = pub_key_pem
+            with open(tmp_pub_key_path, "wb") as f:
+                f.write(pub_key_pem)
+
+            with open(tmp_pub_key_path, "rb") as p:
+                self.tarball = create_encrypted_export_tarball(
+                    f"{dir}/encrypted/file.tar", LocalFileEncryptor(p)
+                ).getvalue()
+                self.relocation_storage.save(f"{dir}/encrypted/file.tar", BytesIO(self.tarball))
+
+    def mock_kms_client(self, fake_kms_client: FakeKeyManagementServiceClient):
+        fake_kms_client.asymmetric_decrypt.call_count = 0
+        fake_kms_client.get_public_key.call_count = 0
+
+        unwrapped = unwrap_encrypted_export_tarball(BytesIO(self.tarball))
+        plaintext_dek = LocalFileDecryptor.from_bytes(
+            self.priv_key_pem
+        ).decrypt_data_encryption_key(unwrapped)
+
+        fake_kms_client.asymmetric_decrypt.return_value = SimpleNamespace(
+            plaintext=plaintext_dek,
+            plaintext_crc32c=crc32c(plaintext_dek),
+        )
+        fake_kms_client.asymmetric_decrypt.side_effect = None
+
+        fake_kms_client.get_public_key.return_value = SimpleNamespace(
+            pem=self.pub_key_pem.decode("utf-8")
+        )
+        fake_kms_client.get_public_key.side_effect = None
+
+    @patch(
+        "sentry.backup.crypto.KeyManagementServiceClient",
+        new_callable=lambda: FakeKeyManagementServiceClient,
+    )
+    def test_good_unencrypted_with_superuser(
+        self, fake_kms_client: FakeKeyManagementServiceClient
+    ) -> None:
+        self.mock_kms_client(fake_kms_client)
+        self.add_user_permission(self.superuser, RELOCATION_ADMIN_PERMISSION)
+        self.login_as(user=self.superuser, superuser=True)
+        response = self.get_success_response(str(self.relocation.uuid), "somedir", "file.json")
+
+        assert fake_kms_client.asymmetric_decrypt.call_count == 0
+        assert (
+            response.data["contents"] == f'"runs/{self.relocation.uuid}/somedir/file.json"'.encode()
+        )
+
+    @patch(
+        "sentry.backup.crypto.KeyManagementServiceClient",
+        new_callable=lambda: FakeKeyManagementServiceClient,
+    )
+    def test_good_encrypted_with_superuser(
+        self, fake_kms_client: FakeKeyManagementServiceClient
+    ) -> None:
+        self.mock_kms_client(fake_kms_client)
+        self.add_user_permission(self.superuser, RELOCATION_ADMIN_PERMISSION)
+        self.login_as(user=self.superuser, superuser=True)
+        response = self.get_success_response(str(self.relocation.uuid), "encrypted", "file.tar")
+
+        assert fake_kms_client.asymmetric_decrypt.call_count == 1
+        assert str(response.data["contents"]) == f'"runs/{self.relocation.uuid}/encrypted/file.tar"'
+
+    @override_options({"staff.ga-rollout": True})
+    @patch(
+        "sentry.backup.crypto.KeyManagementServiceClient",
+        new_callable=lambda: FakeKeyManagementServiceClient,
+    )
+    def test_good_unencrypted_with_staff(
+        self, fake_kms_client: FakeKeyManagementServiceClient
+    ) -> None:
+        self.mock_kms_client(fake_kms_client)
+        self.add_user_permission(self.staff_user, RELOCATION_ADMIN_PERMISSION)
+        self.login_as(user=self.staff_user, staff=True)
+        response = self.get_success_response(str(self.relocation.uuid), "somedir", "file.json")
+
+        assert fake_kms_client.asymmetric_decrypt.call_count == 0
+        assert (
+            response.data["contents"] == f'"runs/{self.relocation.uuid}/somedir/file.json"'.encode()
+        )
+
+    @override_options({"staff.ga-rollout": True})
+    @patch(
+        "sentry.backup.crypto.KeyManagementServiceClient",
+        new_callable=lambda: FakeKeyManagementServiceClient,
+    )
+    def test_good_encrypted_with_staff(
+        self, fake_kms_client: FakeKeyManagementServiceClient
+    ) -> None:
+        self.mock_kms_client(fake_kms_client)
+        self.add_user_permission(self.staff_user, RELOCATION_ADMIN_PERMISSION)
+        self.login_as(user=self.staff_user, staff=True)
+        response = self.get_success_response(str(self.relocation.uuid), "encrypted", "file.tar")
+
+        assert fake_kms_client.asymmetric_decrypt.call_count == 1
+        assert str(response.data["contents"]) == f'"runs/{self.relocation.uuid}/encrypted/file.tar"'
+
+
+class GetRelocationArtifactDetailsBadTest(GetRelocationArtifactDetailsTest):
+    def setUp(self):
+        super().setUp()
+        dir = f"runs/{self.relocation.uuid}"
+        self.relocation_storage = get_relocation_storage()
+
+        # These files are unencrypted, so just save the file name as the content for testing
+        # purposes.
+        self.relocation_storage.save(
+            f"{dir}/somedir/file.json", StringIO(f'"{dir}/somedir/file.json"')
+        )
+
+    @override_options({"staff.ga-rollout": True})
+    def test_bad_unprivileged_user(self):
+        self.login_as(user=self.owner, superuser=False, staff=False)
+
+        # Ensures we don't reveal existence info to improperly authenticated users.
+        does_not_exist_uuid = uuid4().hex
+        self.get_error_response(str(does_not_exist_uuid), "somedir", "file.json", status_code=403)
+
+    def test_bad_superuser_disabled(self):
+        self.add_user_permission(self.superuser, RELOCATION_ADMIN_PERMISSION)
+        self.login_as(user=self.superuser, superuser=False)
+
+        # Ensures we don't reveal existence info to improperly authenticated users.
+        does_not_exist_uuid = uuid4().hex
+        self.get_error_response(str(does_not_exist_uuid), "somedir", "file.json", status_code=403)
+
+    @override_options({"staff.ga-rollout": True})
+    def test_bad_staff_disabled(self):
+        self.add_user_permission(self.staff_user, RELOCATION_ADMIN_PERMISSION)
+        self.login_as(user=self.staff_user, staff=False)
+
+        # Ensures we don't reveal existence info to improperly authenticated users.
+        does_not_exist_uuid = uuid4().hex
+        self.get_error_response(str(does_not_exist_uuid), "somedir", "file.json", status_code=403)
+
+    def test_bad_has_superuser_but_no_relocation_admin_permission(self):
+        self.login_as(user=self.superuser, superuser=True)
+
+        # Ensures we don't reveal existence info to improperly authenticated users.
+        does_not_exist_uuid = uuid4().hex
+        response = self.get_error_response(
+            str(does_not_exist_uuid), "somedir", "file.json", status_code=403
+        )
+
+        assert response.data.get("detail") == ERR_NEED_RELOCATION_ADMIN
+
+    @override_options({"staff.ga-rollout": True})
+    def test_bad_has_staff_but_no_relocation_admin_permission(self):
+        self.login_as(user=self.staff_user, staff=True)
+
+        # Ensures we don't reveal existence info to improperly authenticated users.
+        does_not_exist_uuid = uuid4().hex
+        response = self.get_error_response(
+            str(does_not_exist_uuid), "somedir", "file.json", status_code=403
+        )
+
+        assert response.data.get("detail") == ERR_NEED_RELOCATION_ADMIN
+
+    @override_options({"staff.ga-rollout": True})
+    def test_bad_relocation_not_found(self):
+        self.add_user_permission(self.staff_user, RELOCATION_ADMIN_PERMISSION)
+        self.login_as(user=self.staff_user, staff=True)
+        does_not_exist_uuid = uuid4().hex
+        self.get_error_response(str(does_not_exist_uuid), "somedir", "file.json", status_code=404)
+
+    @override_options({"staff.ga-rollout": True})
+    def test_bad_file_not_found(self):
+        self.add_user_permission(self.staff_user, RELOCATION_ADMIN_PERMISSION)
+        self.login_as(user=self.staff_user, staff=True)
+        self.get_error_response(
+            str(self.relocation.uuid), "nonexistent", "file.json", status_code=404
+        )
