Handle lifetime check for v2 tokens differently

asloobq · asloobq · commit 825dcea61a43 · 2024-04-25T11:07:22.000-07:00
diff --git a/tests/test_bidstream_client.py b/tests/test_bidstream_client.py
@@ -124,10 +124,11 @@ def test_token_lifetime_too_long_for_bidstream(self):  # TokenLifetimeTooLongFor
                 self.assert_fails(result, expected_version, expected_scope)
 
     def test_token_generated_in_the_future_to_simulate_clock_skew(self):  # TokenGeneratedInTheFutureToSimulateClockSkew
+        # Note V2 does not have a "token generated" field, therefore v2 tokens can't have a future "token generated" date and are excluded from this test.
         created_at_future = dt.datetime.now(tz=timezone.utc) + dt.timedelta(minutes=31)  #max allowed clock skew is 30m
-        for expected_scope, expected_version in test_cases_all_scopes_all_versions:
+        for expected_scope, expected_version in test_cases_all_scopes_v3_v4_versions:
             with self.subTest(expected_scope=expected_scope, expected_version=expected_version):
-                token = generate_uid_token(expected_scope, expected_version, created_at=created_at_future)
+                token = generate_uid_token(expected_scope, expected_version, generated_at=created_at_future)
                 refresh_response = self._client._refresh_json(key_bidstream_response_json_default_keys(
                     expected_scope))
                 self.assertTrue(refresh_response.success)
@@ -164,7 +165,7 @@ def test_token_generated_in_the_future_legacy_client(self):  # TokenGeneratedInT
             with self.subTest(expected_scope=expected_scope, expected_version=expected_version):
                 legacy_client.refresh_json(key_bidstream_response_json_default_keys(
                     expected_scope))
-                token = generate_uid_token(expected_scope, expected_version, created_at=created_at_future)
+                token = generate_uid_token(expected_scope, expected_version, generated_at=created_at_future)
                 result = legacy_client.decrypt(token)
                 self.assert_success(result, expected_version, expected_scope)
 
@@ -246,7 +247,7 @@ def test_token_expiry_custom_decryption_time(self):  #TokenExpiryAndCustomNow
         expires_at = now - dt.timedelta(days=60)
         created_at = expires_at - dt.timedelta(minutes=1)
         token = generate_uid_token(IdentityScope.UID2, AdvertisingTokenVersion.ADVERTISING_TOKEN_V4,
-                                   created_at=created_at, expires_at=expires_at)
+                                   identity_established_at=created_at, expires_at=expires_at)
         result = self._client._decrypt_token_into_raw_uid(token, None, expires_at + dt.timedelta(seconds=1))
         self.assertFalse(result.success)
         self.assertEqual(result.status, DecryptionStatus.EXPIRED_TOKEN)
diff --git a/tests/test_encryption.py b/tests/test_encryption.py
@@ -418,7 +418,7 @@ def test_smoke_token_v3(self):
         token_expiry = now + dt.timedelta(days=30) if keys.get_token_expiry_seconds() is None \
             else now + dt.timedelta(seconds=int(keys.get_token_expiry_seconds()))
         result = UID2TokenGenerator.generate_uid2_token_v3(uid2, _master_key, _site_id, _site_key,
-                                                           Params(expiry=token_expiry, token_generated_at=now))
+                                                           Params(expiry=token_expiry, token_generated=now))
         final = decrypt(result, keys, now=now)
 
         self.assertEqual(uid2, final.uid)
diff --git a/tests/test_sharing_client.py b/tests/test_sharing_client.py
@@ -98,10 +98,11 @@ def test_token_lifetime_too_long_for_sharing(self):  # TokenLifetimeTooLongForSh
                 self._test_bidstream_client.assert_fails(result, expected_version, expected_scope)
 
     def test_token_generated_in_the_future_to_simulate_clock_skew(self):  # TokenGeneratedInTheFutureToSimulateClockSkew
+        # Note V2 does not have a "token generated" field, therefore v2 tokens can't have a future "token generated" date and are excluded from this test.
         created_at_future = dt.datetime.now(tz=timezone.utc) + dt.timedelta(minutes=31)  #max allowed clock skew is 30m
-        for expected_scope, expected_version in test_cases_all_scopes_all_versions:
+        for expected_scope, expected_version in test_cases_all_scopes_v3_v4_versions:
             with self.subTest(expected_scope=expected_scope, expected_version=expected_version):
-                token = generate_uid_token(expected_scope, expected_version, created_at=created_at_future)
+                token = generate_uid_token(expected_scope, expected_version, generated_at=created_at_future)
                 refresh_response = self._client._refresh_json(key_sharing_response_json_default_keys(
                     expected_scope))
                 self.assertTrue(refresh_response.success)
@@ -151,7 +152,7 @@ def test_token_generated_in_the_future_legacy_client(self):  # TokenGeneratedInT
             with self.subTest(expected_scope=expected_scope, expected_version=expected_version):
                 legacy_client.refresh_json(key_sharing_response_json_default_keys(
                     expected_scope))
-                token = generate_uid_token(expected_scope, expected_version, created_at=created_at_future)
+                token = generate_uid_token(expected_scope, expected_version, generated_at=created_at_future)
                 result = legacy_client.decrypt(token)
                 self._test_bidstream_client.assert_success(result, expected_version, expected_scope)
 
diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -114,6 +114,6 @@ def create_default_key_collection(key_set):
                                     99999, 2)
 
 
-def generate_uid_token(identity_scope, version, raw_uid=example_uid, created_at=None, expires_at=None):
-    return UID2TokenGenerator.generate_uid_token(raw_uid, master_key, site_id, site_key,
-                                                 identity_scope, version, created_at, expires_at)
+def generate_uid_token(identity_scope, version, raw_uid=example_uid, identity_established_at=None, generated_at=None, expires_at=None):
+    return UID2TokenGenerator.generate_uid_token(raw_uid, master_key, site_id, site_key, identity_scope, version,
+                                                 identity_established_at, generated_at, expires_at)
diff --git a/uid2_client/encryption.py b/uid2_client/encryption.py
@@ -111,18 +111,20 @@ def _decrypt_token(token, keys, domain_name, client_type, now):
         return DecryptedToken.make_error(DecryptionStatus.VERSION_NOT_SUPPORTED)
 
 
-def _token_has_valid_lifetime(keys, client_type, established, expires, now):
+def _token_has_valid_lifetime(keys, client_type, generated_or_now, expires, now):
+    #  generated_or_now allows "now" for token v2, since v2 does not contain a "token generated" field.
+    #  v2 therefore checks against remaining lifetime rather than total lifetime
     if client_type is ClientType.BIDSTREAM:
         max_life_time_seconds = keys.get_max_bidstream_lifetime_seconds()
     elif client_type is ClientType.SHARING:
         max_life_time_seconds = keys.get_max_sharing_lifetime_seconds()
     else:
         return True  # Skip check for legacy clients
 
-    if (expires - established).total_seconds() > max_life_time_seconds:
+    if (expires - generated_or_now).total_seconds() > max_life_time_seconds:
         return False
-    elif established > now:
-        return (established - now).total_seconds() <= keys.get_allow_clock_skew_seconds()
+    elif generated_or_now > now:
+        return (generated_or_now - now).total_seconds() <= keys.get_allow_clock_skew_seconds()
     else:
         return True
 
@@ -165,7 +167,7 @@ def _decrypt_token_v2(token_bytes, keys, domain_name, client_type, now):
     established_ms = int.from_bytes(identity[idx:idx + 8], 'big')
     established = dt.datetime.fromtimestamp(established_ms / 1000.0, tz=timezone.utc)
 
-    if not _token_has_valid_lifetime(keys, client_type, established, expires, now):
+    if not _token_has_valid_lifetime(keys, client_type, now, expires, now):
         return DecryptedToken(DecryptionStatus.INVALID_TOKEN_LIFETIME, id_str, established, site_id, site_key.site_id,
                           keys.get_identity_scope(), None, AdvertisingTokenVersion.ADVERTISING_TOKEN_V2, False, expires)
 
@@ -199,7 +201,7 @@ def _decrypt_token_v3(token_bytes, keys, domain_name, client_type, now, token_ve
         return DecryptedToken(DecryptionStatus.EXPIRED_TOKEN, None, None, None, None,
                               keys.get_identity_scope(), identity_type, token_version, None, expires)
 
-    # created 8:16
+    generated_ms = int.from_bytes(master_payload[8:16], 'big')  # Token Generated
     # operator site id 16:20
     # operator type 20
     # operator version 21:25
@@ -219,6 +221,10 @@ def _decrypt_token_v3(token_bytes, keys, domain_name, client_type, now, token_ve
     # privacy bits 16:20
     privacy_bits = bitarray()
     privacy_bits.frombytes(site_payload[16:20])
+    established_ms = int.from_bytes(site_payload[20:28], 'big')
+    id_bytes = site_payload[36:]
+    id_str = base64.b64encode(id_bytes).decode('ascii')
+
     is_client_side_generated = False
     if privacy_bits[1]:
         is_client_side_generated = True
@@ -227,17 +233,13 @@ def _decrypt_token_v3(token_bytes, keys, domain_name, client_type, now, token_ve
         return DecryptedToken(DecryptionStatus.DOMAIN_NAME_CHECK_FAILED, None, None, site_id, site_key.site_id,
                           keys.get_identity_scope(), identity_type, token_version, is_client_side_generated, expires)
 
-    established_ms = int.from_bytes(site_payload[20:28], 'big')
     established = dt.datetime.fromtimestamp(established_ms / 1000.0, tz=timezone.utc)
-    # refreshed_ms 28:36
+    generated = dt.datetime.fromtimestamp(generated_ms / 1000.0, tz=timezone.utc)
 
-    if not _token_has_valid_lifetime(keys, client_type, established, expires, now):
+    if not _token_has_valid_lifetime(keys, client_type, generated, expires, now):
         return DecryptedToken(DecryptionStatus.INVALID_TOKEN_LIFETIME, None, established, site_id, site_key.site_id,
                           keys.get_identity_scope(), identity_type, token_version, is_client_side_generated, expires)
 
-    id_bytes = site_payload[36:]
-    id_str = base64.b64encode(id_bytes).decode('ascii')
-
     return DecryptedToken(DecryptionStatus.SUCCESS, id_str, established, site_id, site_key.site_id,
                           keys.get_identity_scope(), identity_type, token_version, is_client_side_generated, expires)
 
@@ -286,7 +288,7 @@ def encrypt(uid2, identity_scope, keys, keyset_id=None, **kwargs):
     if identity_scope is None:
         identity_scope = keys.get_identity_scope()
     try:
-        params = Params(expiry=token_expiry, identity_scope=identity_scope, token_generated_at=now)
+        params = Params(expiry=token_expiry, identity_scope=identity_scope, token_generated=now)
         return EncryptionDataResponse.make_success(UID2TokenGenerator.generate_uid2_token_v4(uid2, master_key, site_id, key, params))
     except Exception:
         return EncryptionDataResponse.make_error(EncryptionStatus.ENCRYPTION_FAILURE)
diff --git a/uid2_client/uid2_token_generator.py b/uid2_client/uid2_token_generator.py
@@ -51,10 +51,12 @@ def _encrypt_data_v1(data, key, iv):
 
 class Params:
     def __init__(self, expiry=dt.datetime.now(tz=timezone.utc) + dt.timedelta(hours=1),
-                 identity_scope=IdentityScope.UID2.value, token_generated_at=dt.datetime.now(tz=timezone.utc)):
-        self.identity_scope = identity_scope
+                 identity_scope=IdentityScope.UID2.value, token_generated=dt.datetime.now(tz=timezone.utc),
+                 identity_established=dt.datetime.now(tz=timezone.utc)):
         self.token_expiry = expiry
-        self.token_generated_at = token_generated_at
+        self.identity_scope = identity_scope
+        self.token_generated = token_generated
+        self.identity_established = identity_established
         if not isinstance(expiry, dt.datetime):
             self.token_expiry = dt.datetime.now(tz=timezone.utc) + expiry
 
@@ -66,15 +68,14 @@ def default_params():
 class UID2TokenGenerator:
 
     @staticmethod
-    def generate_uid2_token_v2(id_str, master_key, site_id, site_key, params = default_params(), version=2):
+    def generate_uid2_token_v2(id_str, master_key, site_id, site_key, params=default_params(), version=2):
         id = bytes(id_str, 'utf-8')
         identity = int.to_bytes(site_id, 4, 'big')
         identity += int.to_bytes(len(id), 4, 'big')
         identity += id
         # old privacy_bits
         identity += int.to_bytes(0, 4, 'big')
-        created = params.token_generated_at
-        identity += int.to_bytes(int(created.timestamp()) * 1000, 8, 'big')
+        identity += int.to_bytes(int(params.identity_established.timestamp()) * 1000, 8, 'big')
         identity_iv = bytes([10, 11, 12, 13, 14, 15, 16, 1, 2, 3, 4, 5, 6, 7, 8, 9])
         expiry = params.token_expiry
         master_payload = int.to_bytes(int(expiry.timestamp()) * 1000, 8, 'big')
@@ -98,11 +99,13 @@ def generate_uid2_token_v4(id_str, master_key, site_id, site_key, params=default
 
     @staticmethod
     def generate_uid_token(id_str, master_key, site_id, site_key, identity_scope, token_version,
-                           created_at=None, expires_at=None):
+                           identity_established_at=None, generated_at=None, expires_at=None):
         params = default_params()
         params.identity_scope = identity_scope
-        if created_at is not None:
-            params.token_generated_at = created_at
+        if identity_established_at is not None:
+            params.identity_established = identity_established_at
+        if generated_at is not None:
+            params.token_generated = generated_at
         if expires_at is not None:
             params.token_expiry = expires_at
         if token_version == AdvertisingTokenVersion.ADVERTISING_TOKEN_V2:
@@ -124,13 +127,13 @@ def generate_uid2_token_with_debug_info(id_str, master_key, site_id, site_key, p
 
         # User Identity Data
         site_payload += int.to_bytes(0, length=4, byteorder='big')  # privacy bits
-        generated_at_timestamp = int(params.token_generated_at.timestamp()) * 1000
-        site_payload += int.to_bytes(generated_at_timestamp, length=8, byteorder='big')  # established
+        site_payload += int.to_bytes(int(params.identity_established.timestamp()) * 1000, length=8, byteorder='big')  # established
+        generated_at_timestamp = int(params.token_generated.timestamp()) * 1000
         site_payload += int.to_bytes(generated_at_timestamp, length=8, byteorder='big')  # last refreshed/generated
         site_payload += base64.b64decode(id_str)
 
         master_payload = int.to_bytes(int(params.token_expiry.timestamp()) * 1000, length=8, byteorder='big')  # expiry
-        master_payload += int.to_bytes(generated_at_timestamp, length=8, byteorder='big')  # created
+        master_payload += int.to_bytes(generated_at_timestamp, length=8, byteorder='big')  # generated
 
         # Operator Identity Data
         master_payload += int.to_bytes(0, length=4, byteorder='big')  # site id
