diff --git a/.gitignore b/.gitignore index 26b2b77..d69001d 100644 --- a/.gitignore +++ b/.gitignore @@ -3,3 +3,7 @@ /coti_sdk.egg-info/ /build/ +.venv/ +__pycache__/ +*.pyc +Comparisson.md diff --git a/coti/crypto_utils.py b/coti/crypto_utils.py index 08b69fa..9ce626d 100644 --- a/coti/crypto_utils.py +++ b/coti/crypto_utils.py @@ -5,7 +5,7 @@ from cryptography.hazmat.primitives.asymmetric import padding from cryptography.hazmat.primitives.asymmetric import rsa from eth_keys import keys -from .types import CtString, CtUint, ItString, ItUint +from .types import CtUint256, ItUint256, CtString, CtUint, ItString, ItUint block_size = AES.block_size address_size = 20 @@ -94,6 +94,46 @@ def sign_input_text(sender_address: str, contract_address: str, function_selecto return sign(message, key) +def sign_input_text_256(sender_address: str, contract_address: str, function_selector: str, ct, key): + """ + Sign input text for 256-bit encrypted values. + + Similar to sign_input_text but accepts 64-byte ciphertext (CtUint256 blob). + + Args: + sender_address: Address of the sender (20 bytes, without 0x prefix) + contract_address: Address of the contract (20 bytes, without 0x prefix) + function_selector: Function selector (hex string with 0x prefix, e.g., '0x12345678') + ct: Ciphertext bytes (must be 64 bytes for uint256) + key: Signing key (32 bytes) + + Returns: + bytes: The signature + + Raises: + ValueError: If any input has invalid length + """ + function_selector_bytes = bytes.fromhex(function_selector[2:]) + + if len(sender_address) != address_size: + raise ValueError(f"Invalid sender address length: {len(sender_address)} bytes, must be {address_size} bytes") + if len(contract_address) != address_size: + raise ValueError(f"Invalid contract address length: {len(contract_address)} bytes, must be {address_size} bytes") + if len(function_selector_bytes) != function_selector_size: + raise ValueError(f"Invalid signature size: {len(function_selector_bytes)} bytes, must be {function_selector_size} bytes") + + # 256-bit IT has 64 bytes CT + if len(ct) != 64: + raise ValueError(f"Invalid ct length: {len(ct)} bytes, must be 64 bytes for uint256") + + if len(key) != key_size: + raise ValueError(f"Invalid key length: {len(key)} bytes, must be {key_size} bytes") + + message = sender_address + contract_address + function_selector_bytes + ct + + return sign(message, key) + + def sign(message, key): # Sign the message pk = keys.PrivateKey(key) @@ -122,7 +162,7 @@ def build_input_text(plaintext: int, user_aes_key: str, sender_address: str, con } -def build_string_input_text(plaintext: int, user_aes_key: str, sender_address: str, contract_address: str, function_selector: str, signing_key: str) -> ItString: +def build_string_input_text(plaintext: str, user_aes_key: str, sender_address: str, contract_address: str, function_selector: str, signing_key: str) -> ItString: input_text = { 'ciphertext': { 'value': [] @@ -169,6 +209,130 @@ def decrypt_uint(ciphertext: CtUint, user_aes_key: str) -> int: return decrypted_uint +def create_ciphertext_256(plaintext_int: int, user_aes_key_bytes: bytes) -> bytes: + """ + Create a 256-bit ciphertext by encrypting high and low 128-bit parts separately. + + Args: + plaintext_int: Integer value to encrypt (must fit in 256 bits) + user_aes_key_bytes: AES encryption key (16 bytes) + + Returns: + bytes: 64-byte ciphertext blob formatted as: + [high_ciphertext(16) | high_r(16) | low_ciphertext(16) | low_r(16)] + + Raises: + ValueError: If plaintext exceeds 256 bits + """ + # Convert 256-bit int to 32 bytes (Big Endian) + try: + plaintext_bytes = plaintext_int.to_bytes(32, 'big') + except OverflowError: + raise ValueError("Plaintext size must be 256 bits or smaller.") + + # Split into High and Low 128-bit parts + high_bytes = plaintext_bytes[:16] + low_bytes = plaintext_bytes[16:] + + # Encrypt High + high_ct, high_r = encrypt(user_aes_key_bytes, high_bytes) + + # Encrypt Low + low_ct, low_r = encrypt(user_aes_key_bytes, low_bytes) + + # Construct format: high.ciphertext + high.r + low.ciphertext + low.r + return high_ct + high_r + low_ct + low_r + + +def prepare_it_256(plaintext: int, user_aes_key: str, sender_address: str, contract_address: str, function_selector: str, signing_key: str) -> ItUint256: + """ + Prepare Input Text for a 256-bit encrypted integer. + + Encrypts a 256-bit integer and creates a signed Input Text structure + suitable for smart contract interaction. + + Args: + plaintext: Integer value to encrypt (must fit in 256 bits) + user_aes_key: AES encryption key (hex string without 0x prefix, 32 hex chars) + sender_address: Address of the sender (hex string with 0x prefix, 40 hex chars) + contract_address: Address of the contract (hex string with 0x prefix, 40 hex chars) + function_selector: Function selector (hex string with 0x prefix, e.g., '0x12345678') + signing_key: Private key for signing (32 bytes) + + Returns: + ItUint256: Dictionary containing ciphertext (ciphertextHigh, ciphertextLow) and signature + + Raises: + ValueError: If plaintext exceeds 256 bits + """ + if plaintext.bit_length() > 256: + raise ValueError("Plaintext size must be 256 bits or smaller.") + + user_aes_key_bytes = bytes.fromhex(user_aes_key) + + ct_blob = create_ciphertext_256(plaintext, user_aes_key_bytes) + + # Split for types + # ct_blob is 64 bytes: [high_ct(16) | high_r(16) | low_ct(16) | low_r(16)] + high_blob = ct_blob[:32] + low_blob = ct_blob[32:] + + # Sign the full 64-byte blob + signature = sign_input_text_256( + bytes.fromhex(sender_address[2:]), + bytes.fromhex(contract_address[2:]), + function_selector, + ct_blob, + signing_key + ) + + return { + 'ciphertext': { + 'ciphertextHigh': int.from_bytes(high_blob, 'big'), + 'ciphertextLow': int.from_bytes(low_blob, 'big') + }, + 'signature': signature + } + + +def decrypt_uint256(ciphertext: CtUint256, user_aes_key: str) -> int: + """ + Decrypt a 256-bit encrypted integer. + + Decrypts both high and low 128-bit parts and combines them back into + a single 256-bit integer. + + Args: + ciphertext: CtUint256 dictionary containing ciphertextHigh and ciphertextLow + user_aes_key: AES decryption key (hex string without 0x prefix, 32 hex chars) + + Returns: + int: The decrypted 256-bit integer value + """ + user_aes_key_bytes = bytes.fromhex(user_aes_key) + + # Process High + ct_high_int = ciphertext['ciphertextHigh'] + ct_high_bytes = ct_high_int.to_bytes(32, 'big') + cipher_high = ct_high_bytes[:block_size] + r_high = ct_high_bytes[block_size:] + + plaintext_high = decrypt(user_aes_key_bytes, r_high, cipher_high) + + # Process Low + ct_low_int = ciphertext['ciphertextLow'] + ct_low_bytes = ct_low_int.to_bytes(32, 'big') + cipher_low = ct_low_bytes[:block_size] + r_low = ct_low_bytes[block_size:] + + plaintext_low = decrypt(user_aes_key_bytes, r_low, cipher_low) + + # Combine back to 256-bit int + # High part is MSB + full_bytes = plaintext_high + plaintext_low + return int.from_bytes(full_bytes, 'big') + + def decrypt_string(ciphertext: CtString, user_aes_key: str) -> str: if 'value' in ciphertext or hasattr(ciphertext, 'value'): # format when reading ciphertext from an event __ciphertext = ciphertext['value'] @@ -231,6 +395,7 @@ def decrypt_rsa(private_key_bytes: bytes, ciphertext: bytes): ) return plaintext + #This function recovers a user's key by decrypting two encrypted key shares with the given private key, #and then XORing the two key shares together. def recover_user_key(private_key_bytes: bytes, encrypted_key_share0: bytes, encrypted_key_share1: bytes): @@ -239,3 +404,4 @@ def recover_user_key(private_key_bytes: bytes, encrypted_key_share0: bytes, encr # XOR both key shares to get the user key return bytes([a ^ b for a, b in zip(key_share0, key_share1)]) + diff --git a/coti/types.py b/coti/types.py index 487c5eb..657e897 100644 --- a/coti/types.py +++ b/coti/types.py @@ -19,4 +19,19 @@ class ItStringCiphertext(TypedDict): class ItString(TypedDict): ciphertext: ItStringCiphertext - signature: List[bytes] \ No newline at end of file + signature: List[bytes] + + +class CtUint256(TypedDict): + ciphertextHigh: int + ciphertextLow: int + + +class ItUint256Ciphertext(TypedDict): + ciphertextHigh: int + ciphertextLow: int + + +class ItUint256(TypedDict): + ciphertext: ItUint256Ciphertext + signature: bytes diff --git a/tests/__pycache__/conftest.cpython-314-pytest-9.0.2.pyc b/tests/__pycache__/conftest.cpython-314-pytest-9.0.2.pyc new file mode 100644 index 0000000..07cb215 Binary files /dev/null and b/tests/__pycache__/conftest.cpython-314-pytest-9.0.2.pyc differ diff --git a/tests/__pycache__/test_crypto_utils.cpython-314-pytest-9.0.2.pyc b/tests/__pycache__/test_crypto_utils.cpython-314-pytest-9.0.2.pyc new file mode 100644 index 0000000..e219eb1 Binary files /dev/null and b/tests/__pycache__/test_crypto_utils.cpython-314-pytest-9.0.2.pyc differ diff --git a/tests/__pycache__/test_crypto_utils_256.cpython-314-pytest-9.0.2.pyc b/tests/__pycache__/test_crypto_utils_256.cpython-314-pytest-9.0.2.pyc new file mode 100644 index 0000000..da6d98f Binary files /dev/null and b/tests/__pycache__/test_crypto_utils_256.cpython-314-pytest-9.0.2.pyc differ diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..e36375a --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,34 @@ +import pytest +import os +from eth_keys import keys +from coti.crypto_utils import generate_aes_key + +@pytest.fixture +def user_key(): + # Return hex string of the key + return os.environ.get("TEST_USER_KEY") or generate_aes_key().hex() + +@pytest.fixture +def private_key_bytes(): + pk_hex = os.environ.get("TEST_PRIVATE_KEY") + if pk_hex: + # Handle 0x prefix if present + clean_hex = pk_hex[2:] if pk_hex.startswith("0x") else pk_hex + return bytes.fromhex(clean_hex) + return os.urandom(32) + +@pytest.fixture +def sender_address(private_key_bytes): + # Derive address from private key to ensure consistency in signing checks if needed + pk = keys.PrivateKey(private_key_bytes) + return pk.public_key.to_checksum_address() + +@pytest.fixture +def contract_address(): + # Dummy contract address + return "0x1000000000000000000000000000000000000001" + +@pytest.fixture +def function_selector(): + # Dummy function selector + return "0x11223344" diff --git a/tests/test_crypto_utils.py b/tests/test_crypto_utils.py new file mode 100644 index 0000000..741fa77 --- /dev/null +++ b/tests/test_crypto_utils.py @@ -0,0 +1,134 @@ +from coti.crypto_utils import ( + encrypt, decrypt, build_input_text, build_string_input_text, + decrypt_uint, decrypt_string, block_size +) +from Crypto.Random import get_random_bytes +import pytest + +def test_encrypt_decrypt_basic(user_key): + # Test basic encryption and decryption + user_key_bytes = bytes.fromhex(user_key) + plaintext = b"0" * 15 + b"1" # 16 bytes + + ciphertext, r = encrypt(user_key_bytes, plaintext) + + assert len(ciphertext) == block_size + assert len(r) == block_size + + decrypted = decrypt(user_key_bytes, r, ciphertext) + assert decrypted == plaintext + +def test_encrypt_invalid_plaintext_size(user_key): + user_key_bytes = bytes.fromhex(user_key) + plaintext = b"0" * 17 # 17 bytes + + with pytest.raises(ValueError, match="Plaintext size must be 128 bits or smaller"): + encrypt(user_key_bytes, plaintext) + +def test_encrypt_invalid_key_size(): + user_key_bytes = get_random_bytes(15) # 15 bytes + plaintext = b"0" * 16 + + with pytest.raises(ValueError, match="Key size must be 128 bits"): + encrypt(user_key_bytes, plaintext) + +def test_decrypt_invalid_ciphertext_size(user_key): + user_key_bytes = bytes.fromhex(user_key) + r = get_random_bytes(16) + ciphertext = b"0" * 15 + + with pytest.raises(ValueError, match="Ciphertext size must be 128 bits"): + decrypt(user_key_bytes, r, ciphertext) + +def test_build_input_text_128(user_key, sender_address, contract_address, function_selector, private_key_bytes): + # Test 128-bit UINT Input Text + plaintext = 123456789 + + it = build_input_text( + plaintext, + user_key, + sender_address, + contract_address, + function_selector, + private_key_bytes + ) + + assert 'ciphertext' in it + assert 'signature' in it + assert isinstance(it['ciphertext'], int) + + # Decrypt to verify + decrypted = decrypt_uint(it['ciphertext'], user_key) + assert decrypted == plaintext + +def test_build_input_text_max_128(user_key, sender_address, contract_address, function_selector, private_key_bytes): + # Test max 128-bit value + plaintext = (1 << 128) - 1 + + it = build_input_text( + plaintext, + user_key, + sender_address, + contract_address, + function_selector, + private_key_bytes + ) + + decrypted = decrypt_uint(it['ciphertext'], user_key) + assert decrypted == plaintext + +def test_build_string_input_text_basic(user_key, sender_address, contract_address, function_selector, private_key_bytes): + plaintext = "Hello, World!" + + it = build_string_input_text( + plaintext, + user_key, + sender_address, + contract_address, + function_selector, + private_key_bytes + ) + + assert 'ciphertext' in it + assert 'value' in it['ciphertext'] + assert isinstance(it['ciphertext']['value'], list) + assert len(it['ciphertext']['value']) > 0 + + # Decrypt + decrypted = decrypt_string(it['ciphertext'], user_key) + assert decrypted == plaintext + +def test_build_string_input_text_empty(user_key, sender_address, contract_address, function_selector, private_key_bytes): + plaintext = "" + it = build_string_input_text( + plaintext, + user_key, + sender_address, + contract_address, + function_selector, + private_key_bytes + ) + + assert len(it['ciphertext']['value']) == 0 + + decrypted = decrypt_string(it['ciphertext'], user_key) + assert decrypted == plaintext + +def test_build_string_input_text_long(user_key, sender_address, contract_address, function_selector, private_key_bytes): + # Long string that spans multiple blocks + plaintext = "A" * 100 + + it = build_string_input_text( + plaintext, + user_key, + sender_address, + contract_address, + function_selector, + private_key_bytes + ) + + # 100 bytes / 8 bytes per chunk = 13 chunks (12 full, 1 partial) + assert len(it['ciphertext']['value']) == 13 + + decrypted = decrypt_string(it['ciphertext'], user_key) + assert decrypted == plaintext diff --git a/tests/test_crypto_utils_256.py b/tests/test_crypto_utils_256.py new file mode 100644 index 0000000..1325bfa --- /dev/null +++ b/tests/test_crypto_utils_256.py @@ -0,0 +1,241 @@ + +from coti.crypto_utils import ( + build_input_text, decrypt_uint, prepare_it_256, decrypt_uint256, + generate_aes_key, sign_input_text_256 +) +import pytest +import os + +# Mock keys usually 32 bytes hex +MOCK_AES_KEY = generate_aes_key().hex() +MOCK_SENDER = "0x" + "1" * 40 +MOCK_CONTRACT = "0x" + "2" * 40 +MOCK_SELECTOR = "0x12345678" +# Using a dummy signing key (private key) - we can use an account from brownie or generate one +# For simple unit testing of crypto logic without full eth keys lib dependency, we might need to mock signature +# But crypto_utils imports 'keys' from 'eth_keys'. +# Let's try to use a dummy 32-byte key for signing +MOCK_SIGNING_KEY = os.urandom(32) + + +def test_encryption_decryption_256(): + # Test with a value > 128 bits + plaintext_256 = (1 << 240) + 123456789 + + # Prepare IT (Encrypt) + # We use a mocked signing key, actual signature validity doesn't matter for decrypt unit test + it = prepare_it_256( + plaintext_256, + MOCK_AES_KEY, + MOCK_SENDER, + MOCK_CONTRACT, + MOCK_SELECTOR, + MOCK_SIGNING_KEY + ) + + # Check structure + assert 'ciphertext' in it + assert 'ciphertextHigh' in it['ciphertext'] + assert 'ciphertextLow' in it['ciphertext'] + assert 'signature' in it + + # Decrypt + decrypted_val = decrypt_uint256(it['ciphertext'], MOCK_AES_KEY) + + assert decrypted_val == plaintext_256 + + +def test_encryption_decryption_256_small_value(): + # Test with a small value fitting in 128 bits, but using 256 pipeline + plaintext_small = 42 + + it = prepare_it_256( + plaintext_small, + MOCK_AES_KEY, + MOCK_SENDER, + MOCK_CONTRACT, + MOCK_SELECTOR, + MOCK_SIGNING_KEY + ) + + decrypted_val = decrypt_uint256(it['ciphertext'], MOCK_AES_KEY) + + assert decrypted_val == plaintext_small + + +def test_overflow_check(): + # Test value > 256 bits + plaintext_large = 1 << 257 + + with pytest.raises(ValueError): + prepare_it_256( + plaintext_large, + MOCK_AES_KEY, + MOCK_SENDER, + MOCK_CONTRACT, + MOCK_SELECTOR, + MOCK_SIGNING_KEY + ) + + +def test_max_256_bit_value(): + """Test with exactly 256-bit value (2^256 - 1)""" + plaintext_max = (1 << 256) - 1 + + it = prepare_it_256( + plaintext_max, + MOCK_AES_KEY, + MOCK_SENDER, + MOCK_CONTRACT, + MOCK_SELECTOR, + MOCK_SIGNING_KEY + ) + + decrypted_val = decrypt_uint256(it['ciphertext'], MOCK_AES_KEY) + + assert decrypted_val == plaintext_max + + +def test_various_bit_lengths(): + """Test round-trip with various bit lengths (129-bit, 200-bit, 255-bit)""" + test_values = [ + (1 << 128) + 1, # 129-bit value + (1 << 199) + 12345, # 200-bit value + (1 << 254) + 9876543210, # 255-bit value + ] + + for plaintext in test_values: + it = prepare_it_256( + plaintext, + MOCK_AES_KEY, + MOCK_SENDER, + MOCK_CONTRACT, + MOCK_SELECTOR, + MOCK_SIGNING_KEY + ) + + decrypted_val = decrypt_uint256(it['ciphertext'], MOCK_AES_KEY) + assert decrypted_val == plaintext, f"Failed for {plaintext.bit_length()}-bit value" + + +def test_sign_input_text_256_invalid_sender_address(): + """Test sign_input_text_256 with invalid sender address length""" + ct_blob = b'0' * 64 + + # Invalid sender address (19 bytes instead of 20) + invalid_sender = b'1' * 19 + valid_contract = b'2' * 20 + + with pytest.raises(ValueError, match="Invalid sender address length"): + sign_input_text_256( + invalid_sender, + valid_contract, + MOCK_SELECTOR, + ct_blob, + MOCK_SIGNING_KEY + ) + + +def test_sign_input_text_256_invalid_contract_address(): + """Test sign_input_text_256 with invalid contract address length""" + ct_blob = b'0' * 64 + + valid_sender = b'1' * 20 + invalid_contract = b'2' * 21 # 21 bytes instead of 20 + + with pytest.raises(ValueError, match="Invalid contract address length"): + sign_input_text_256( + valid_sender, + invalid_contract, + MOCK_SELECTOR, + ct_blob, + MOCK_SIGNING_KEY + ) + + +def test_sign_input_text_256_invalid_ct_length(): + """Test sign_input_text_256 with invalid ciphertext length""" + ct_blob = b'0' * 32 # 32 bytes instead of 64 + + valid_sender = b'1' * 20 + valid_contract = b'2' * 20 + + with pytest.raises(ValueError, match="Invalid ct length.*must be 64 bytes"): + sign_input_text_256( + valid_sender, + valid_contract, + MOCK_SELECTOR, + ct_blob, + MOCK_SIGNING_KEY + ) + + +def test_sign_input_text_256_invalid_key_length(): + """Test sign_input_text_256 with invalid signing key length""" + ct_blob = b'0' * 64 + + valid_sender = b'1' * 20 + valid_contract = b'2' * 20 + invalid_key = os.urandom(31) # 31 bytes instead of 32 + + with pytest.raises(ValueError, match="Invalid key length"): + sign_input_text_256( + valid_sender, + valid_contract, + MOCK_SELECTOR, + ct_blob, + invalid_key + ) + + +def test_boundary_values(): + """Test boundary values around 128-bit threshold""" + boundary_values = [ + (1 << 128) - 1, # Max 128-bit value + 1 << 128, # Min 129-bit value (boundary) + 0, # Zero + 1, # One + ] + + for plaintext in boundary_values: + it = prepare_it_256( + plaintext, + MOCK_AES_KEY, + MOCK_SENDER, + MOCK_CONTRACT, + MOCK_SELECTOR, + MOCK_SIGNING_KEY + ) + + decrypted_val = decrypt_uint256(it['ciphertext'], MOCK_AES_KEY) + assert decrypted_val == plaintext + + +def test_high_low_split(): + """Test that high and low parts are correctly split and combined""" + # Value where high part is non-zero and low part is non-zero + high_value = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF # Max 128-bit value + low_value = 0x123456789ABCDEF0123456789ABCDEF0 # Another 128-bit value + + # Combine: high_value << 128 | low_value + plaintext = (high_value << 128) | low_value + + it = prepare_it_256( + plaintext, + MOCK_AES_KEY, + MOCK_SENDER, + MOCK_CONTRACT, + MOCK_SELECTOR, + MOCK_SIGNING_KEY + ) + + decrypted_val = decrypt_uint256(it['ciphertext'], MOCK_AES_KEY) + assert decrypted_val == plaintext + + # Verify the split is correct by checking bit patterns + decrypted_high = decrypted_val >> 128 + decrypted_low = decrypted_val & ((1 << 128) - 1) + + assert decrypted_high == high_value + assert decrypted_low == low_value + diff --git a/tests/test_integration_256.py b/tests/test_integration_256.py new file mode 100644 index 0000000..b3d019b --- /dev/null +++ b/tests/test_integration_256.py @@ -0,0 +1,202 @@ +""" +Integration tests for 256-bit encryption compatibility. + +These tests verify that the Python SDK's 256-bit implementation is compatible +with the TypeScript SDK by using known test vectors and edge cases. +""" + +import pytest +from coti.crypto_utils import prepare_it_256, decrypt_uint256, generate_aes_key, create_ciphertext_256 +import os + + +# Test vectors for cross-SDK compatibility +# These should match the TypeScript SDK's expected behavior +class TestCrossSdkCompatibility: + """Tests to ensure compatibility with TypeScript SDK""" + + def test_ciphertext_structure(self): + """Verify ciphertext structure matches TypeScript SDK format""" + plaintext = 123456789012345678901234567890 + aes_key = generate_aes_key().hex() + sender = "0x" + "1" * 40 + contract = "0x" + "2" * 40 + selector = "0x12345678" + signing_key = os.urandom(32) + + it = prepare_it_256(plaintext, aes_key, sender, contract, selector, signing_key) + + # Verify structure matches expected TypeScript format + assert 'ciphertext' in it + assert 'signature' in it + assert 'ciphertextHigh' in it['ciphertext'] + assert 'ciphertextLow' in it['ciphertext'] + + # Verify both parts are integers + assert isinstance(it['ciphertext']['ciphertextHigh'], int) + assert isinstance(it['ciphertext']['ciphertextLow'], int) + assert isinstance(it['signature'], bytes) + + def test_ciphertext_blob_size(self): + """Verify the internal ciphertext blob is 64 bytes as expected""" + plaintext = (1 << 200) + 42 + aes_key_bytes = generate_aes_key() + + ct_blob = create_ciphertext_256(plaintext, aes_key_bytes) + + # Should be 64 bytes: high_ct(16) + high_r(16) + low_ct(16) + low_r(16) + assert len(ct_blob) == 64 + + # Verify it can be split into 4 equal parts + assert len(ct_blob[:16]) == 16 # high_ct + assert len(ct_blob[16:32]) == 16 # high_r + assert len(ct_blob[32:48]) == 16 # low_ct + assert len(ct_blob[48:64]) == 16 # low_r + + def test_deterministic_encryption_with_same_plaintext(self): + """Verify that encrypting the same plaintext produces different ciphertexts (due to random 'r')""" + plaintext = 987654321 + aes_key = generate_aes_key().hex() + sender = "0x" + "a" * 40 + contract = "0x" + "b" * 40 + selector = "0x11223344" + signing_key = os.urandom(32) + + it1 = prepare_it_256(plaintext, aes_key, sender, contract, selector, signing_key) + it2 = prepare_it_256(plaintext, aes_key, sender, contract, selector, signing_key) + + # Ciphertexts should be different due to random 'r' + assert (it1['ciphertext']['ciphertextHigh'] != it2['ciphertext']['ciphertextHigh'] or + it1['ciphertext']['ciphertextLow'] != it2['ciphertext']['ciphertextLow']) + + # But both should decrypt to the same plaintext + decrypted1 = decrypt_uint256(it1['ciphertext'], aes_key) + decrypted2 = decrypt_uint256(it2['ciphertext'], aes_key) + + assert decrypted1 == plaintext + assert decrypted2 == plaintext + + def test_known_test_vectors(self): + """ + Test with known values that could be shared between SDKs. + + Note: In a real integration test, these would be test vectors + generated by the TypeScript SDK and verified here. + """ + test_cases = [ + 0, + 1, + 255, + 256, + 65535, + 65536, + (1 << 128) - 1, # Max 128-bit + 1 << 128, # Min 129-bit + (1 << 256) - 1, # Max 256-bit + ] + + aes_key = generate_aes_key().hex() + + for plaintext in test_cases: + sender = "0x" + "1" * 40 + contract = "0x" + "2" * 40 + selector = "0x12345678" + signing_key = os.urandom(32) + + it = prepare_it_256(plaintext, aes_key, sender, contract, selector, signing_key) + decrypted = decrypt_uint256(it['ciphertext'], aes_key) + + assert decrypted == plaintext, f"Failed for test case: {plaintext}" + + def test_high_part_only_nonzero(self): + """Test value where only high 128 bits are non-zero""" + # Value = 0xFF...FF << 128 (high part all ones, low part all zeros) + plaintext = ((1 << 128) - 1) << 128 + + aes_key = generate_aes_key().hex() + sender = "0x" + "1" * 40 + contract = "0x" + "2" * 40 + selector = "0x12345678" + signing_key = os.urandom(32) + + it = prepare_it_256(plaintext, aes_key, sender, contract, selector, signing_key) + decrypted = decrypt_uint256(it['ciphertext'], aes_key) + + assert decrypted == plaintext + + def test_low_part_only_nonzero(self): + """Test value where only low 128 bits are non-zero""" + # Value = 0xFF...FF (low part all ones, high part all zeros) + plaintext = (1 << 128) - 1 + + aes_key = generate_aes_key().hex() + sender = "0x" + "1" * 40 + contract = "0x" + "2" * 40 + selector = "0x12345678" + signing_key = os.urandom(32) + + it = prepare_it_256(plaintext, aes_key, sender, contract, selector, signing_key) + decrypted = decrypt_uint256(it['ciphertext'], aes_key) + + assert decrypted == plaintext + + def test_alternating_bit_pattern(self): + """Test with alternating bit patterns""" + + high = 0xAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA + low = 0x55555555555555555555555555555555 + plaintext = (high << 128) | low + + aes_key = generate_aes_key().hex() + sender = "0x" + "1" * 40 + contract = "0x" + "2" * 40 + selector = "0x12345678" + signing_key = os.urandom(32) + + it = prepare_it_256(plaintext, aes_key, sender, contract, selector, signing_key) + decrypted = decrypt_uint256(it['ciphertext'], aes_key) + + assert decrypted == plaintext + + +class TestMultipleKeyScenarios: + """Test encryption/decryption with different AES keys""" + + def test_different_keys_produce_different_results(self): + """Verify that different AES keys produce different ciphertexts""" + plaintext = 123456789 + key1 = generate_aes_key().hex() + key2 = generate_aes_key().hex() + + sender = "0x" + "1" * 40 + contract = "0x" + "2" * 40 + selector = "0x12345678" + signing_key = os.urandom(32) + + it1 = prepare_it_256(plaintext, key1, sender, contract, selector, signing_key) + it2 = prepare_it_256(plaintext, key2, sender, contract, selector, signing_key) + + # Different keys should produce different ciphertexts + assert (it1['ciphertext']['ciphertextHigh'] != it2['ciphertext']['ciphertextHigh'] or + it1['ciphertext']['ciphertextLow'] != it2['ciphertext']['ciphertextLow']) + + def test_wrong_key_produces_wrong_plaintext(self): + """Verify that decrypting with wrong key doesn't recover original plaintext""" + plaintext = 987654321 + correct_key = generate_aes_key().hex() + wrong_key = generate_aes_key().hex() + + sender = "0x" + "1" * 40 + contract = "0x" + "2" * 40 + selector = "0x12345678" + signing_key = os.urandom(32) + + it = prepare_it_256(plaintext, correct_key, sender, contract, selector, signing_key) + + # Decrypt with correct key + decrypted_correct = decrypt_uint256(it['ciphertext'], correct_key) + assert decrypted_correct == plaintext + + # Decrypt with wrong key + decrypted_wrong = decrypt_uint256(it['ciphertext'], wrong_key) + assert decrypted_wrong != plaintext