feat: Implement 256-bit encryption and expand test coverage

Author: plucena

.gitignore

@@ -3,3 +3,4 @@
 /coti_sdk.egg-info/
 
 /build/
+.venv/

Comparisson.md

@@ -0,0 +1,46 @@
+# SDK Comparison: Python vs TypeScript
+
+This document compares the functionality available in `coti-sdk-python` against `coti-sdk-typescript`.
+
+## Overview
+
+Overall, both SDKs provide core capabilities for:
+- AES encryption/decryption (128-bit)
+- RSA key pair generation and decryption
+- User key recovery
+- Signing and building Input Text (IT) for smart contracts
+
+**Key Difference:** The TypeScript SDK currently includes explicit support for **256-bit integer encryption/decryption** (splitting into high/low 128-bit blocks), whereas the Python SDK primarily targets 128-bit integers (32-byte ciphertext blobs containing cipher + random).
+
+## Function Parity Table
+
+| Feature / Functionality    | Python SDK (`coti/crypto_utils.py`) | TypeScript SDK (`src/crypto_utils.ts`)                  | Notes                                                                                                      |
+| :------------------------- | :---------------------------------- | :------------------------------------------------------ | :--------------------------------------------------------------------------------------------------------- |
+| **AES Encryption**         | `encrypt`                           | `encrypt`                                               | Both use AES-128 (16-byte block/key).                                                                      |
+| **AES Decryption**         | `decrypt`                           | `decrypt`                                               | TS `decrypt` supports optional `r2`/`ciphertext2` args for multi-block. Python matches strict 128-bit sig. |
+| **Generate AES Key**       | `generate_aes_key`                  | `generateRandomAesKeySizeNumber`                        | Naming differs, same functionality (16 random bytes).                                                      |
+| **RSA KeyGen**             | `generate_rsa_keypair`              | `generateRSAKeyPair`                                    |                                                                                                            |
+| **RSA Decryption**         | `decrypt_rsa`                       | `decryptRSA`                                            |                                                                                                            |
+| **Recover User Key**       | `recover_user_key`                  | `recoverUserKey`                                        | Decrypts shares & XORs them.                                                                               |
+| **Sign Message**           | `sign`                              | `sign`                                                  |                                                                                                            |
+| **Sign Input Text**        | `sign_input_text`                   | `signInputText`                                         |                                                                                                            |
+| **Build IT (Uint)**        | `build_input_text` (128-bit)        | `buildInputText` (64-bit), `prepareIT` (128-bit)        | TS has separate methods for 64 and 128 bit.                                                                |
+| **Build IT (String)**      | `build_string_input_text`           | `buildStringInputText`                                  | Handles string chunking into 64-bit blocks.                                                                |
+| **Decrypt Uint (128-bit)** | `decrypt_uint`                      | `decryptUint`                                           | Standard 128-bit decryption (Cipher + R).                                                                  |
+| **Decrypt Uint (256-bit)** | `decrypt_uint256`                   | `decryptUint256`                                        | Splitting into High/Low parts.                                                                             |
+| **Prepare IT (256-bit)**   | `prepare_it_256`                    | `prepareIT256`                                          | Primary entry point for 256-bit plaintexts.                                                                |
+| **Decrypt String**         | `decrypt_string`                    | `decryptString`                                         |                                                                                                            |
+| **Encoding Helpers**       | (Internal/Standard Lib)             | `encodeString`, `encodeKey`, `encodeUint`, `decodeUint` | TS exposes these utils; Python relies on Standard Lib.                                                     |
+| **Encrypt Number**         | (Internal steps)                    | `encryptNumber`                                         | TS helper exposed (used internally in both).                                                               |
+
+## Detailed Differences
+
+### 256-bit Support
+The TypeScript SDK has specific functions (`prepareIT256`, `decryptUint256`) and logic in `decrypt` to handle 256-bit integers by splitting them into two 128-bit AES blocks. The Python SDK currently treats `CtUint` as a single 32-byte blob (16 bytes cipher + 16 bytes random), effectively limiting it to 128-bit integers.
+
+### Naming Conventions
+- Python uses `snake_case` (e.g., `recover_user_key`).
+- TypeScript uses `camelCase` (e.g., `recoverUserKey`).
+
+### Helper Functions
+The TypeScript SDK exports several low-level helper functions (`encode*`, `decode*`) which are useful for JS/TS type coercion. The Python SDK tends to keep these internal or rely on Python's rich standard library methods for bytes/int conversion.

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)])
+

coti/types.py

@@ -19,4 +19,19 @@ class ItStringCiphertext(TypedDict):
 
 class ItString(TypedDict): 
     ciphertext: ItStringCiphertext
-    signature: List[bytes]
+    signature: List[bytes]
+
+
+class CtUint256(TypedDict):
+    ciphertextHigh: int
+    ciphertextLow: int
+
+
+class ItUint256Ciphertext(TypedDict):
+    ciphertextHigh: int
+    ciphertextLow: int
+
+
+class ItUint256(TypedDict):
+    ciphertext: ItUint256Ciphertext
+    signature: bytes

tests/__pycache__/conftest.cpython-314-pytest-9.0.2.pyc

@@ -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"