feat: add create_storage_content_signature

Author: danpoletaev

src/apify_shared/utils.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
+import base64
 import contextlib
 import hashlib
 import hmac
 import io
 import json
 import re
 import string
+import time
 from datetime import datetime, timezone
 from enum import Enum
 from typing import Any, TypeVar, cast
@@ -153,3 +155,24 @@ def create_hmac_signature(secret_key: str, message: str) -> str:
     decimal_signature = int(signature, 16)
 
     return encode_base62(decimal_signature)
+
+
+def create_storage_content_signature(
+    resource_id: str, url_signing_secret_key: str, expires_in_millis: int | None = None, version: int = 0
+) -> str:
+    """Create a secure signature for a resource like a dataset or key-value store.
+
+    This signature is used to generate a signed URL for authenticated access, which can be expiring or permanent.
+    The signature is created using HMAC with the provided secret key and includes
+    the resource ID, expiration time, and version.
+
+    Note: expires_in_millis is optional. If not provided, the signature will not expire.
+
+    """
+    expires_at = int(time.time() * 1000) + expires_in_millis if expires_in_millis else 0
+
+    message_to_sign = f'{version}.{expires_at}.{resource_id}'
+    hmac = create_hmac_signature(url_signing_secret_key, message_to_sign)
+
+    base64url_encoded_payload = base64.urlsafe_b64encode(f'{version}.{expires_at}.{hmac}'.encode())
+    return base64url_encoded_payload.decode('utf-8')

tests/unit/test_utils.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
+import base64
 import io
 from datetime import datetime, timezone
 from enum import Enum
 
 from apify_shared.utils import (
     create_hmac_signature,
+    create_storage_content_signature,
     encode_base62,
     filter_out_none_values_recursively,
     filter_out_none_values_recursively_internal,
@@ -171,3 +173,38 @@ def test_create_same_hmac() -> None:
     message = 'hmac-same-message-to-be-authenticated'
     assert create_hmac_signature(secret_key, message) == 'FYMcmTIm3idXqleF1Sw5'
     assert create_hmac_signature(secret_key, message) == 'FYMcmTIm3idXqleF1Sw5'
+
+
+# This test ensures compatibility with the JavaScript version of the same method.
+# https://github.com/apify/apify-shared-js/blob/master/packages/utilities/src/storages.ts
+def test_create_storage_content_signature() -> None:
+    # This test uses the same parameters as in JS tests.
+    secret_key = 'hmac-secret-key'
+    message = 'resource-id'
+
+    signature = create_storage_content_signature(
+        resource_id=message,
+        url_signing_secret_key=secret_key,
+    )
+
+    version, expires_at, hmac = base64.urlsafe_b64decode(signature).decode('utf-8').split('.')
+
+    assert signature == 'MC4wLjNUd2ZFRTY1OXVmU05zbVM0N2xS'
+    assert version == '0'
+    assert expires_at == '0'
+    assert hmac == '3TwfEE659ufSNsmS47lR'
+
+
+def test_create_storage_content_signature_with_expiration() -> None:
+    secret_key = 'hmac-secret-key'
+    message = 'resource-id'
+
+    signature = create_storage_content_signature(
+        resource_id=message,
+        url_signing_secret_key=secret_key,
+        expires_in_millis=10000,
+    )
+
+    version, expires_at, hmac = base64.urlsafe_b64decode(signature).decode('utf-8').split('.')
+    assert version == '0'
+    assert expires_at != '0'