Finished token_v3 encryption

Author: cody-constine-ttd

.gitignore

@@ -1,3 +1,5 @@
 *.swp
+.idea/*
+*__pycache__*
 dist/*
 *egg-info/*

pyproject.toml

@@ -19,7 +19,7 @@ classifiers = [
     "Operating System :: OS Independent",
 ]
 dependencies = [
-    "pycrypto"
+    "pycryptodome"
 ]
 
 [project.urls]

tests/test_encryption.py

@@ -4,7 +4,7 @@
 import unittest
 
 from tests.uid2_token_generator import UID2TokenGenerator, Params
-from uid2_client import decrypt_token, encrypt_data, decrypt_data, encryption_block_size, EncryptionError, Uid2Base64UrlCoder, AdvertisingTokenVersion
+from uid2_client import *
 from uid2_client.identity_scope import IdentityScope
 from uid2_client.identity_type import IdentityType
 from uid2_client.keys import *
@@ -19,8 +19,9 @@
 
 _example_id = 'ywsvDNINiZOVSsfkHpLpSJzXzhr6Jx9Z/4Q0+lsEUvM='
 _now = dt.datetime.now(tz=timezone.utc)
-_master_key = EncryptionKey(_master_key_id, -1, _now - dt.timedelta(days=-1), _now, _now + dt.timedelta(days=1), _master_secret)
+_master_key = EncryptionKey(_master_key_id, -1, _now - dt.timedelta(days=-1), _now, _now + dt.timedelta(days=1), _master_secret, keyset_id=9999)
 _site_key = EncryptionKey(_site_key_id, _site_id, _now - dt.timedelta(days=-1), _now, _now + dt.timedelta(days=1), _site_secret)
+_keyset_key = EncryptionKey(_site_key_id, _site_id, _now - dt.timedelta(days=-1), _now, _now + dt.timedelta(days=1), _site_secret, keyset_id=20)
 _test_site_key = EncryptionKey(_test_site_key_id, _site_id, dt.datetime(2020, 1, 1, tzinfo=timezone.utc), dt.datetime(2020, 1, 1, tzinfo=timezone.utc), _now + dt.timedelta(days=1), encryption_block_size * b'9')
 
 class TestEncryptionFunctions(unittest.TestCase):
@@ -317,6 +318,20 @@ def test_decrypt_token_v2_custom_now(self):
         self.assertEqual(_example_id, result.uid2)
 
 
+    def test_encrypt_token_v3(self):
+        uid2 = "Y29keWlzZ3JlYXQ="
+        identity_scope = IdentityScope.UID2
+        now = dt.datetime.now(tz=timezone.utc)
+
+        keys = EncryptionKeysCollection([_master_key, _site_key, _keyset_key], default_keyset_id=20,
+                                        master_keyset_id=9999, caller_site_id=20)
+
+        result = encrypt_key(uid2, identity_scope, keys, now=now)
+        print(result)
+        final = decrypt_token(result, keys, now=now)
+
+        self.assertEqual(uid2, final.uid2)
+
     def test_decrypt_token_v2_invalid_payload(self):
         params = Params(dt.datetime.now(tz=timezone.utc) + dt.timedelta(seconds=-1))
         token = UID2TokenGenerator.generate_uid2_token_v2(_example_id, _master_key, _site_id, _site_key, params)

uid2_client/client.py

@@ -5,7 +5,6 @@
 >>> from uid2_client import Uid2Client
 """
 
-
 import base64
 import datetime as dt
 from datetime import timezone
@@ -53,7 +52,6 @@ def __init__(self, base_url, auth_key, secret_key):
         self._auth_key = auth_key
         self._secret_key = base64.b64decode(secret_key)
 
-
     def refresh_keys(self):
         """Get the latest encryption keys for advertising tokens.
 
@@ -66,20 +64,21 @@ def refresh_keys(self):
         """
         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']))
-            for k in json.loads(self._parse_v2_response(resp.read(), nonce)).get('body')]
-        return EncryptionKeysCollection(keys)
+        resp = self._post('/v2/key/sharing', headers=self._auth_headers(), data=req)
+        resp_body = json.loads(self._parse_v2_response(resp.read(), nonce)).get('body')
+        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']), k['keyset_id'])
+                for k in resp_body["keys"]]
 
+        return EncryptionKeysCollection(keys, resp_body["caller_site_id"], resp_body["master_keyset_id"],
+                                        resp_body["default_keyset_id"], resp_body["token_expiry_seconds"])
 
     def _make_url(self, path):
         return self._base_url + path
 
-
     def _auth_headers(self):
         return {'Authorization': 'Bearer ' + self._auth_key}
 
-
     def _make_v2_request(self, now):
         payload = int.to_bytes(int(now.timestamp() * 1000), 8, 'big')
         nonce = os.urandom(8)
@@ -90,14 +89,12 @@ def _make_v2_request(self, now):
 
         return base64.b64encode(envelope), nonce
 
-
     def _parse_v2_response(self, encrypted, nonce):
         payload = _decrypt_gcm(base64.b64decode(encrypted), self._secret_key)
         if nonce != payload[8:16]:
             raise ValueError("nonce mismatch")
         return payload[16:]
 
-
     def _post(self, path, headers, data):
         req = request.Request(self._make_url(path), headers=headers, method='POST', data=data)
         return request.urlopen(req)

uid2_client/encryption.py

@@ -146,6 +146,71 @@ def _decrypt_token_v3(token_bytes, keys, now):
 
     return DecryptedToken(id_str, established, site_id, site_key.site_id)
 
+def _encrypt_token_v3(uid2, identity_scope, master_key, site_key, site_id, now, token_expiry):
+    site_payload = bytearray(128)
+    #Site id
+    site_payload[0:4] = int.to_bytes(site_id, byteorder='big', length=4)
+    site_payload[4:12] = int.to_bytes(0, byteorder='big', length=8)
+    site_payload[12:16] = int.to_bytes(0, byteorder='big', length=4)
+    site_payload[16:20] = int.to_bytes(0, byteorder='big', length=4)
+    site_payload[20:28] = int.to_bytes(int(now.timestamp())*1000, byteorder='big', length=8)
+    site_payload[36:] = bytes(base64.b64decode(uid2))
+
+    id_payload = _encrypt_gcm(bytes(site_payload), None, site_key.secret)
+
+    master_payload = bytearray(256)
+    master_payload[:8] = int.to_bytes(int(token_expiry.timestamp())*1000, byteorder='big', length=8)
+    master_payload[8:16] = int.to_bytes(int(now.timestamp()), byteorder='big', length=8)
+    master_payload[16:20] = int.to_bytes(0, byteorder='big', length=4)
+    master_payload[20:21] = int.to_bytes(1, byteorder='big', length=1)
+    master_payload[21:25] = int.to_bytes(0, byteorder='big', length=4)
+    master_payload[25:29] = int.to_bytes(0, byteorder='big', length=4)
+    master_payload[29:33] = int.to_bytes(site_key.key_id, byteorder='big', length=4)
+    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)
+    root_writer[0:1] = int.to_bytes(0, byteorder='big', length=1)
+    root_writer[1:2] = int.to_bytes(112, 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)
+
+    return base64.b64encode(root_writer)
+
+
+
+def encrypt_key(uid2, indentity_scope, keys, keyset_id=None, **kwargs):
+    """ Encrypt an uid2 into a sharing token
+
+    Args:
+        uid2: the uid2 to be encrypted
+        keys (EncryptionKeysCollection): collection of keys to choose from for encryption
+        keyset_id (int) : An optional keyset id to use for the encryption. Will use default keyset if left blank
+
+    Returns (str): Sharing Token
+
+    """
+    now = kwargs.get("now")
+    if now is None:
+        now = dt.datetime.now(tz=timezone.utc)
+    key = keys.get_default_keyset_key(now) if keyset_id is None else keys.get_by_keyset_key(keyset_id, now)
+    master_key = keys.get_by_keyset_key(9999, now)
+
+    token_expiry = now + dt.timedelta(days=30) if keys.get_token_expiry_seconds() is None \
+        else now + dt.timedelta(seconds=keys.get_token_expiry_seconds())
+
+    site_id = keys.get_caller_site_id()
+    if site_id is None:
+        print("No Site ID in keys")
+        return
+
+    if key is None:
+        print("No Keyset Key found")
+        return
+
+    return _encrypt_token_v3(uid2, indentity_scope, master_key, key, site_id, now, token_expiry)
+
 
 def encrypt_data(data, identity_scope, **kwargs):
     """Encrypt arbitrary binary data.
@@ -307,7 +372,7 @@ def _add_pkcs7_padding(data, block_size):
 
 
 def _encrypt(data, iv, key):
-    cipher = AES.new(key.secret, AES.MODE_CBC, iv=iv)
+    cipher = AES.new(key.secret, AES.MODE_CBC, IV=iv)
     return cipher.encrypt(_add_pkcs7_padding(data, AES.block_size))
 
 

uid2_client/keys.py

@@ -6,7 +6,7 @@
 
 
 import datetime as dt
-from bisect import bisect_right
+from bisect import bisect_right, bisect_left
 
 
 class EncryptionKey:
@@ -21,14 +21,15 @@ class EncryptionKey:
         secret (bytes): the actual encryption key
     """
 
-    def __init__(self, key_id, site_id, created, activates, expires, secret):
+    def __init__(self, key_id, site_id, created, activates, expires, secret, keyset_id=None):
         """Create a new encryption key."""
         self._id = key_id
         self._site_id = site_id
         self._created = created
         self._activates = activates
         self._expires = expires
         self._secret = secret
+        self._keyset_id = keyset_id
 
 
     @property
@@ -42,6 +43,11 @@ def site_id(self):
         """int: id of the site the key belongs to."""
         return self._site_id
 
+    @property
+    def keyset_id(self):
+        """int: keyset id, can be None"""
+        return self._keyset_id
+
 
     @property
     def created(self):
@@ -94,20 +100,29 @@ class EncryptionKeysCollection:
     used for decoding UID2 advertising tokens.
     """
 
-    def __init__(self, keys):
+    def __init__(self, keys, caller_site_id=None, master_keyset_id=None, default_keyset_id=None, token_expiry_seconds=None):
         self._latest_expires = None
         self._keys = dict()
         self._keys_by_site = dict()
+        self._keys_by_keyset = dict()
+        self.set_keys(keys)
+        self._caller_site_id = caller_site_id
+        self._master_keyset_id = master_keyset_id
+        self._default_keyset_id = default_keyset_id
+        self._token_expiry_seconds = token_expiry_seconds
+
+    def set_keys(self, keys):
         for key in keys:
             self._keys[key.key_id] = key
             if key.site_id > 0:
                 self._keys_by_site.setdefault(key.site_id, []).append(key)
+            if key.keyset_id is not None:
+                self._keys_by_keyset.setdefault(key.keyset_id, []).append(key)
             if self._latest_expires is None or key.expires > self._latest_expires:
                 self._latest_expires = key.expires
         for _, site_keys in self._keys_by_site.items():
             site_keys.sort(key=lambda x: x.activates)
 
-
     def __len__(self):
         return len(self._keys)
 
@@ -119,6 +134,18 @@ def __contains__(self, key_id):
     def __getitem__(self, key_id):
         return self._keys[key_id]
 
+    def get_caller_site_id(self):
+        return self._caller_site_id
+
+    def get_master_keyset_id(self):
+        return self._master_keyset_id
+
+    def get_default_keyset_id(self):
+        return self._default_keyset_id
+
+    def get_token_expiry_seconds(self):
+        return self._token_expiry_seconds
+
 
     def get(self, key_id, default=None):
         """Get encryption key with the specified id, else default."""
@@ -135,6 +162,29 @@ def values(self):
         return self._keys.values()
 
 
+    def get_default_keyset_key(self, now):
+        return self.get_by_keyset_key(self._default_keyset_id, now)
+
+    def get_by_keyset_key(self, keyset_id, now):
+        """ Gets Active Key by keyset_id
+
+        Args:
+            keyset_id: the keyset id to get
+
+        Returns: EncryptionKey: active keyset key or None
+
+        """
+        keyset_keys = self._keys_by_keyset.get(keyset_id)
+        if keyset_keys is None or len(keyset_keys) == 0:
+            return None
+        i = bisect_right(_SiteKeyActivatesList(keyset_keys), now)
+        while i > 0:
+            i -= 1
+            key = keyset_keys[i]
+            if key.is_active(now):
+                return key
+        return None
+
     def get_active_site_key(self, site_id, now):
         """Get active encryption key for the specified site, else None.