feat: Decrypt input secrets if there are some (#45)

Co-authored-by: František Nesveda <[email protected]>

Author: drobnikj

setup.py

@@ -59,6 +59,7 @@
         'websockets ~= 10.4',
         'aiofiles ~= 22.1.0',
         'aioshutil ~= 1.2',
+        'cryptography ~= 39.0.0',
     ],
     extras_require={
         'dev': [

src/apify/_crypto.py

@@ -0,0 +1,144 @@
+import base64
+import secrets
+from typing import Any
+
+from cryptography.exceptions import InvalidTag as InvalidTagException
+from cryptography.hazmat.primitives import hashes, serialization
+from cryptography.hazmat.primitives.asymmetric import padding, rsa
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+
+from .consts import ENCRYPTED_INPUT_VALUE_REGEXP
+
+ENCRYPTION_KEY_LENGTH = 32
+ENCRYPTION_IV_LENGTH = 16
+ENCRYPTION_AUTH_TAG_LENGTH = 16
+
+
+def public_encrypt(value: str, *, public_key: rsa.RSAPublicKey) -> dict:
+    """Encrypts the given value using AES cipher and the password for encryption using the public key.
+
+    The encryption password is a string of encryption key and initial vector used for cipher.
+    It returns the encrypted password and encrypted value in BASE64 format.
+
+    Args:
+        value (str): Password used to encrypt the private key encoded as base64 string.
+        public_key (RSAPublicKey): Private key to use for decryption.
+
+    Returns:
+        disc: Encrypted password and value.
+    """
+    key_bytes = _crypto_random_object_id(ENCRYPTION_KEY_LENGTH).encode('utf-8')
+    initialized_vector_bytes = _crypto_random_object_id(ENCRYPTION_IV_LENGTH).encode('utf-8')
+    value_bytes = value.encode('utf-8')
+
+    password_bytes = key_bytes + initialized_vector_bytes
+
+    # NOTE: Auth Tag is appended to the end of the encrypted data, it has length of 16 bytes and ensures integrity of the data.
+    cipher = Cipher(algorithms.AES(key_bytes), modes.GCM(initialized_vector_bytes, min_tag_length=ENCRYPTION_AUTH_TAG_LENGTH))
+    encryptor = cipher.encryptor()
+    encrypted_value_bytes = encryptor.update(value_bytes) + encryptor.finalize()
+    encrypted_password_bytes = public_key.encrypt(
+        password_bytes,
+        padding.OAEP(
+            mgf=padding.MGF1(algorithm=hashes.SHA1()),
+            algorithm=hashes.SHA1(),
+            label=None,
+        ),
+    )
+    return {
+        'encrypted_value': base64.b64encode(encrypted_value_bytes + encryptor.tag).decode('utf-8'),
+        'encrypted_password': base64.b64encode(encrypted_password_bytes).decode('utf-8'),
+    }
+
+
+def private_decrypt(
+    encrypted_password: str,
+    encrypted_value: str,
+    *,
+    private_key: rsa.RSAPrivateKey,
+) -> str:
+    """Decrypts the given encrypted value using the private key and password.
+
+    Args:
+        encrypted_password (str): Password used to encrypt the private key encoded as base64 string.
+        encrypted_value (str): Encrypted value to decrypt as base64 string.
+        private_key (RSAPrivateKey): Private key to use for decryption.
+
+    Returns:
+        str: Decrypted value.
+    """
+    encrypted_password_bytes = base64.b64decode(encrypted_password.encode('utf-8'))
+    encrypted_value_bytes = base64.b64decode(encrypted_value.encode('utf-8'))
+
+    # Decrypt the password
+    password_bytes = private_key.decrypt(
+        encrypted_password_bytes,
+        padding.OAEP(
+            mgf=padding.MGF1(algorithm=hashes.SHA1()),
+            algorithm=hashes.SHA1(),
+            label=None,
+        ),
+    )
+
+    if len(password_bytes) != ENCRYPTION_KEY_LENGTH + ENCRYPTION_IV_LENGTH:
+        raise ValueError('Decryption failed, invalid password length!')
+
+    # Slice the encrypted into cypher and authentication tag
+    authentication_tag_bytes = encrypted_value_bytes[-ENCRYPTION_AUTH_TAG_LENGTH:]
+    encrypted_data_bytes = encrypted_value_bytes[:len(encrypted_value_bytes) - ENCRYPTION_AUTH_TAG_LENGTH]
+    encryption_key_bytes = password_bytes[:ENCRYPTION_KEY_LENGTH]
+    initialization_vector_bytes = password_bytes[ENCRYPTION_KEY_LENGTH:]
+
+    try:
+        cipher = Cipher(algorithms.AES(encryption_key_bytes), modes.GCM(initialization_vector_bytes, authentication_tag_bytes))
+        decryptor = cipher.decryptor()
+        decipher_bytes = decryptor.update(encrypted_data_bytes) + decryptor.finalize()
+    except InvalidTagException:
+        raise ValueError('Decryption failed, malformed encrypted value or password.')
+    except Exception as err:
+        raise err
+
+    return decipher_bytes.decode('utf-8')
+
+
+def _load_private_key(private_key_file_base64: str, private_key_password: str) -> rsa.RSAPrivateKey:
+    private_key = serialization.load_pem_private_key(base64.b64decode(
+        private_key_file_base64.encode('utf-8')), password=private_key_password.encode('utf-8'))
+    if not isinstance(private_key, rsa.RSAPrivateKey):
+        raise ValueError('Invalid private key.')
+
+    return private_key
+
+
+def _load_public_key(public_key_file_base64: str) -> rsa.RSAPublicKey:
+    public_key = serialization.load_pem_public_key(base64.b64decode(public_key_file_base64.encode('utf-8')))
+    if not isinstance(public_key, rsa.RSAPublicKey):
+        raise ValueError('Invalid public key.')
+
+    return public_key
+
+
+def _crypto_random_object_id(length: int = 17) -> str:
+    """Python reimplementation of cryptoRandomObjectId from `@apify/utilities`."""
+    chars = 'abcdefghijklmnopqrstuvwxyzABCEDFGHIJKLMNOPQRSTUVWXYZ0123456789'
+    return ''.join(secrets.choice(chars) for _ in range(length))
+
+
+def _decrypt_input_secrets(private_key: rsa.RSAPrivateKey, input: Any) -> Any:
+    """Decrypt input secrets."""
+    if not isinstance(input, dict):
+        return input
+
+    for key, value in input.items():
+        if isinstance(value, str):
+            match = ENCRYPTED_INPUT_VALUE_REGEXP.fullmatch(value)
+            if match:
+                encrypted_password = match.group(1)
+                encrypted_value = match.group(2)
+                input[key] = private_decrypt(
+                    encrypted_password,
+                    encrypted_value,
+                    private_key=private_key,
+                )
+
+    return input

src/apify/_utils.py

@@ -10,7 +10,6 @@
 import mimetypes
 import os
 import re
-import secrets
 import sys
 import time
 from collections import OrderedDict
@@ -324,12 +323,6 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
     return cast(MetadataType, wrapper)
 
 
-def _crypto_random_object_id(length: int = 17) -> str:
-    """Python reimplementation of cryptoRandomObjectId from `@apify/utilities`."""
-    chars = 'abcdefghijklmnopqrstuvwxyzABCEDFGHIJKLMNOPQRSTUVWXYZ0123456789'
-    return ''.join(secrets.choice(chars) for _ in range(length))
-
-
 T = TypeVar('T')
 
 

src/apify/actor.py

@@ -10,6 +10,7 @@
 from apify_client import ApifyClientAsync
 from apify_client.consts import WebhookEventType
 
+from ._crypto import _decrypt_input_secrets, _load_private_key
 from ._utils import (
     _fetch_and_parse_env_var,
     _get_cpu_usage_percent,
@@ -584,9 +585,17 @@ async def get_input(cls) -> Any:
     async def _get_input_internal(self) -> Any:
         self._raise_if_not_initialized()
 
-        # TODO: decryption
+        input_value = await self.get_value(self._config.input_key)
+        input_secrets_private_key = self._config.input_secrets_private_key_file
+        input_secrets_key_passphrase = self._config.input_secrets_private_key_passphrase
+        if input_secrets_private_key and input_secrets_key_passphrase:
+            private_key = _load_private_key(
+                input_secrets_private_key,
+                input_secrets_key_passphrase,
+            )
+            input_value = _decrypt_input_secrets(private_key, input_value)
 
-        return await self.get_value(self._config.input_key)
+        return input_value
 
     @classmethod
     async def get_value(cls, key: str) -> Any:

src/apify/consts.py

@@ -1,3 +1,4 @@
+import re
 from enum import Enum
 from typing import List, Literal, get_args
 
@@ -159,3 +160,7 @@ class StorageTypes(str, Enum):
 REQUEST_QUEUE_HEAD_MAX_LIMIT = 1000
 
 EVENT_LISTENERS_TIMEOUT_SECS = 5
+
+BASE64_REGEXP = '[-A-Za-z0-9+/]*={0,3}'
+ENCRYPTED_INPUT_VALUE_PREFIX = 'ENCRYPTED_VALUE'
+ENCRYPTED_INPUT_VALUE_REGEXP = re.compile(f'^{ENCRYPTED_INPUT_VALUE_PREFIX}:({BASE64_REGEXP}):({BASE64_REGEXP})$')

src/apify/storages/request_queue.py

@@ -10,7 +10,8 @@
 from apify_client import ApifyClientAsync
 from apify_client.clients import RequestQueueClientAsync
 
-from .._utils import LRUCache, _budget_ow, _crypto_random_object_id, _unique_key_to_request_id
+from .._crypto import _crypto_random_object_id
+from .._utils import LRUCache, _budget_ow, _unique_key_to_request_id
 from ..config import Configuration
 from ..consts import REQUEST_QUEUE_HEAD_MAX_LIMIT
 from ..memory_storage import MemoryStorage

tests/integration/_utils.py

@@ -1,4 +1,4 @@
-from apify._utils import _crypto_random_object_id
+from apify._crypto import _crypto_random_object_id
 
 
 def generate_unique_resource_name(label: str) -> str:

tests/integration/test_actor_api_helpers.py

@@ -2,7 +2,7 @@
 import json
 
 from apify import Actor
-from apify._utils import _crypto_random_object_id
+from apify._crypto import _crypto_random_object_id
 from apify_client import ApifyClientAsync
 
 from ._utils import generate_unique_resource_name

tests/integration/test_fixtures.py

@@ -1,7 +1,7 @@
 from datetime import datetime, timezone
 
 from apify import Actor
-from apify._utils import _crypto_random_object_id
+from apify._crypto import _crypto_random_object_id
 from apify_client import ApifyClientAsync
 
 from .conftest import ActorFactory

tests/unit/actor/test_actor_key_value_store.py

@@ -1,9 +1,13 @@
 import pytest
 
 from apify import Actor
+from apify._crypto import public_encrypt
 from apify._utils import _json_dumps
+from apify.consts import ENCRYPTED_INPUT_VALUE_PREFIX, ApifyEnvVars
 from apify.memory_storage import MemoryStorage
 
+from ..test_crypto import PRIVATE_KEY_PASSWORD, PRIVATE_KEY_PEM_BASE64, PUBLIC_KEY
+
 
 # NOTE: We only test the key-value store methond available on Actor class/instance. Actual tests for the implementations are in storages/.
 class TestOpenKeyValueStore:
@@ -42,3 +46,24 @@ async def test_get_input(self, memory_storage: MemoryStorage) -> None:
         async with Actor() as my_actor:
             input = await my_actor.get_input()
             assert input['foo'] == test_input['foo']
+
+    async def test_get_input_with_secrets(self, memory_storage: MemoryStorage, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.setenv(ApifyEnvVars.INPUT_SECRETS_PRIVATE_KEY_FILE, PRIVATE_KEY_PEM_BASE64)
+        monkeypatch.setenv(ApifyEnvVars.INPUT_SECRETS_PRIVATE_KEY_PASSPHRASE, PRIVATE_KEY_PASSWORD)
+        input_key = 'INPUT'
+        secret_string = 'secret-string'
+        encrypted_secret = public_encrypt(secret_string, public_key=PUBLIC_KEY)
+        input_with_secret = {
+            'foo': 'bar',
+            'secret': f'{ENCRYPTED_INPUT_VALUE_PREFIX}:{encrypted_secret["encrypted_password"]}:{encrypted_secret["encrypted_value"]}',
+        }
+        kvs_info = await memory_storage.key_value_stores().get_or_create(name='default')
+        await memory_storage.key_value_store(kvs_info['id']).set_record(
+            key=input_key,
+            value=_json_dumps(input_with_secret),
+            content_type='application/json',
+        )
+        async with Actor() as my_actor:
+            input = await my_actor.get_input()
+            assert input['foo'] == input_with_secret['foo']
+            assert input['secret'] == secret_string