From c19a01ee6748aa2b298dd692dfad089f92c2198e Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 13 Nov 2024 02:09:30 +0000 Subject: [PATCH 1/3] Add example crypt-kv --- examples/python/crypt-kv/.gitignore | 1 + examples/python/crypt-kv/README.md | 79 ++++++++++ examples/python/crypt-kv/crypt_kv.py | 182 ++++++++++++++++++++++ examples/python/crypt-kv/test_crypt_kv.py | 107 +++++++++++++ examples/python/crypt-kv/test_s3.py | 39 +++++ 5 files changed, 408 insertions(+) create mode 100644 examples/python/crypt-kv/.gitignore create mode 100644 examples/python/crypt-kv/README.md create mode 100644 examples/python/crypt-kv/crypt_kv.py create mode 100644 examples/python/crypt-kv/test_crypt_kv.py create mode 100644 examples/python/crypt-kv/test_s3.py diff --git a/examples/python/crypt-kv/.gitignore b/examples/python/crypt-kv/.gitignore new file mode 100644 index 000000000..1d17dae13 --- /dev/null +++ b/examples/python/crypt-kv/.gitignore @@ -0,0 +1 @@ +.venv diff --git a/examples/python/crypt-kv/README.md b/examples/python/crypt-kv/README.md new file mode 100644 index 000000000..b0de9e244 --- /dev/null +++ b/examples/python/crypt-kv/README.md @@ -0,0 +1,79 @@ +# Crypt-KV + +A secure key-value store with encryption and AWS S3 storage support. This library provides encrypted storage of key-value pairs with key versioning and flexible storage backends. + +## Features + +- Encrypted key-value storage using Fernet symmetric encryption +- Merkle tree for data integrity verification +- AWS S3 storage backend +- Flexible storage interface supporting different backends + +## Quick Start + +```python +import boto3 +from crypt_kv import Boto3Db, SimpleKeyProvider, CryptClient + +# Initialize S3 client +client = boto3.client('s3', region_name='us-east-1') + +# Create storage backend +db = Boto3Db("my-bucket", client) + +# Initialize key provider and crypto client +key_provider = SimpleKeyProvider(b"my-root-key") +crypt_client = CryptClient(key_provider, db=db) + +# Store encrypted values +crypt_client.set("my-key", b"secret data") + +# Retrieve and decrypt values +value = crypt_client.get("my-key") +``` + +## Examples + +### Basic Usage + +```python +# Create client with in-memory storage +provider = SimpleKeyProvider(b"test-root-key") +client = CryptClient(provider) + +# Store and retrieve values +client.set("key1", b"secret value") +value = client.get("key1") +``` + +### Persistence Across Sessions + +```python +# First session +client1 = CryptClient(provider, db=storage) +client1.set("key1", b"value1") +root_hash = client1.root_hash + +# Later session +client2 = CryptClient(provider, root_hash=root_hash, db=storage) +value = client2.get("key1") # Retrieves original value +``` + +### AWS S3 Storage + +```python +# Initialize S3 backend +s3_client = boto3.client('s3', region_name='us-east-1') +db = Boto3Db("my-bucket", s3_client) + +# Create crypto client with S3 storage +client = CryptClient(key_provider, db=db) +``` + +## Testing + +Run tests using pytest: + +```bash +pytest test_crypt_kv.py test_s3.py +``` diff --git a/examples/python/crypt-kv/crypt_kv.py b/examples/python/crypt-kv/crypt_kv.py new file mode 100644 index 000000000..598c803df --- /dev/null +++ b/examples/python/crypt-kv/crypt_kv.py @@ -0,0 +1,182 @@ +import hashlib +import base64 + +from typing import Optional +from trie import HexaryTrie +from cryptography.fernet import Fernet + + +class SimpleKeyProvider: + """Manages encryption key versioning and derivation. + + Attributes: + current_key_version: Current version number for key rotation + root_key: Master key used to derive encryption keys + """ + + current_key_version: int + root_key: bytes + + def __init__(self, root_key: bytes): + self.current_key_version = 0 + self.root_key = root_key + + def get_key(self, path: bytes, version: Optional[int] = None) -> tuple[int, bytes]: + """Gets the encryption key for a given path and version. + + Args: + path: The path/identifier to derive a key for + version: Optional key version, defaults to current version + + Returns: + Tuple of (key version number, derived encryption key) + """ + version = self.current_key_version if version is None else version + derived_key = self._derive_key(path, version) + return (version, derived_key) + + def _derive_key(self, path: bytes, version: int) -> bytes: + salt = str(version).encode() + b"|" + path + return hashlib.pbkdf2_hmac( + 'sha256', + password=self.root_key, + salt=salt, + iterations=100000, # OWASP recommended minimum + dklen=32 # 256-bit key + ) + + +class DstackKeyProvider(SimpleKeyProvider): + """Key provider that derives encryption keys from Dstack's key management service.""" + + def __init__(self, key_id: str): + from tappd_client import TappdClient + self.client = TappdClient() + pem = self.client.derive_key(key_id).key + super().__init__(hashlib.sha256(pem.encode()).digest()) + + +class Boto3Db: + """Database adapter for storing encrypted data in S3 using boto3. + + Provides a dict-like interface for storing and retrieving bytes from S3. + Keys and values are stored as bytes and base64 encoded for S3 compatibility. + + Attributes: + bucket: Name of the S3 bucket to store data in + boto3: Boto3 S3 client instance for making API calls + """ + + def __init__(self, bucket, boto3_client): + self.bucket = bucket + self.boto3 = boto3_client + + def __getitem__(self, key: bytes) -> bytes: + key = base64.urlsafe_b64encode(key).decode() + return self.boto3.get_object(Bucket=self.bucket, Key=key)["Body"].read() + + def __setitem__(self, key: bytes, value: bytes): + key = base64.urlsafe_b64encode(key).decode() + self.boto3.put_object(Bucket=self.bucket, Key=key, Body=value) + + +class FernetCryptor: + """Encrypts and decrypts data using Fernet symmetric encryption.""" + + def encrypt(self, key: bytes, value: bytes) -> bytes: + f = Fernet(base64.urlsafe_b64encode(key)) + return f.encrypt(value) + + def decrypt(self, key: bytes, value: bytes) -> bytes: + f = Fernet(base64.urlsafe_b64encode(key)) + return f.decrypt(value) + + +class CryptClient: + """Client for encrypting and storing key-value pairs. + + Attributes: + db: HexaryTrie instance for storing encrypted data + key_provider: SimpleKeyProvider instance for key management + """ + + db: HexaryTrie + key_provider: SimpleKeyProvider + cryptor: FernetCryptor + + def __init__(self, key_provider: SimpleKeyProvider, root_hash=None, db=None, cryptor=None): + self.key_provider = key_provider + self.cryptor = cryptor or FernetCryptor() + if db is None: + db = {} + if root_hash is None: + self.db = HexaryTrie(db=db) + else: + self.db = HexaryTrie(db=db, root_hash=root_hash) + + @property + def root_hash(self) -> bytes: + return self.db.root_hash + + def _encode_key(self, key: 'str | bytes') -> bytes: + if isinstance(key, str): + return key.encode() + return key + + def get(self, key: str) -> bytes: + """Get the decrypted value for a key. + + Args: + key: The key to look up + + Returns: + The decrypted value as bytes + + Raises: + KeyError: If the key does not exist + """ + key = self._encode_key(key) + if not self.exists(key): + raise KeyError(f"Key '{key}' not found") + packed_data = self.db.get(key) + # Unpack version (first 4 bytes) and encrypted value + kver = int.from_bytes(packed_data[:4], 'big') + crypted_value = packed_data[4:] + (_, crypt_key) = self.key_provider.get_key(key, version=kver) + decrypted_value = self.cryptor.decrypt(crypt_key, crypted_value) + return decrypted_value + + def set(self, key: str, value: bytes): + """Set a key-value pair, encrypting the value. + + Args: + key: The key to set + value: The value to encrypt and store + """ + key = self._encode_key(key) + (kver, crypt_key) = self.key_provider.get_key(key) + crypted_value = self.cryptor.encrypt(crypt_key, value) + # Pack version and encrypted value into single bytes object + packed_data = kver.to_bytes(4, 'big') + crypted_value + self.db.set(key, packed_data) + + def exists(self, key: str) -> bool: + """Check if a key exists. + + Args: + key: The key to check + + Returns: + True if the key exists, False otherwise + """ + key = self._encode_key(key) + return self.db.exists(key) + + def delete(self, key: str): + """Delete a key-value pair. + + Args: + key: The key to delete + """ + key = self._encode_key(key) + self.db.delete(key) diff --git a/examples/python/crypt-kv/test_crypt_kv.py b/examples/python/crypt-kv/test_crypt_kv.py new file mode 100644 index 000000000..e8c26c6ba --- /dev/null +++ b/examples/python/crypt-kv/test_crypt_kv.py @@ -0,0 +1,107 @@ +import pytest +from crypt_kv import SimpleKeyProvider, CryptClient + +def test_key_provider(): + # Test key derivation and versioning + provider = SimpleKeyProvider(b"test-root-key") + + # Test default version + version, key1 = provider.get_key(b"test/path") + assert version == 0 + assert isinstance(key1, bytes) + assert len(key1) == 32 # SHA-256 produces 32 byte keys + + # Test key derivation is deterministic + version, key2 = provider.get_key(b"test/path") + assert key1 == key2 + + # Test different paths produce different keys + _, other_key = provider.get_key(b"other/path") + assert other_key != key1 + +def test_crypt_client_basic(): + provider = SimpleKeyProvider(b"test-root-key") + client = CryptClient(provider) + + # Test basic set/get + test_value = b"hello world" + client.set("test_key", test_value) + retrieved = client.get("test_key") + assert retrieved == test_value + +def test_crypt_client_multiple_values(): + provider = SimpleKeyProvider(b"test-root-key") + client = CryptClient(provider) + + # Test multiple values + values = { + "key1": b"value1", + "key2": b"value2", + "key3": b"value3" + } + + for k, v in values.items(): + client.set(k, v) + + for k, v in values.items(): + assert client.get(k) == v + +def test_persistence(): + provider = SimpleKeyProvider(b"test-root-key") + db = {} # Shared db dictionary + + # First client sets values + client1 = CryptClient(provider, db=db) + client1.set("persistent", b"test data") + client1.set("key2", b"test data 2") + client1.set("key3", b"test data 3") + root_hash = client1.root_hash + # Second client with same root hash should read same values + client2 = CryptClient(provider, root_hash=root_hash, db=db) + assert client2.get("persistent") == b"test data" + assert client2.get("key2") == b"test data 2" + assert client2.get("key3") == b"test data 3" + +def test_encryption_actually_encrypts(): + provider = SimpleKeyProvider(b"test-root-key") + client = CryptClient(provider) + + test_value = b"secret message" + client.set("secret", test_value) + + # Get raw encrypted value from trie + packed_data = client.db.get(b"secret") + assert isinstance(packed_data, bytes) + + # Unpack version and encrypted value + version = int.from_bytes(packed_data[:4], 'big') + encrypted_value = packed_data[4:] + + # Verify the stored data is actually encrypted + assert version == provider.current_key_version + assert encrypted_value != test_value + assert b"secret" not in encrypted_value + +def test_invalid_key(): + provider = SimpleKeyProvider(b"test-root-key") + client = CryptClient(provider) + + with pytest.raises(KeyError): + client.get("nonexistent_key") + +def test_delete(): + provider = SimpleKeyProvider(b"test-root-key") + client = CryptClient(provider) + + # Set and verify value exists + test_value = b"to be deleted" + client.set("temp", test_value) + assert client.exists("temp") == True + + # Delete and verify it's gone + client.delete("temp") + assert client.exists("temp") == False + + # Verify get raises KeyError + with pytest.raises(KeyError): + client.get("temp") diff --git a/examples/python/crypt-kv/test_s3.py b/examples/python/crypt-kv/test_s3.py new file mode 100644 index 000000000..110754c1d --- /dev/null +++ b/examples/python/crypt-kv/test_s3.py @@ -0,0 +1,39 @@ +import boto3 +from moto import mock_aws +from crypt_kv import Boto3Db, SimpleKeyProvider, CryptClient + + +@mock_aws +def test_s3_put(): + conn = boto3.resource("s3", region_name="us-east-1") + conn.create_bucket(Bucket="mybucket") + + client = boto3.client("s3", region_name="us-east-1") + db = Boto3Db("mybucket", client) + key_provider = SimpleKeyProvider(b"keyring1") + crypt_client = CryptClient(key_provider, db=db) + + # set some values + crypt_client.set(b"TEE", b"is awesome") + crypt_client.set(b"foo", b"bar") + + # list objects + response = client.list_objects_v2(Bucket="mybucket") + n_keys = len([_ for _ in response["Contents"]]) + assert n_keys == 4 + + # Create a new client instance + root_hash = crypt_client.root_hash + client = boto3.client("s3", region_name="us-east-1") + db = Boto3Db("mybucket", client) + crypt_client2 = CryptClient(SimpleKeyProvider(b"keyring1"), root_hash=root_hash, db=db) + + # get the set values + val_tee = crypt_client2.get(b"TEE") + assert val_tee == b"is awesome" + val_foo = crypt_client2.get(b"foo") + assert val_foo == b"bar" + print(f"tee={val_tee}") + print(f"foo={val_foo}") + +test_s3_put() From 45d1d555c58b5f5c77780199ca294a8fac66655d Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 13 Nov 2024 02:45:22 +0000 Subject: [PATCH 2/3] Reduce directory level --- examples/{python => }/crypt-kv/.gitignore | 0 examples/{python => }/crypt-kv/README.md | 0 examples/{python => }/crypt-kv/crypt_kv.py | 0 examples/{python => }/crypt-kv/test_crypt_kv.py | 0 examples/{python => }/crypt-kv/test_s3.py | 0 5 files changed, 0 insertions(+), 0 deletions(-) rename examples/{python => }/crypt-kv/.gitignore (100%) rename examples/{python => }/crypt-kv/README.md (100%) rename examples/{python => }/crypt-kv/crypt_kv.py (100%) rename examples/{python => }/crypt-kv/test_crypt_kv.py (100%) rename examples/{python => }/crypt-kv/test_s3.py (100%) diff --git a/examples/python/crypt-kv/.gitignore b/examples/crypt-kv/.gitignore similarity index 100% rename from examples/python/crypt-kv/.gitignore rename to examples/crypt-kv/.gitignore diff --git a/examples/python/crypt-kv/README.md b/examples/crypt-kv/README.md similarity index 100% rename from examples/python/crypt-kv/README.md rename to examples/crypt-kv/README.md diff --git a/examples/python/crypt-kv/crypt_kv.py b/examples/crypt-kv/crypt_kv.py similarity index 100% rename from examples/python/crypt-kv/crypt_kv.py rename to examples/crypt-kv/crypt_kv.py diff --git a/examples/python/crypt-kv/test_crypt_kv.py b/examples/crypt-kv/test_crypt_kv.py similarity index 100% rename from examples/python/crypt-kv/test_crypt_kv.py rename to examples/crypt-kv/test_crypt_kv.py diff --git a/examples/python/crypt-kv/test_s3.py b/examples/crypt-kv/test_s3.py similarity index 100% rename from examples/python/crypt-kv/test_s3.py rename to examples/crypt-kv/test_s3.py From 3c62d92a4a8a7c87005415c0c83c289ca63a95e6 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 13 Nov 2024 07:59:07 +0000 Subject: [PATCH 3/3] crypt-kv: Add ipynb --- examples/crypt-kv/crypt_kv.ipynb | 91 ++++++++++++++++++++++++++++++++ 1 file changed, 91 insertions(+) create mode 100644 examples/crypt-kv/crypt_kv.ipynb diff --git a/examples/crypt-kv/crypt_kv.ipynb b/examples/crypt-kv/crypt_kv.ipynb new file mode 100644 index 000000000..645dcfd3c --- /dev/null +++ b/examples/crypt-kv/crypt_kv.ipynb @@ -0,0 +1,91 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "3edbe311-20c8-43ae-a315-f6ec5102afc2", + "metadata": {}, + "outputs": [], + "source": [ + "!pip install boto3 moto tappd-client cryptography trie \"eth-hash[pycryptodome]\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e5ecef0a-518c-489d-8a43-fadadc6cf33f", + "metadata": {}, + "outputs": [], + "source": [ + "import boto3\n", + "from moto import mock_aws\n", + "from crypt_kv import Boto3Db, DstackKeyProvider, CryptClient\n", + "\n", + "\n", + "@mock_aws\n", + "def test_s3_put():\n", + " conn = boto3.resource(\"s3\", region_name=\"us-east-1\")\n", + " conn.create_bucket(Bucket=\"mybucket\")\n", + "\n", + " client = boto3.client(\"s3\", region_name=\"us-east-1\")\n", + " db = Boto3Db(\"mybucket\", client)\n", + " key_provider = DstackKeyProvider(\"key\")\n", + " crypt_client = CryptClient(key_provider, db=db)\n", + "\n", + " # set some values\n", + " crypt_client.set(b\"TEE\", b\"is awesome\")\n", + " crypt_client.set(b\"foo\", b\"bar\")\n", + "\n", + " # list objects\n", + " response = client.list_objects_v2(Bucket=\"mybucket\")\n", + " n_keys = len([_ for _ in response[\"Contents\"]])\n", + " assert n_keys == 4\n", + "\n", + " # Create a new client instance\n", + " root_hash = crypt_client.root_hash\n", + " client = boto3.client(\"s3\", region_name=\"us-east-1\")\n", + " db = Boto3Db(\"mybucket\", client)\n", + " crypt_client2 = CryptClient(DstackKeyProvider(\"key\"), root_hash=root_hash, db=db)\n", + "\n", + " # get the set values\n", + " val_tee = crypt_client2.get(b\"TEE\")\n", + " assert val_tee == b\"is awesome\"\n", + " val_foo = crypt_client2.get(b\"foo\")\n", + " assert val_foo == b\"bar\"\n", + " print(f\"tee={val_tee}\")\n", + " print(f\"foo={val_foo}\")\n", + "\n", + "test_s3_put()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "712d5dea-de1e-42f1-bd00-d0a5648ef180", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.7" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}