From 877298f8d7e38faad81933cf58a0b9d54f186535 Mon Sep 17 00:00:00 2001
From: Matthew <20070360+mdtro@users.noreply.github.com>
Date: Fri, 12 Jan 2024 17:02:20 -0600
Subject: [PATCH] feat: apitoken last characters option (#62972)

Take two of https://github.com/getsentry/sentry/pull/59455.

- Add an option to toggle on/off automatically populating the
`token_last_characters` for the `ApiToken` model.

This option will ensure tokens created from here on have the field
populated. It will also allow me to thoroughly test the backfill
migration needed for existing tokens that will be coming in a future PR
by disabling the option in tests, creating a bunch of API tokens, and
then running the backfill migration test.

Tracking Issue: https://github.com/getsentry/sentry/issues/58918
RFC: https://github.com/getsentry/rfcs/pull/32
---
 src/sentry/backup/comparators.py     |  3 ++-
 src/sentry/models/apitoken.py        | 16 +++++++---------
 src/sentry/options/defaults.py       |  7 +++++++
 tests/sentry/models/test_apitoken.py | 13 +++++++++++++
 4 files changed, 29 insertions(+), 10 deletions(-)

diff --git a/src/sentry/backup/comparators.py b/src/sentry/backup/comparators.py
index fb4a4e49bd5d93..250b8a5df5991e 100644
--- a/src/sentry/backup/comparators.py
+++ b/src/sentry/backup/comparators.py
@@ -764,7 +764,8 @@ def get_default_comparators():
         list,
         {
             "sentry.apitoken": [
-                HashObfuscatingComparator("refresh_token", "token", "token_last_characters"),
+                HashObfuscatingComparator("refresh_token", "token"),
+                IgnoredComparator("token_last_characters"),
                 UnorderedListComparator("scope_list"),
             ],
             "sentry.apiapplication": [HashObfuscatingComparator("client_id", "client_secret")],
diff --git a/src/sentry/models/apitoken.py b/src/sentry/models/apitoken.py
index d2b50cbe5fc9a9..d1f6b3926f358e 100644
--- a/src/sentry/models/apitoken.py
+++ b/src/sentry/models/apitoken.py
@@ -2,12 +2,13 @@
 
 import secrets
 from datetime import timedelta
-from typing import ClassVar, Collection, Optional, Tuple
+from typing import Any, ClassVar, Collection, Optional, Tuple
 
 from django.db import models, router, transaction
 from django.utils import timezone
 from django.utils.encoding import force_str
 
+from sentry import options
 from sentry.backup.dependencies import ImportKind
 from sentry.backup.helpers import ImportFlags
 from sentry.backup.scopes import ImportScope, RelocationScope
@@ -57,15 +58,12 @@ class Meta:
     def __str__(self):
         return force_str(self.token)
 
-    # TODO(mdtro): uncomment this function after 0583_apitoken_add_name_and_last_chars migration has been applied
-    # def save(self, *args: Any, **kwargs: Any) -> None:
-    #     # when a new ApiToken is created we take the last four characters of the token
-    #     # and save them in the `token_last_characters` field so users can identify
-    #     # tokens in the UI where they're mostly obfuscated
-    #     token_last_characters = self.token[-4:]
-    #     self.token_last_characters = token_last_characters
+    def save(self, *args: Any, **kwargs: Any) -> None:
+        if options.get("apitoken.auto-add-last-chars"):
+            token_last_characters = self.token[-4:]
+            self.token_last_characters = token_last_characters
 
-    #     return super().save(**kwargs)
+        return super().save(**kwargs)
 
     def outbox_region_names(self) -> Collection[str]:
         return list(find_all_region_names())
diff --git a/src/sentry/options/defaults.py b/src/sentry/options/defaults.py
index 704172e9be1a92..b6057d036d456e 100644
--- a/src/sentry/options/defaults.py
+++ b/src/sentry/options/defaults.py
@@ -265,6 +265,13 @@
     flags=FLAG_ALLOW_EMPTY | FLAG_PRIORITIZE_DISK | FLAG_REQUIRED,
 )
 
+# API Tokens
+register(
+    "apitoken.auto-add-last-chars",
+    default=True,
+    type=Bool,
+    flags=FLAG_ALLOW_EMPTY | FLAG_PRIORITIZE_DISK | FLAG_AUTOMATOR_MODIFIABLE,
+)
 
 register(
     "api.rate-limit.org-create",
diff --git a/tests/sentry/models/test_apitoken.py b/tests/sentry/models/test_apitoken.py
index 731d1f462a8de3..f7a9f777c8abe5 100644
--- a/tests/sentry/models/test_apitoken.py
+++ b/tests/sentry/models/test_apitoken.py
@@ -9,6 +9,7 @@
 from sentry.models.integrations.sentry_app_installation_token import SentryAppInstallationToken
 from sentry.silo import SiloMode
 from sentry.testutils.cases import TestCase
+from sentry.testutils.helpers import override_options
 from sentry.testutils.outbox import outbox_runner
 from sentry.testutils.silo import assume_test_silo_mode, control_silo_test
 
@@ -61,6 +62,18 @@ def test_organization_id_for_non_internal(self):
             assert ApiTokenReplica.objects.get(apitoken_id=token.id).organization_id is None
         assert token.organization_id is None
 
+    @override_options({"apitoken.auto-add-last-chars": True})
+    def test_last_chars_are_set(self):
+        user = self.create_user()
+        token = ApiToken.objects.create(user_id=user.id)
+        assert token.token_last_characters == token.token[-4:]
+
+    @override_options({"apitoken.auto-add-last-chars": False})
+    def test_last_chars_are_not_set(self):
+        user = self.create_user()
+        token = ApiToken.objects.create(user_id=user.id)
+        assert token.token_last_characters is None
+
 
 class ApiTokenInternalIntegrationTest(TestCase):
     def setUp(self):