Created ad token v4 that uses Base64URLEncoding (#5)

* Created ad token v4 that uses Base64URLEncoding
* Changed ADVERTISING_TOKEN_V4 enum to 128 for readability in Base64
* removed use of datetime.utcnow/datetime.utcfromtimestamp as it's problematic when not run in utc timzone see https://blog.ganssle.io/articles/2019/11/utcnow.html
* further refactoring to make it more similar to other language sdk's class structure
* added cross platform unit tests and further refactoring to make it more similar to other language sdk's class structure
* created uid2_client/identity_scope.py  uid2_client/identity_type.py

Author: sunnywu

tests/test_encryption.py

@@ -1,4 +1,5 @@
 import datetime as dt
+from datetime import timezone
 import unittest
 
 from uid2_client.keys import *
@@ -22,19 +23,19 @@ def test_key_is_active(self):
 
     def test_empty_keys_collection(self):
         keys = EncryptionKeysCollection([])
-        self.assertFalse(keys.valid(dt.datetime.utcnow()))
+        self.assertFalse(keys.valid(dt.datetime.now(tz=timezone.utc)))
         self.assertEqual(len(keys), 0)
         self.assertEqual(len(keys.key_ids()), 0)
         self.assertEqual(len(keys.values()), 0)
         self.assertNotIn(123, keys)
         self.assertIsNone(keys.get(123))
         with self.assertRaises(KeyError):
             keys[123]
-        self.assertIsNone(keys.get_active_site_key(22, dt.datetime.utcnow()))
+        self.assertIsNone(keys.get_active_site_key(22, dt.datetime.now(tz=timezone.utc)))
 
 
     def test_multiple_keys_in_collection(self):
-        now = dt.datetime.utcnow()
+        now = dt.datetime.now(tz=timezone.utc)
         keys = EncryptionKeysCollection([
             EncryptionKey(123, 201, now - dt.timedelta(days=5), now - dt.timedelta(days=4), now - dt.timedelta(days=3), b'123456'),
             EncryptionKey(124, 202, now - dt.timedelta(days=1), now, now + dt.timedelta(days=1), b'234567')])
@@ -59,7 +60,7 @@ def test_multiple_keys_in_collection(self):
 
 
     def test_site_keys(self):
-        now = dt.datetime.utcnow()
+        now = dt.datetime.now(tz=timezone.utc)
         keys = EncryptionKeysCollection([
             EncryptionKey(122, 200, now, now - dt.timedelta(days=1), now + dt.timedelta(days=3), b'000000'),
             EncryptionKey(123, 201, now, now - dt.timedelta(days=5), now + dt.timedelta(days=3), b'111111'),
@@ -89,7 +90,7 @@ def test_site_keys(self):
 
 
     def test_non_site_keys(self):
-        now = dt.datetime.utcnow()
+        now = dt.datetime.now(tz=timezone.utc)
         keys = EncryptionKeysCollection([
             EncryptionKey(122, -1, now, now - dt.timedelta(days=9), now + dt.timedelta(days=3), b'000000'),
             EncryptionKey(123, -1, now, now - dt.timedelta(days=5), now + dt.timedelta(days=3), b'111111'),
@@ -103,7 +104,7 @@ def test_non_site_keys(self):
 
 
     def test_all_keys_in_collection_expired(self):
-        now = dt.datetime.utcnow()
+        now = dt.datetime.now(tz=timezone.utc)
         keys = EncryptionKeysCollection([
             EncryptionKey(123, 201, now, now - dt.timedelta(days=5), now - dt.timedelta(days=3), b'123456'),
             EncryptionKey(124, 202, now, now - dt.timedelta(days=1), now - dt.timedelta(days=1), b'234567')])

tests/test_keys.py

@@ -0,0 +1,105 @@
+import base64
+
+from datetime import timezone
+import os
+from uid2_client import encryption_block_size
+from uid2_client.advertising_token_version import AdvertisingTokenVersion
+from uid2_client.encryption import _encrypt_data_v1, _encrypt_gcm, _PayloadType
+from uid2_client.identity_scope import IdentityScope
+from uid2_client.identity_type import IdentityType
+from uid2_client.keys import *
+from uid2_client.uid2_base64_url_coder import Uid2Base64UrlCoder
+
+
+class Params:
+    def __init__(self, expiry=dt.datetime.now(tz=timezone.utc) + dt.timedelta(hours=1),
+                 identity_scope=IdentityScope.UID2.value, identity_type=IdentityType.Email.value):
+        self.identity_scope = identity_scope
+        self.identity_type = identity_type
+        self.token_expiry = expiry
+        if not isinstance(expiry, dt.datetime):
+            self.token_expiry = dt.datetime.now(tz=timezone.utc) + expiry
+
+
+def default_params():
+    return Params()
+
+
+class UID2TokenGenerator:
+    @staticmethod
+    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')
+        identity += int.to_bytes(int((dt.datetime.now(tz=timezone.utc) - dt.timedelta(hours=1)).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')
+        master_payload += _encrypt_data_v1(identity, key=site_key, iv=identity_iv)
+        master_iv = bytes([21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36])
+
+        token = int.to_bytes(version, 1, 'big')
+        token += _encrypt_data_v1(master_payload, key=master_key, iv=master_iv)
+
+        return base64.b64encode(token).decode('ascii')
+
+    @staticmethod
+    def generate_uid2_token_v3(id_str, master_key, site_id, site_key, params = default_params()):
+        return UID2TokenGenerator.generate_uid2_token_with_debug_info(id_str, master_key, site_id, site_key, params,
+                                                    AdvertisingTokenVersion.ADVERTISING_TOKEN_V3.value)
+
+    @staticmethod
+    def generate_uid2_token_v4(id_str, master_key, site_id, site_key, params = default_params()):
+        return UID2TokenGenerator.generate_uid2_token_with_debug_info(id_str, master_key, site_id, site_key, params,
+                                                    AdvertisingTokenVersion.ADVERTISING_TOKEN_V4.value)
+
+    @staticmethod
+    def generate_uid2_token_with_debug_info(id_str, master_key, site_id, site_key, params, version):
+        id = base64.b64decode(id_str)
+
+        site_payload = int.to_bytes(site_id, 4, 'big')
+        site_payload += int.to_bytes(0, 8, 'big')  # publisher id
+        site_payload += int.to_bytes(0, 4, 'big')  # client key id
+
+        site_payload += int.to_bytes(0, 4, 'big')  # privacy bits
+        site_payload += int.to_bytes(int((dt.datetime.now(tz=timezone.utc) - dt.timedelta(hours=1)).timestamp()) * 1000,
+                                     8, 'big')  # established
+        site_payload += int.to_bytes(int((dt.datetime.now(tz=timezone.utc) - dt.timedelta(hours=1)).timestamp()) * 1000,
+                                     8, 'big')  # refreshed
+        site_payload += id
+
+        expiry = params.token_expiry
+        master_payload = int.to_bytes(int(expiry.timestamp()) * 1000, 8, 'big')
+        master_payload += int.to_bytes(int((dt.datetime.now(tz=timezone.utc)).timestamp()) * 1000, 8, 'big')  # created
+
+        master_payload += int.to_bytes(0, 4, 'big')  # operator site id
+        master_payload += int.to_bytes(0, 1, 'big')  # operator type
+        master_payload += int.to_bytes(0, 4, 'big')  # operator version
+        master_payload += int.to_bytes(0, 4, 'big')  # operator key id
+
+        master_payload += int.to_bytes(site_key.key_id, 4, 'big')
+        master_payload += _encrypt_gcm(site_payload, None, site_key.secret)
+
+        token = int.to_bytes((params.identity_scope << 4 | params.identity_type << 2), 1, 'big')
+        token += int.to_bytes(version, 1, 'big')
+        token += int.to_bytes(master_key.key_id, 4, 'big')
+        token += _encrypt_gcm(master_payload, None, master_key.secret)
+
+        if version == AdvertisingTokenVersion.ADVERTISING_TOKEN_V4.value:
+            return Uid2Base64UrlCoder.encode(token)
+        else:
+            return base64.b64encode(token).decode('ascii')
+
+    @staticmethod
+    def encrypt_data_v2(data, key, site_id, now):
+        iv = os.urandom(encryption_block_size)
+        result = int.to_bytes(_PayloadType.ENCRYPTED_DATA.value, 1, 'big')
+        result += int.to_bytes(1, 1, 'big')  # version
+        result += int.to_bytes(int(now.timestamp() * 1000), 8, 'big')
+        result += int.to_bytes(site_id, 4, 'big')
+        result += _encrypt_data_v1(data, key, iv)
+        return base64.b64encode(result).decode('ascii')

tests/uid2_token_generator.py

@@ -13,10 +13,4 @@
 from .encryption import *
 from .keys import *
 
-from enum import Enum
 
-
-class IdentityScope(Enum):
-    """Enum for types of unified ID"""
-    UID2 = 0
-    EUID = 1

uid2_client/__init__.py

@@ -0,0 +1,9 @@
+from enum import Enum
+
+class AdvertisingTokenVersion(Enum):
+    # showing as "AHA..." in the Base64 Encoding (Base64 'H' is 000111 and 112 is 01110000)
+    ADVERTISING_TOKEN_V3 = 112
+    # showing as "AIA..." in the Base64URL Encoding ('H' is followed by 'I' hence
+    # this choice for the next token version) (Base64 'I' is 001000 and 128 is 10000000)
+    ADVERTISING_TOKEN_V4 = 128
+

uid2_client/advertising_token_version.py

@@ -6,6 +6,7 @@
 
 
 import datetime as dt
+from datetime import timezone
 import sys
 import threading
 
@@ -131,7 +132,7 @@ def _make_error_result(self, err):
 
 
     def _make_success_result(self, keys):
-        return EncryptionKeysAutoRefreshResult(keys, None, dt.datetime.utcnow())
+        return EncryptionKeysAutoRefreshResult(keys, None, dt.datetime.now(tz=timezone.utc))
 
 
     def __enter__(self):

uid2_client/auto_refresh.py

@@ -8,6 +8,7 @@
 
 import base64
 import datetime as dt
+from datetime import timezone
 import json
 import os
 import urllib.request as request
@@ -17,7 +18,7 @@
 
 
 def _make_dt(timestamp):
-    return dt.datetime.utcfromtimestamp(timestamp)
+    return dt.datetime.fromtimestamp(timestamp, tz=timezone.utc)
 
 
 class Uid2Client:
@@ -63,7 +64,7 @@ def refresh_keys(self):
         Returns:
             EncryptionKeysCollection containing the keys
         """
-        req, nonce = self._make_v2_request(dt.datetime.utcnow())
+        req, nonce = self._make_v2_request(dt.datetime.now(tz=timezone.utc))
         print(req)
         resp = self._post('/v2/key/latest', headers=self._auth_headers(), data=req)
         keys = [EncryptionKey(k['id'], k.get('site_id', -1), _make_dt(k['created']), _make_dt(k['activates']), _make_dt(k['expires']), base64.b64decode(k['secret']))

uid2_client/client.py

@@ -7,10 +7,13 @@
 
 import base64
 import datetime as dt
+from datetime import timezone
 import os
 from Crypto.Cipher import AES
 from enum import Enum
 
+from uid2_client.advertising_token_version import AdvertisingTokenVersion
+from uid2_client.uid2_base64_url_coder import Uid2Base64UrlCoder
 
 encryption_block_size = AES.block_size
 """int: block size for encryption routines
@@ -24,8 +27,9 @@ class _PayloadType(Enum):
     ENCRYPTED_DATA = 128
     ENCRYPTED_DATA_V3 = 96
 
+base64_url_special_chars = {"-", "_"}
 
-def decrypt_token(token, keys, now=dt.datetime.utcnow()):
+def decrypt_token(token, keys, now=dt.datetime.now(tz=timezone.utc)):
     """Decrypt advertising token to extract UID2 details.
 
     Args:
@@ -53,12 +57,18 @@ def _decrypt_token(token, keys, now):
     if not keys.valid(now):
         raise EncryptionError('no keys available or all keys have expired; refresh the latest keys from UID2 service')
 
-    token_bytes = base64.b64decode(token)
+    header_str = token[0:4]
+    index = next((i for i, ch in enumerate(header_str) if ch in base64_url_special_chars), None)
+    is_base64_url_encoding = (index is not None)
+    token_bytes = Uid2Base64UrlCoder.decode(header_str) if is_base64_url_encoding else base64.b64decode(header_str)
 
     if token_bytes[0] == 2:
-        return _decrypt_token_v2(token_bytes, keys, now)
-    elif token_bytes[1] == 112:
-        return _decrypt_token_v3(token_bytes, keys, now)
+        return _decrypt_token_v2(base64.b64decode(token), keys, now)
+    elif token_bytes[1] == AdvertisingTokenVersion.ADVERTISING_TOKEN_V3.value:
+        return _decrypt_token_v3(base64.b64decode(token), keys, now)
+    elif token_bytes[1] == AdvertisingTokenVersion.ADVERTISING_TOKEN_V4.value:
+        # same as V3 but use Base64URL encoding
+        return _decrypt_token_v3(Uid2Base64UrlCoder.decode(token), keys, now)
     else:
         raise EncryptionError('token version not supported')
 
@@ -73,7 +83,7 @@ def _decrypt_token_v2(token_bytes, keys, now):
     master_payload = _decrypt(token_bytes[21:], master_iv, master_key)
 
     expires_ms = int.from_bytes(master_payload[:8], 'big')
-    expires = dt.datetime.utcfromtimestamp(expires_ms / 1000.0)
+    expires = dt.datetime.fromtimestamp(expires_ms / 1000.0, tz=timezone.utc)
     if expires < now:
         raise EncryptionError("token expired")
 
@@ -92,7 +102,7 @@ def _decrypt_token_v2(token_bytes, keys, now):
 
     idx = 8 + id_len + 4
     established_ms = int.from_bytes(identity[idx:idx+8], 'big')
-    established = dt.datetime.utcfromtimestamp(established_ms / 1000.0)
+    established = dt.datetime.fromtimestamp(established_ms / 1000.0, tz=timezone.utc)
 
     return DecryptedToken(id_str, established, site_id, site_key.site_id)
 
@@ -106,7 +116,7 @@ def _decrypt_token_v3(token_bytes, keys, now):
     master_payload = _decrypt_gcm(token_bytes[6:], master_key.secret)
 
     expires_ms = int.from_bytes(master_payload[:8], 'big')
-    expires = dt.datetime.utcfromtimestamp(expires_ms / 1000.0)
+    expires = dt.datetime.fromtimestamp(expires_ms / 1000.0, tz=timezone.utc)
     if expires < now:
         raise EncryptionError("token expired")
 
@@ -128,7 +138,7 @@ def _decrypt_token_v3(token_bytes, keys, now):
     # client key id 12:16
     # privacy bits 16:20
     established_ms = int.from_bytes(site_payload[20:28], 'big')
-    established = dt.datetime.utcfromtimestamp(established_ms / 1000.0)
+    established = dt.datetime.fromtimestamp(established_ms / 1000.0, tz=timezone.utc)
     # refreshed_ms 28:36
 
     id_bytes = site_payload[36:]
@@ -178,7 +188,7 @@ def encrypt_data(data, identity_scope, **kwargs):
     """
     now = kwargs.get("now")
     if now is None:
-        now = dt.datetime.utcnow()
+        now = dt.datetime.now(tz=timezone.utc)
     keys = kwargs.get("keys")
     key = kwargs.get("key")
     if keys is not None and key is not None:
@@ -268,7 +278,7 @@ def _decrypt_data_v2(encrypted_bytes, keys):
     iv = encrypted_bytes[18:34]
     data = _decrypt(encrypted_bytes[34:], iv, key)
     encrypted_ms = int.from_bytes(encrypted_bytes[2:10], 'big')
-    encrypted_at = dt.datetime.utcfromtimestamp(encrypted_ms / 1000.0)
+    encrypted_at = dt.datetime.fromtimestamp(encrypted_ms / 1000.0, tz=timezone.utc)
     return DecryptedData(data, encrypted_at)
 
 
@@ -284,7 +294,7 @@ def _decrypt_data_v3(encrypted_bytes, keys):
 
     payload = _decrypt_gcm(encrypted_bytes[6:], key.secret)
     encrypted_ms = int.from_bytes(payload[:8], 'big')
-    encrypted_at = dt.datetime.utcfromtimestamp(encrypted_ms / 1000.0)
+    encrypted_at = dt.datetime.fromtimestamp(encrypted_ms / 1000.0, tz=timezone.utc)
 
     # site id 8:12
 

uid2_client/encryption.py

@@ -0,0 +1,8 @@
+from enum import Enum
+
+
+class IdentityScope(Enum):
+    """Enum for types of unified ID"""
+    UID2 = 0
+    EUID = 1
+

uid2_client/identity_scope.py

@@ -0,0 +1,7 @@
+from enum import Enum
+
+class IdentityType(Enum):
+    """Enum for types of ID source"""
+    Email = 0
+    Phone = 1
+