Skip to content

Commit d019440

Browse files
asloobqasloobq
authored
Merge pull request #34 from IABTechLab/aaq-UID2-1649-token-generator-refactor
[UID2-1649] Move UID2 Token generator into source code
2 parents f1329a4 + cec2ca2 commit d019440

File tree

5 files changed

+81
-123
lines changed

5 files changed

+81
-123
lines changed

‎tests/test_encryption.py‎

Lines changed: 5 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,7 @@
11
import unittest
22

33
from test_utils import *
4-
from tests.uid2_token_generator import Params
54
from uid2_client import *
6-
from uid2_client.encryption import _encrypt_token
75

86
_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])
97
_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):
417415

418416
keys = EncryptionKeysCollection([_master_key, _site_key, _keyset_key], default_keyset_id=20,
419417
master_keyset_id=9999, caller_site_id=20)
420-
421-
result = _encrypt_token(uid2, identity_scope, _master_key, _site_key, _site_id, now=now,
422-
token_expiry=now + dt.timedelta(days=30) if keys.get_token_expiry_seconds() is None \
423-
else now + dt.timedelta(seconds=int(keys.get_token_expiry_seconds())),
424-
ad_token_version=AdvertisingTokenVersion.ADVERTISING_TOKEN_V3)
425-
final = decrypt(result.encrypted_data, keys, now=now)
418+
token_expiry = now + dt.timedelta(days=30) if keys.get_token_expiry_seconds() is None \
419+
else now + dt.timedelta(seconds=int(keys.get_token_expiry_seconds()))
420+
result = UID2TokenGenerator.generate_uid2_token_v3(uid2, _master_key, _site_id, _site_key,
421+
Params(expiry=token_expiry, token_generated_at=now))
422+
final = decrypt(result, keys, now=now)
426423

427424
self.assertEqual(uid2, final.uid)
428425

‎tests/test_utils.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,8 @@
22
import datetime as dt
33
from datetime import timezone
44

5-
from tests.uid2_token_generator import UID2TokenGenerator
65
from uid2_client import EncryptionKey, encrypt, IdentityScope, IdentityType, EncryptionKeysCollection, \
7-
AdvertisingTokenVersion
6+
AdvertisingTokenVersion, UID2TokenGenerator
87

98
master_secret = bytes(
109
[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‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,5 +23,6 @@
2323
from .encryption_status import *
2424
from .encryption_data_response import *
2525
from .refresh_response import *
26+
from .uid2_token_generator import *
2627

2728

‎uid2_client/encryption.py‎

Lines changed: 3 additions & 81 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,8 @@
1010
from datetime import timezone
1111
import os
1212
from Crypto.Cipher import AES
13-
from enum import Enum
1413

14+
from uid2_client.uid2_token_generator import Params, UID2TokenGenerator, _encrypt_gcm, _PayloadType
1515
from uid2_client.advertising_token_version import AdvertisingTokenVersion
1616
from uid2_client.client_type import ClientType
1717
from uid2_client.decryption_status import DecryptionStatus
@@ -22,19 +22,6 @@
2222
from uid2_client.identity_scope import IdentityScope
2323

2424

25-
encryption_block_size = AES.block_size
26-
"""int: block size for encryption routines
27-
28-
This determines the size of initialization vectors (IV), required data padding, etc.
29-
"""
30-
31-
32-
class _PayloadType(Enum):
33-
"""Enum for types of payload that can be encoded in opaque strings"""
34-
ENCRYPTED_DATA = 128
35-
ENCRYPTED_DATA_V3 = 96
36-
37-
3825
base64_url_special_chars = {"-", "_"}
3926

4027

@@ -255,48 +242,6 @@ def _decrypt_token_v3(token_bytes, keys, domain_name, client_type, now, token_ve
255242
keys.get_identity_scope(), identity_type, token_version, is_client_side_generated, expires)
256243

257244

258-
def _encrypt_token(uid2, identity_scope, master_key, site_key, site_id, now, token_expiry, ad_token_version):
259-
site_payload = bytearray(128)
260-
# Publisher Data
261-
site_payload[0:4] = int.to_bytes(site_id, byteorder='big', length=4) # Site id
262-
site_payload[4:12] = int.to_bytes(0, byteorder='big', length=8) # Publisher ID
263-
site_payload[12:16] = int.to_bytes(0, byteorder='big', length=4) # Client Key ID
264-
# User Identity Data
265-
site_payload[16:20] = int.to_bytes(0, byteorder='big', length=4) # Privacy Bits
266-
site_payload[20:28] = int.to_bytes(int(now.timestamp()) * 1000, byteorder='big',
267-
length=8) # Established
268-
site_payload[28:36] = int.to_bytes(int(now.timestamp()) * 1000, byteorder='big', length=8) # last refresh
269-
site_payload[36:] = bytes(base64.b64decode(uid2))
270-
271-
id_payload = _encrypt_gcm(bytes(site_payload), None, site_key.secret)
272-
273-
# Operator Identity Data
274-
master_payload = bytearray(256)
275-
master_payload[:8] = int.to_bytes(int(token_expiry.timestamp()) * 1000, byteorder='big', length=8) # Expiry
276-
master_payload[8:16] = int.to_bytes(int(now.timestamp()), byteorder='big', length=8) # Token Created
277-
master_payload[16:20] = int.to_bytes(0, byteorder='big', length=4) # Site ID
278-
master_payload[20:21] = int.to_bytes(1, byteorder='big', length=1) # Operator Type
279-
master_payload[21:25] = int.to_bytes(0, byteorder='big', length=4) # Operator Version
280-
master_payload[25:29] = int.to_bytes(0, byteorder='big', length=4) # Operator Key ID
281-
master_payload[29:33] = int.to_bytes(site_key.key_id, byteorder='big', length=4) # Site Key ID
282-
master_payload[33:] = bytes(id_payload)
283-
284-
encrypted_master_payload = _encrypt_gcm(bytes(master_payload), None, master_key.secret)
285-
286-
root_writer = bytearray(len(encrypted_master_payload) + 6)
287-
first_char = uid2[0]
288-
identity_type = IdentityType.Phone if first_char == 'F' or first_char == 'B' else IdentityType.Email
289-
root_writer[0:1] = int.to_bytes((int(identity_scope) << 4 | int(identity_type) << 2) | 3, byteorder='big', length=1)
290-
root_writer[1:2] = int.to_bytes(ad_token_version.value, byteorder='big', length=1)
291-
root_writer[2:6] = int.to_bytes(master_key.key_id, byteorder='big', length=4)
292-
root_writer[6:] = bytes(encrypted_master_payload)
293-
294-
if ad_token_version == AdvertisingTokenVersion.ADVERTISING_TOKEN_V4:
295-
return EncryptionDataResponse.make_success(Uid2Base64UrlCoder.encode(root_writer))
296-
297-
return EncryptionDataResponse.make_success(base64.b64encode(root_writer))
298-
299-
300245
# DEPRECATED, DO NOT CALL DIRECTLY. PLEASE USE Uid2Client's client.encrypt()
301246
def encrypt(uid2, identity_scope, keys, keyset_id=None, **kwargs):
302247
""" Encrypt an UID2 into a sharing token
@@ -341,7 +286,8 @@ def encrypt(uid2, identity_scope, keys, keyset_id=None, **kwargs):
341286
if identity_scope is None:
342287
identity_scope = keys.get_identity_scope()
343288
try:
344-
return _encrypt_token(uid2, identity_scope, master_key, key, site_id, now, token_expiry, ad_token_version)
289+
params = Params(expiry=token_expiry, identity_scope=identity_scope, token_generated_at=now)
290+
return EncryptionDataResponse.make_success(UID2TokenGenerator.generate_uid2_token_v4(uid2, master_key, site_id, key, params))
345291
except Exception:
346292
return EncryptionDataResponse.make_error(EncryptionStatus.ENCRYPTION_FAILURE)
347293

@@ -431,10 +377,6 @@ def encrypt_data(data, identity_scope, **kwargs):
431377
return base64.b64encode(result).decode('ascii')
432378

433379

434-
def _encrypt_data_v1(data, key, iv):
435-
return int.to_bytes(key.key_id, 4, 'big') + iv + _encrypt(data, iv, key)
436-
437-
438380
# DEPRECATED, DO NOT CALL
439381
def decrypt_data(encrypted_data, keys):
440382
"""Decrypt data encrypted with encrypt_data().
@@ -503,16 +445,6 @@ def _decrypt_data_v3(encrypted_bytes, keys):
503445
return DecryptedData(payload[12:], encrypted_at)
504446

505447

506-
def _add_pkcs7_padding(data, block_size):
507-
pad_len = block_size - (len(data) % block_size)
508-
return data + bytes([pad_len]) * pad_len
509-
510-
511-
def _encrypt(data, iv, key):
512-
cipher = AES.new(key.secret, AES.MODE_CBC, IV=iv)
513-
return cipher.encrypt(_add_pkcs7_padding(data, AES.block_size))
514-
515-
516448
def _decrypt(encrypted, iv, key):
517449
cipher = AES.new(key.secret, AES.MODE_CBC, iv=iv)
518450
data = cipher.decrypt(encrypted)
@@ -521,16 +453,6 @@ def _decrypt(encrypted, iv, key):
521453
return data[:-pad_len]
522454

523455

524-
def _encrypt_gcm(data, iv, secret):
525-
if iv is None:
526-
iv = os.urandom(12)
527-
elif len(iv) != 12:
528-
raise ValueError("iv must be 12 bytes")
529-
cipher = AES.new(secret, AES.MODE_GCM, nonce=iv)
530-
ciphertext, tag = cipher.encrypt_and_digest(data)
531-
return cipher.nonce + ciphertext + tag
532-
533-
534456
def _decrypt_gcm(encrypted, secret):
535457
cipher = AES.new(secret, AES.MODE_GCM, nonce=encrypted[:12])
536458
return cipher.decrypt_and_verify(encrypted[12:-16], encrypted[-16:])

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

Lines changed: 71 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,60 @@
11
import base64
2+
import os
3+
from enum import Enum
24

5+
from Crypto.Cipher import AES
36
from datetime import timezone
4-
import os
5-
from uid2_client import encryption_block_size
6-
from uid2_client.advertising_token_version import AdvertisingTokenVersion
7-
from uid2_client.encryption import _encrypt_data_v1, _encrypt_gcm, _PayloadType
7+
import datetime as dt
88
from uid2_client.identity_scope import IdentityScope
9+
10+
11+
from uid2_client.advertising_token_version import AdvertisingTokenVersion
912
from uid2_client.identity_type import IdentityType
10-
from uid2_client.keys import *
1113
from uid2_client.uid2_base64_url_coder import Uid2Base64UrlCoder
1214

15+
encryption_block_size = AES.block_size
16+
"""int: block size for encryption routines
17+
18+
This determines the size of initialization vectors (IV), required data padding, etc.
19+
"""
20+
21+
22+
class _PayloadType(Enum):
23+
"""Enum for types of payload that can be encoded in opaque strings"""
24+
ENCRYPTED_DATA = 128
25+
ENCRYPTED_DATA_V3 = 96
26+
27+
28+
def _add_pkcs7_padding(data, block_size):
29+
pad_len = block_size - (len(data) % block_size)
30+
return data + bytes([pad_len]) * pad_len
31+
32+
33+
def _encrypt_gcm(data, iv, secret):
34+
if iv is None:
35+
iv = os.urandom(12)
36+
elif len(iv) != 12:
37+
raise ValueError("iv must be 12 bytes")
38+
cipher = AES.new(secret, AES.MODE_GCM, nonce=iv)
39+
ciphertext, tag = cipher.encrypt_and_digest(data)
40+
return cipher.nonce + ciphertext + tag
41+
42+
43+
def _encrypt(data, iv, key):
44+
cipher = AES.new(key.secret, AES.MODE_CBC, IV=iv)
45+
return cipher.encrypt(_add_pkcs7_padding(data, AES.block_size))
46+
47+
48+
def _encrypt_data_v1(data, key, iv):
49+
return int.to_bytes(key.key_id, 4, 'big') + iv + _encrypt(data, iv, key)
50+
1351

1452
class Params:
1553
def __init__(self, expiry=dt.datetime.now(tz=timezone.utc) + dt.timedelta(hours=1),
16-
identity_scope=IdentityScope.UID2.value, token_created_at=dt.datetime.now(tz=timezone.utc)):
54+
identity_scope=IdentityScope.UID2.value, token_generated_at=dt.datetime.now(tz=timezone.utc)):
1755
self.identity_scope = identity_scope
1856
self.token_expiry = expiry
19-
self.token_created_at = token_created_at
57+
self.token_generated_at = token_generated_at
2058
if not isinstance(expiry, dt.datetime):
2159
self.token_expiry = dt.datetime.now(tz=timezone.utc) + expiry
2260

@@ -26,6 +64,7 @@ def default_params():
2664

2765

2866
class UID2TokenGenerator:
67+
2968
@staticmethod
3069
def generate_uid2_token_v2(id_str, master_key, site_id, site_key, params = default_params(), version=2):
3170
id = bytes(id_str, 'utf-8')
@@ -34,7 +73,7 @@ def generate_uid2_token_v2(id_str, master_key, site_id, site_key, params = defau
3473
identity += id
3574
# old privacy_bits
3675
identity += int.to_bytes(0, 4, 'big')
37-
created = params.token_created_at
76+
created = params.token_generated_at
3877
identity += int.to_bytes(int(created.timestamp()) * 1000, 8, 'big')
3978
identity_iv = bytes([10, 11, 12, 13, 14, 15, 16, 1, 2, 3, 4, 5, 6, 7, 8, 9])
4079
expiry = params.token_expiry
@@ -48,12 +87,12 @@ def generate_uid2_token_v2(id_str, master_key, site_id, site_key, params = defau
4887
return base64.b64encode(token).decode('ascii')
4988

5089
@staticmethod
51-
def generate_uid2_token_v3(id_str, master_key, site_id, site_key, params = default_params()):
90+
def generate_uid2_token_v3(id_str, master_key, site_id, site_key, params=default_params()):
5291
return UID2TokenGenerator.generate_uid2_token_with_debug_info(id_str, master_key, site_id, site_key, params,
5392
AdvertisingTokenVersion.ADVERTISING_TOKEN_V3.value)
5493

5594
@staticmethod
56-
def generate_uid2_token_v4(id_str, master_key, site_id, site_key, params = default_params()):
95+
def generate_uid2_token_v4(id_str, master_key, site_id, site_key, params=default_params()):
5796
return UID2TokenGenerator.generate_uid2_token_with_debug_info(id_str, master_key, site_id, site_key, params,
5897
AdvertisingTokenVersion.ADVERTISING_TOKEN_V4.value)
5998

@@ -63,7 +102,7 @@ def generate_uid_token(id_str, master_key, site_id, site_key, identity_scope, to
63102
params = default_params()
64103
params.identity_scope = identity_scope
65104
if created_at is not None:
66-
params.token_created_at = created_at
105+
params.token_generated_at = created_at
67106
if expires_at is not None:
68107
params.token_expiry = expires_at
69108
if token_version == AdvertisingTokenVersion.ADVERTISING_TOKEN_V2:
@@ -77,29 +116,29 @@ def generate_uid_token(id_str, master_key, site_id, site_key, identity_scope, to
77116

78117
@staticmethod
79118
def generate_uid2_token_with_debug_info(id_str, master_key, site_id, site_key, params, version):
80-
id = base64.b64decode(id_str)
81119

82-
site_payload = int.to_bytes(site_id, 4, 'big')
83-
site_payload += int.to_bytes(0, 8, 'big') # publisher id
84-
site_payload += int.to_bytes(0, 4, 'big') # client key id
120+
# Publisher Data
121+
site_payload = int.to_bytes(site_id, length=4, byteorder='big')
122+
site_payload += int.to_bytes(0, length=8, byteorder='big') # publisher id
123+
site_payload += int.to_bytes(0, length=4, byteorder='big') # client key id
85124

86-
site_payload += int.to_bytes(0, 4, 'big') # privacy bits
87-
created = params.token_created_at
88-
site_payload += int.to_bytes(int(created.timestamp()) * 1000, 8, 'big') # established
89-
site_payload += int.to_bytes(int(created.timestamp()) * 1000, 8, 'big') # refreshed
90-
site_payload += id
125+
# User Identity Data
126+
site_payload += int.to_bytes(0, length=4, byteorder='big') # privacy bits
127+
generated_at_timestamp = int(params.token_generated_at.timestamp()) * 1000
128+
site_payload += int.to_bytes(generated_at_timestamp, length=8, byteorder='big') # established
129+
site_payload += int.to_bytes(generated_at_timestamp, length=8, byteorder='big') # last refreshed/generated
130+
site_payload += base64.b64decode(id_str)
91131

92-
expiry = params.token_expiry
93-
94-
master_payload = int.to_bytes(int(expiry.timestamp()) * 1000, 8, 'big')
95-
master_payload += int.to_bytes(int(created.timestamp()) * 1000, 8, 'big') # created
132+
master_payload = int.to_bytes(int(params.token_expiry.timestamp()) * 1000, length=8, byteorder='big') # expiry
133+
master_payload += int.to_bytes(generated_at_timestamp, length=8, byteorder='big') # created
96134

97-
master_payload += int.to_bytes(0, 4, 'big') # operator site id
98-
master_payload += int.to_bytes(0, 1, 'big') # operator type
99-
master_payload += int.to_bytes(0, 4, 'big') # operator version
100-
master_payload += int.to_bytes(0, 4, 'big') # operator key id
135+
# Operator Identity Data
136+
master_payload += int.to_bytes(0, length=4, byteorder='big') # site id
137+
master_payload += int.to_bytes(1, length=1, byteorder='big') # operator type
138+
master_payload += int.to_bytes(0, length=4, byteorder='big') # operator version
139+
master_payload += int.to_bytes(0, length=4, byteorder='big') # operator key id
101140

102-
master_payload += int.to_bytes(site_key.key_id, 4, 'big')
141+
master_payload += int.to_bytes(site_key.key_id, length=4, byteorder='big') # Site Key ID
103142
master_payload += _encrypt_gcm(site_payload, None, site_key.secret)
104143

105144
first_char = id_str[0]
@@ -108,9 +147,9 @@ def generate_uid2_token_with_debug_info(id_str, master_key, site_id, site_key, p
108147
if (first_char == 'F') or (first_char == 'B'):
109148
identity_type = IdentityType.Phone.value
110149

111-
token = int.to_bytes((params.identity_scope << 4 | identity_type << 2) | 3, 1, 'big')
112-
token += int.to_bytes(version, 1, 'big')
113-
token += int.to_bytes(master_key.key_id, 4, 'big')
150+
token = int.to_bytes((params.identity_scope << 4 | identity_type << 2) | 3, length=1, byteorder='big')
151+
token += int.to_bytes(version, length=1, byteorder='big')
152+
token += int.to_bytes(master_key.key_id, length=4, byteorder='big')
114153
token += _encrypt_gcm(master_payload, None, master_key.secret)
115154

116155
if version == AdvertisingTokenVersion.ADVERTISING_TOKEN_V4.value:

0 commit comments

Comments
 (0)