Move uid2_token_generator in source code

Author: asloobq

tests/test_encryption.py

@@ -1,9 +1,7 @@
 import unittest
 
 from test_utils import *
-from tests.uid2_token_generator import Params
 from uid2_client import *
-from uid2_client.encryption import _encrypt_token
 
 _master_secret = bytes([139, 37, 241, 173, 18, 92, 36, 232, 165, 168, 23, 18, 38, 195, 123, 92, 160, 136, 185, 40, 91, 173, 165, 221, 168, 16, 169, 164, 38, 139, 8, 155])
 _site_secret =   bytes([32, 251, 7, 194, 132, 154, 250, 86, 202, 116, 104, 29, 131, 192, 139, 215, 48, 164, 11, 65, 226, 110, 167, 14, 108, 51, 254, 125, 65, 24, 23, 133])
@@ -417,12 +415,11 @@ def test_smoke_token_v3(self):
 
         keys = EncryptionKeysCollection([_master_key, _site_key, _keyset_key], default_keyset_id=20,
                                         master_keyset_id=9999, caller_site_id=20)
-
-        result = _encrypt_token(uid2, identity_scope, _master_key, _site_key, _site_id, now=now,
-                                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())),
-                         ad_token_version=AdvertisingTokenVersion.ADVERTISING_TOKEN_V3)
-        final = decrypt(result.encrypted_data, keys, now=now)
+        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_created_at=now))
+        final = decrypt(result, keys, now=now)
 
         self.assertEqual(uid2, final.uid)
 

tests/test_utils.py

@@ -2,9 +2,8 @@
 import datetime as dt
 from datetime import timezone
 
-from tests.uid2_token_generator import UID2TokenGenerator
 from uid2_client import EncryptionKey, encrypt, IdentityScope, IdentityType, EncryptionKeysCollection, \
-    AdvertisingTokenVersion
+    AdvertisingTokenVersion, UID2TokenGenerator
 
 master_secret = bytes(
     [139, 37, 241, 173, 18, 92, 36, 232, 165, 168, 23, 18, 38, 195, 123, 92, 160, 136, 185, 40, 91, 173, 165, 221, 168,

uid2_client/__init__.py

@@ -23,5 +23,6 @@
 from .encryption_status import *
 from .encryption_data_response import *
 from .refresh_response import *
+from .uid2_token_generator import *
 
 

uid2_client/encryption.py

@@ -10,8 +10,8 @@
 from datetime import timezone
 import os
 from Crypto.Cipher import AES
-from enum import Enum
 
+from uid2_client.uid2_token_generator import Params, UID2TokenGenerator, _encrypt_gcm, _PayloadType
 from uid2_client.advertising_token_version import AdvertisingTokenVersion
 from uid2_client.client_type import ClientType
 from uid2_client.decryption_status import DecryptionStatus
@@ -22,19 +22,6 @@
 from uid2_client.identity_scope import IdentityScope
 
 
-encryption_block_size = AES.block_size
-"""int: block size for encryption routines
-
-This determines the size of initialization vectors (IV), required data padding, etc.
-"""
-
-
-class _PayloadType(Enum):
-    """Enum for types of payload that can be encoded in opaque strings"""
-    ENCRYPTED_DATA = 128
-    ENCRYPTED_DATA_V3 = 96
-
-
 base64_url_special_chars = {"-", "_"}
 
 
@@ -255,48 +242,6 @@ def _decrypt_token_v3(token_bytes, keys, domain_name, client_type, now, token_ve
                           keys.get_identity_scope(), identity_type, token_version, is_client_side_generated, expires)
 
 
-def _encrypt_token(uid2, identity_scope, master_key, site_key, site_id, now, token_expiry, ad_token_version):
-    site_payload = bytearray(128)
-    # Publisher Data
-    site_payload[0:4] = int.to_bytes(site_id, byteorder='big', length=4)  # Site id
-    site_payload[4:12] = int.to_bytes(0, byteorder='big', length=8)  # Publisher ID
-    site_payload[12:16] = int.to_bytes(0, byteorder='big', length=4)  # Client Key ID
-    # User Identity Data
-    site_payload[16:20] = int.to_bytes(0, byteorder='big', length=4)  # Privacy Bits
-    site_payload[20:28] = int.to_bytes(int(now.timestamp()) * 1000, byteorder='big',
-                                       length=8)  # Established
-    site_payload[28:36] = int.to_bytes(int(now.timestamp()) * 1000, byteorder='big', length=8)  # last refresh
-    site_payload[36:] = bytes(base64.b64decode(uid2))
-
-    id_payload = _encrypt_gcm(bytes(site_payload), None, site_key.secret)
-
-    # Operator Identity Data
-    master_payload = bytearray(256)
-    master_payload[:8] = int.to_bytes(int(token_expiry.timestamp()) * 1000, byteorder='big', length=8)  # Expiry
-    master_payload[8:16] = int.to_bytes(int(now.timestamp()), byteorder='big', length=8)  # Token Created
-    master_payload[16:20] = int.to_bytes(0, byteorder='big', length=4)  # Site ID
-    master_payload[20:21] = int.to_bytes(1, byteorder='big', length=1)  # Operator Type
-    master_payload[21:25] = int.to_bytes(0, byteorder='big', length=4)  # Operator Version
-    master_payload[25:29] = int.to_bytes(0, byteorder='big', length=4)  # Operator Key ID
-    master_payload[29:33] = int.to_bytes(site_key.key_id, byteorder='big', length=4)  # Site Key ID
-    master_payload[33:] = bytes(id_payload)
-
-    encrypted_master_payload = _encrypt_gcm(bytes(master_payload), None, master_key.secret)
-
-    root_writer = bytearray(len(encrypted_master_payload) + 6)
-    first_char = uid2[0]
-    identity_type = IdentityType.Phone if first_char == 'F' or first_char == 'B' else IdentityType.Email
-    root_writer[0:1] = int.to_bytes((int(identity_scope) << 4 | int(identity_type) << 2) | 3, byteorder='big', length=1)
-    root_writer[1:2] = int.to_bytes(ad_token_version.value, byteorder='big', length=1)
-    root_writer[2:6] = int.to_bytes(master_key.key_id, byteorder='big', length=4)
-    root_writer[6:] = bytes(encrypted_master_payload)
-
-    if ad_token_version == AdvertisingTokenVersion.ADVERTISING_TOKEN_V4:
-        return EncryptionDataResponse.make_success(Uid2Base64UrlCoder.encode(root_writer))
-
-    return EncryptionDataResponse.make_success(base64.b64encode(root_writer))
-
-
 # DEPRECATED, DO NOT CALL DIRECTLY. PLEASE USE Uid2Client's client.encrypt()
 def encrypt(uid2, identity_scope, keys, keyset_id=None, **kwargs):
     """ Encrypt an UID2 into a sharing token
@@ -341,7 +286,8 @@ def encrypt(uid2, identity_scope, keys, keyset_id=None, **kwargs):
     if identity_scope is None:
         identity_scope = keys.get_identity_scope()
     try:
-        return _encrypt_token(uid2, identity_scope, master_key, key, site_id, now, token_expiry, ad_token_version)
+        params = Params(expiry=token_expiry, identity_scope=identity_scope, token_created_at=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)
 
@@ -431,10 +377,6 @@ def encrypt_data(data, identity_scope, **kwargs):
     return base64.b64encode(result).decode('ascii')
 
 
-def _encrypt_data_v1(data, key, iv):
-    return int.to_bytes(key.key_id, 4, 'big') + iv + _encrypt(data, iv, key)
-
-
 # DEPRECATED, DO NOT CALL
 def decrypt_data(encrypted_data, keys):
     """Decrypt data encrypted with encrypt_data().
@@ -503,16 +445,6 @@ def _decrypt_data_v3(encrypted_bytes, keys):
     return DecryptedData(payload[12:], encrypted_at)
 
 
-def _add_pkcs7_padding(data, block_size):
-    pad_len = block_size - (len(data) % block_size)
-    return data + bytes([pad_len]) * pad_len
-
-
-def _encrypt(data, iv, key):
-    cipher = AES.new(key.secret, AES.MODE_CBC, IV=iv)
-    return cipher.encrypt(_add_pkcs7_padding(data, AES.block_size))
-
-
 def _decrypt(encrypted, iv, key):
     cipher = AES.new(key.secret, AES.MODE_CBC, iv=iv)
     data = cipher.decrypt(encrypted)
@@ -521,16 +453,6 @@ def _decrypt(encrypted, iv, key):
     return data[:-pad_len]
 
 
-def _encrypt_gcm(data, iv, secret):
-    if iv is None:
-        iv = os.urandom(12)
-    elif len(iv) != 12:
-        raise ValueError("iv must be 12 bytes")
-    cipher = AES.new(secret, AES.MODE_GCM, nonce=iv)
-    ciphertext, tag = cipher.encrypt_and_digest(data)
-    return cipher.nonce + ciphertext + tag
-
-
 def _decrypt_gcm(encrypted, secret):
     cipher = AES.new(secret, AES.MODE_GCM, nonce=encrypted[:12])
     return cipher.decrypt_and_verify(encrypted[12:-16], encrypted[-16:])

tests/uid2_token_generator.py uid2_client/uid2_token_generator.pytests/uid2_token_generator.py renamed to uid2_client/uid2_token_generator.py

@@ -1,15 +1,53 @@
 import base64
+import os
+from enum import Enum
 
+from Crypto.Cipher import AES
 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
+import datetime as dt
 from uid2_client.identity_scope import IdentityScope
+
+
+from uid2_client.advertising_token_version import AdvertisingTokenVersion
 from uid2_client.identity_type import IdentityType
-from uid2_client.keys import *
 from uid2_client.uid2_base64_url_coder import Uid2Base64UrlCoder
 
+encryption_block_size = AES.block_size
+"""int: block size for encryption routines
+
+This determines the size of initialization vectors (IV), required data padding, etc.
+"""
+
+
+class _PayloadType(Enum):
+    """Enum for types of payload that can be encoded in opaque strings"""
+    ENCRYPTED_DATA = 128
+    ENCRYPTED_DATA_V3 = 96
+
+
+def _add_pkcs7_padding(data, block_size):
+    pad_len = block_size - (len(data) % block_size)
+    return data + bytes([pad_len]) * pad_len
+
+
+def _encrypt_gcm(data, iv, secret):
+    if iv is None:
+        iv = os.urandom(12)
+    elif len(iv) != 12:
+        raise ValueError("iv must be 12 bytes")
+    cipher = AES.new(secret, AES.MODE_GCM, nonce=iv)
+    ciphertext, tag = cipher.encrypt_and_digest(data)
+    return cipher.nonce + ciphertext + tag
+
+
+def _encrypt(data, iv, key):
+    cipher = AES.new(key.secret, AES.MODE_CBC, IV=iv)
+    return cipher.encrypt(_add_pkcs7_padding(data, AES.block_size))
+
+
+def _encrypt_data_v1(data, key, iv):
+    return int.to_bytes(key.key_id, 4, 'big') + iv + _encrypt(data, iv, key)
+
 
 class Params:
     def __init__(self, expiry=dt.datetime.now(tz=timezone.utc) + dt.timedelta(hours=1),
@@ -26,6 +64,7 @@ def default_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')
@@ -48,12 +87,12 @@ def generate_uid2_token_v2(id_str, master_key, site_id, site_key, params = defau
         return base64.b64encode(token).decode('ascii')
 
     @staticmethod
-    def generate_uid2_token_v3(id_str, master_key, site_id, site_key, params = default_params()):
+    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()):
+    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)