feat(backup): Support decryption for sentry compare (#58796)

Encrypted files can now be compared as well.

Issue: getsentry/team-ospo#203

src/sentry/backup/helpers.py

@@ -148,7 +148,7 @@ def decrypt_data_encryption_key_local(
         format=serialization.PublicFormat.SubjectPublicKeyInfo,
     )
     if unwrapped.plain_public_key_pem != generated_public_key_pem:
-        raise ValueError(
+        raise DecryptionError(
             "The public key does not match that generated by the `decrypt_with` private key."
         )
 
@@ -180,7 +180,7 @@ class CryptoKeyVersion(NamedTuple):
     version: str
 
 
-class KMSDecryptionError(Exception):
+class DecryptionError(Exception):
     pass
 
 
@@ -192,7 +192,7 @@ def decrypt_data_encryption_key_using_gcp_kms(
     try:
         crypto_key_version = CryptoKeyVersion(**json_gcp_kms_config)
     except TypeError:
-        raise KMSDecryptionError(
+        raise DecryptionError(
             """Your supplied KMS configuration did not have the correct fields - please ensure that
             it is a single, top-level object with the fields `project_id`, `location`, `key_ring`,
             `key`, and `version`, with all values as strings."""
@@ -216,14 +216,14 @@ def decrypt_data_encryption_key_using_gcp_kms(
         }
     )
     if not decrypt_response.plaintext_crc32c == crc32c(decrypt_response.plaintext):
-        raise Exception("The response received from the server was corrupted in-transit.")
+        raise DecryptionError("The response received from the server was corrupted in-transit.")
 
     return decrypt_response.plaintext
 
 
 def decrypt_encrypted_tarball(
     tarball: BinaryIO, decryption_using_gcp_kms: bool, decrypt_with: BinaryIO
-) -> str:
+) -> bytes:
     """
     A tarball encrypted by a call to `_export` with `encrypt_with` set has some specific properties
     (filenames, etc). This method handles all of those, and decrypts using the provided private key
@@ -241,7 +241,7 @@ def decrypt_encrypted_tarball(
     )
     decrypted_dek = decryption_handler(unwrapped, decryption_info)
     decryptor = Fernet(decrypted_dek)
-    return decryptor.decrypt(unwrapped.encrypted_json_blob).decode("utf-8")
+    return decryptor.decrypt(unwrapped.encrypted_json_blob)
 
 
 def get_final_derivations_of(model: Type) -> set[Type]:

src/sentry/runner/commands/backup.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
 
-from typing import BinaryIO, Callable, Sequence, TextIO
+from io import BytesIO
+from typing import Callable, Sequence, TextIO
 
 import click
 
 from sentry.backup.comparators import get_default_comparators
 from sentry.backup.findings import Finding, FindingJSONEncoder
-from sentry.backup.helpers import ImportFlags
+from sentry.backup.helpers import DecryptionError, ImportFlags, Side, decrypt_encrypted_tarball
 from sentry.backup.validate import validate
 from sentry.runner.decorators import configuration
 from sentry.utils import json
@@ -88,9 +89,9 @@ def parse_filter_arg(filter_arg: str) -> set[str] | None:
     return filter_by
 
 
-def get_decryptor(
-    decrypt_with: BinaryIO | None, decrypt_with_gcp_kms: BinaryIO | None
-) -> BinaryIO | None:
+def get_decryptor_io_from_flags(
+    decrypt_with: BytesIO | None, decrypt_with_gcp_kms: BytesIO | None
+) -> BytesIO | None:
     if decrypt_with is not None and decrypt_with_gcp_kms is not None:
         raise click.UsageError(
             """`--decrypt-with` and `--decrypt-with-gcp-kms` are mutually exclusive options - you may use one or the other, but not both."""
@@ -128,23 +129,72 @@ def write_findings(
     required=False,
     help=FINDINGS_FILE_HELP,
 )
+@click.option(
+    "--decrypt-left-with",
+    type=click.File("rb"),
+    help=DECRYPT_WITH_HELP,
+)
+@click.option(
+    "--decrypt-left-with-gcp-kms",
+    type=click.File("rb"),
+    help=DECRYPT_WITH_GCP_KMS_HELP,
+)
+@click.option(
+    "--decrypt-right-with",
+    type=click.File("rb"),
+    help="Identical to `--decrypt-left-with`, but for the 2nd input argument.",
+)
+@click.option(
+    "--decrypt-right-with-gcp-kms",
+    type=click.File("rb"),
+    help="Identical to `--decrypt-left-with-gcp-kms`, but for the 2nd input argument.",
+)
 @configuration
-def compare(left, right, findings_file):
+def compare(
+    left,
+    right,
+    decrypt_left_with,
+    decrypt_left_with_gcp_kms,
+    decrypt_right_with,
+    decrypt_right_with_gcp_kms,
+    findings_file,
+):
     """
-    Compare two exports generated by the `export` command for equality, modulo certain necessary changed like `date_updated` timestamps, unique tokens, and the like.
+    Compare two exports generated by the `export` command for equality, modulo certain necessary
+    expected differences like `date_updated` timestamps, unique tokens, and the like.
     """
 
-    with left:
+    # Helper function that loads data from one of the two sides, decrypting it if necessary along
+    # the way.
+    def load_data(
+        side: Side, src: BytesIO, decrypt_with: BytesIO, decrypt_with_gcp_kms: BytesIO
+    ) -> json.JSONData:
+        decrypt_io = get_decryptor_io_from_flags(decrypt_with, decrypt_with_gcp_kms)
+
+        # Decrypt the tarball, if the user has indicated that this is one by using either of the
+        # `--decrypt...` flags.
+        if decrypt_io is not None:
+            try:
+                input = BytesIO(
+                    decrypt_encrypted_tarball(src, decrypt_with_gcp_kms is not None, decrypt_io)
+                )
+            except DecryptionError as e:
+                click.echo(f"Invalid {side.name} side tarball: {str(e)}", err=True)
+        else:
+            input = src
+
+        # Now read the input string into memory as JSONData.
         try:
-            left_data = json.load(left)
+            data = json.load(input)
         except json.JSONDecodeError:
-            click.echo("Invalid left JSON", err=True)
+            click.echo(f"Invalid {side.name} JSON", err=True)
+
+        return data
 
+    with left:
+        left_data = load_data(Side.left, left, decrypt_left_with, decrypt_left_with_gcp_kms)
     with right:
-        try:
-            right_data = json.load(right)
-        except json.JSONDecodeError:
-            click.echo("Invalid right JSON", err=True)
+        right_data = load_data(Side.right, right, decrypt_right_with, decrypt_right_with_gcp_kms)
 
     res = validate(left_data, right_data, get_default_comparators())
     if res:
@@ -214,7 +264,7 @@ def import_users(
     try:
         import_in_user_scope(
             src,
-            decrypt_with=get_decryptor(decrypt_with, decrypt_with_gcp_kms),
+            decrypt_with=get_decryptor_io_from_flags(decrypt_with, decrypt_with_gcp_kms),
             flags=ImportFlags(
                 merge_users=merge_users,
                 decrypt_using_gcp_kms=decrypt_with_gcp_kms is not None,
@@ -284,7 +334,7 @@ def import_organizations(
     try:
         import_in_organization_scope(
             src,
-            decrypt_with=get_decryptor(decrypt_with, decrypt_with_gcp_kms),
+            decrypt_with=get_decryptor_io_from_flags(decrypt_with, decrypt_with_gcp_kms),
             flags=ImportFlags(
                 merge_users=merge_users,
                 decrypt_using_gcp_kms=decrypt_with_gcp_kms is not None,
@@ -352,7 +402,7 @@ def import_config(
     try:
         import_in_config_scope(
             src,
-            decrypt_with=get_decryptor(decrypt_with, decrypt_with_gcp_kms),
+            decrypt_with=get_decryptor_io_from_flags(decrypt_with, decrypt_with_gcp_kms),
             flags=ImportFlags(
                 merge_users=merge_users,
                 overwrite_configs=overwrite_configs,
@@ -412,7 +462,7 @@ def import_global(
     try:
         import_in_global_scope(
             src,
-            decrypt_with=get_decryptor(decrypt_with, decrypt_with_gcp_kms),
+            decrypt_with=get_decryptor_io_from_flags(decrypt_with, decrypt_with_gcp_kms),
             flags=ImportFlags(
                 overwrite_configs=overwrite_configs,
                 decrypt_using_gcp_kms=decrypt_with_gcp_kms is not None,