From 2ea65f7f3745ef4aaae01545084ab2a202b638e6 Mon Sep 17 00:00:00 2001
From: mdtro <20070360+mdtro@users.noreply.github.com>
Date: Tue, 2 Apr 2024 23:11:14 -0500
Subject: [PATCH] custom exception when attempting to read more than once

---
 src/sentry/models/apitoken.py        | 21 ++++++++++++++++++++-
 tests/sentry/models/test_apitoken.py | 12 ++++++------
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/src/sentry/models/apitoken.py b/src/sentry/models/apitoken.py
index c2231084157452..2fe3cb83d8b806 100644
--- a/src/sentry/models/apitoken.py
+++ b/src/sentry/models/apitoken.py
@@ -33,6 +33,14 @@ def generate_token():
     return secrets.token_hex(nbytes=32)
 
 
+class PlaintextSecretAlreadyRead(Exception):
+    def __init__(
+        self,
+        message="the secret you are trying to read is read-once and cannot be accessed directly again",
+    ):
+        super().__init__(message)
+
+
 class ApiTokenManager(ControlOutboxProducingManager):
     def create(self, *args, **kwargs):
         token_type: AuthTokenType | None = kwargs.get("token_type", None)
@@ -128,6 +136,8 @@ def _plaintext_token(self):
 
         if plaintext_token is not None:
             setattr(self, f"_{manager_class_name}__plaintext_token", None)
+        else:
+            raise PlaintextSecretAlreadyRead()
 
         return plaintext_token
 
@@ -144,9 +154,18 @@ def _plaintext_refresh_token(self):
             self, f"_{manager_class_name}__plaintext_refresh_token", None
         )
 
-        if plaintext_refresh_token is not None:
+        if plaintext_refresh_token:
             setattr(self, f"_{manager_class_name}__plaintext_refresh_token", None)
 
+        # some token types do not have refresh tokens, so we check to see
+        # if there's a hash value that exists for the refresh token.
+        #
+        # if there is a hash value, then a refresh token is expected
+        # and if the plaintext_refresh_token is None, then it has already
+        # been read once so we should throw the exception
+        if not plaintext_refresh_token and self.refresh_token:
+            raise PlaintextSecretAlreadyRead()
+
         return plaintext_refresh_token
 
     def save(self, *args: Any, **kwargs: Any) -> None:
diff --git a/tests/sentry/models/test_apitoken.py b/tests/sentry/models/test_apitoken.py
index ddcbc5b8236160..a84adfbccd31d8 100644
--- a/tests/sentry/models/test_apitoken.py
+++ b/tests/sentry/models/test_apitoken.py
@@ -1,11 +1,12 @@
 import hashlib
 from datetime import timedelta
 
+import pytest
 from django.utils import timezone
 
 from sentry.conf.server import SENTRY_SCOPE_HIERARCHY_MAPPING, SENTRY_SCOPES
 from sentry.hybridcloud.models import ApiTokenReplica
-from sentry.models.apitoken import ApiToken
+from sentry.models.apitoken import ApiToken, PlaintextSecretAlreadyRead
 from sentry.models.integrations.sentry_app_installation import SentryAppInstallation
 from sentry.models.integrations.sentry_app_installation_token import SentryAppInstallationToken
 from sentry.silo import SiloMode
@@ -98,11 +99,10 @@ def test_plaintext_values_only_available_immediately_after_create(self):
         assert token._plaintext_token is not None
         assert token._plaintext_refresh_token is None  # user auth tokens don't have refresh tokens
 
-        _ = token._plaintext_token
-
-        # we read the value above so now it should
-        # now be None as it is a "read once" property
-        assert token._plaintext_token is None
+        # we accessed the plaintext token above when we asserted it was not None
+        # accessing it again should throw an exception
+        with pytest.raises(PlaintextSecretAlreadyRead):
+            _ = token._plaintext_token
 
     @override_options({"apitoken.save-hash-on-create": True})
     def test_user_auth_token_hash(self):