Add optional parameter for encrypt/decrypt to specify timestamp

Author: asloobq

uid2_client/bid_stream_client.py

@@ -2,9 +2,10 @@
 >>> from uid2_client import BidStreamClient
 """
 import base64
+import datetime as dt
 
-from .encryption import decrypt_token
 from .client_type import ClientType
+from .encryption import decrypt_token
 from .refresh_keys_util import refresh_bidstream_keys
 
 
@@ -35,12 +36,13 @@ def __init__(self, base_url, auth_key, secret_key):
         self._auth_key = auth_key
         self._secret_key = base64.b64decode(secret_key)
 
-    def decrypt_ad_token_into_raw_uid(self, token, domain_name):
+    def decrypt_ad_token_into_raw_uid(self, token, domain_name, now=dt.datetime.now(tz=dt.timezone.utc)):
         """Decrypt advertising token to extract UID2 details.
 
             Args:
                 token (str): advertising token to decrypt
                 domain_name (str) : domain name from bid request
+                now (datetime): date/time to use as "now" when doing token expiration check
 
             Returns:
                 DecryptedToken: details extracted from the advertising token
@@ -49,7 +51,7 @@ def decrypt_ad_token_into_raw_uid(self, token, domain_name):
                 EncryptionError: if token version is not supported, the token has expired,
                                  or no required decryption keys present in the keys collection
         """
-        return decrypt_token(token, self._keys, domain_name, ClientType.Bidstream)
+        return decrypt_token(token, self._keys, domain_name, ClientType.Bidstream, now)
 
     def refresh_keys(self):
         """Get the latest encryption keys for advertising tokens.

uid2_client/encryption.py

@@ -88,6 +88,8 @@ def decrypt_token(token, keys, domain_name, client_type, now=dt.datetime.now(tz=
 
 
 def _decrypt_token(token, keys, domain_name, client_type, now):
+    if keys is None:
+        raise EncryptionError('keys not initialized')
     if not keys.valid(now):
         raise EncryptionError('no keys available or all keys have expired; refresh the latest keys from UID2 service')
 

uid2_client/sharing_client.py

@@ -2,9 +2,10 @@
 >>> from uid2_client import SharingClient
 """
 import base64
+import datetime as dt
 
-from .encryption import encrypt, decrypt_token
 from .client_type import ClientType
+from .encryption import encrypt, decrypt_token
 from .refresh_keys_util import refresh_sharing_keys
 
 
@@ -36,23 +37,25 @@ def __init__(self, base_url, auth_key, secret_key):
         self._auth_key = auth_key
         self._secret_key = base64.b64decode(secret_key)
 
-    def encrypt_raw_uid_into_sharing_token(self, uid2, keyset_id=None):
+    def encrypt_raw_uid_into_sharing_token(self, uid2, keyset_id=None, now=dt.datetime.now(tz=dt.timezone.utc)):
         """ Encrypt a UID2 into a sharing token
 
             Args:
                 uid2: the UID2 or EUID 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
+                now (Datetime): the datettime to use for now. Defaults to utc now
 
             Returns (str): Sharing Token
             """
-        return encrypt(uid2, None, self._keys, keyset_id)
+        return encrypt(uid2, None, self._keys, keyset_id, now=now)
 
-    def decrypt_sharing_token_into_raw_uid(self, token):
+    def decrypt_sharing_token_into_raw_uid(self, token, now=dt.datetime.now(tz=dt.timezone.utc)):
         """Decrypt sharing token to extract UID2 details.
 
             Args:
                 token (str): sharing token to decrypt
+                now (datetime): date/time to use as "now" when doing token expiration check
 
             Returns:
                 DecryptedToken: details extracted from the sharing token
@@ -61,7 +64,7 @@ def decrypt_sharing_token_into_raw_uid(self, token):
                 EncryptionError: if token version is not supported, the token has expired,
                                  or no required decryption keys present in the keys collection
         """
-        return decrypt_token(token, self._keys, None, ClientType.Sharing)
+        return decrypt_token(token, self._keys, None, ClientType.Sharing, now)
 
     def refresh_keys(self):
         """Get the latest encryption keys for sharing tokens.