From c2b86beb4cab9f7300a10f48f96ce5eec865dcdf Mon Sep 17 00:00:00 2001 From: Chris Cassano <1285652+glitch003@users.noreply.github.com> Date: Thu, 9 Jan 2025 09:24:22 +0700 Subject: [PATCH] python test works (#27) --- python/lit_python_sdk/client.py | 82 +++++++++++++++++++++++++++++++++ python/tests/test_client.py | 60 +++++++++++++++++++++++- 2 files changed, 141 insertions(+), 1 deletion(-) diff --git a/python/lit_python_sdk/client.py b/python/lit_python_sdk/client.py index 85f20f8..40d63e6 100644 --- a/python/lit_python_sdk/client.py +++ b/python/lit_python_sdk/client.py @@ -179,6 +179,88 @@ def generate_auth_sig(self, to_sign: str) -> Dict[str, Any]: """Generates an auth signature""" return self._post("/authHelpers/generateAuthSig", {"toSign": to_sign}) + def encrypt_string( + self, + data_to_encrypt: str, + access_control_conditions: Optional[List[Dict[str, Any]]] = None, + evm_contract_conditions: Optional[List[Dict[str, Any]]] = None, + sol_rpc_conditions: Optional[List[Dict[str, Any]]] = None, + unified_access_control_conditions: Optional[List[Dict[str, Any]]] = None, + ) -> Dict[str, Any]: + """Encrypts a string using Lit Protocol + + Args: + data_to_encrypt: The string to encrypt + access_control_conditions: Optional list of access control conditions + evm_contract_conditions: Optional list of EVM contract conditions + sol_rpc_conditions: Optional list of Solana RPC conditions + unified_access_control_conditions: Optional list of unified access control conditions + + Returns: + Dict containing the encrypted data (ciphertext) and dataToEncryptHash + """ + payload = { + "dataToEncrypt": data_to_encrypt, + } + if access_control_conditions is not None: + payload["accessControlConditions"] = access_control_conditions + if evm_contract_conditions is not None: + payload["evmContractConditions"] = evm_contract_conditions + if sol_rpc_conditions is not None: + payload["solRpcConditions"] = sol_rpc_conditions + if unified_access_control_conditions is not None: + payload["unifiedAccessControlConditions"] = unified_access_control_conditions + + return self._post("/litNodeClient/encryptString", payload) + + def decrypt_string( + self, + ciphertext: str, + data_to_encrypt_hash: str, + chain: str, + access_control_conditions: Optional[List[Dict[str, Any]]] = None, + evm_contract_conditions: Optional[List[Dict[str, Any]]] = None, + sol_rpc_conditions: Optional[List[Dict[str, Any]]] = None, + unified_access_control_conditions: Optional[List[Dict[str, Any]]] = None, + auth_sig: Optional[Dict[str, Any]] = None, + session_sigs: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Decrypts a string using Lit Protocol + + Args: + ciphertext: The encrypted string to decrypt + data_to_encrypt_hash: The hash of the original data + chain: The blockchain network (e.g. 'ethereum') + access_control_conditions: Optional list of access control conditions + evm_contract_conditions: Optional list of EVM contract conditions + sol_rpc_conditions: Optional list of Solana RPC conditions + unified_access_control_conditions: Optional list of unified access control conditions + auth_sig: Optional authentication signature + session_sigs: Optional session signatures + + Returns: + Dict containing the decrypted string + """ + payload = { + "ciphertext": ciphertext, + "dataToEncryptHash": data_to_encrypt_hash, + "chain": chain, + } + if access_control_conditions is not None: + payload["accessControlConditions"] = access_control_conditions + if evm_contract_conditions is not None: + payload["evmContractConditions"] = evm_contract_conditions + if sol_rpc_conditions is not None: + payload["solRpcConditions"] = sol_rpc_conditions + if unified_access_control_conditions is not None: + payload["unifiedAccessControlConditions"] = unified_access_control_conditions + if auth_sig is not None: + payload["authSig"] = auth_sig + if session_sigs is not None: + payload["sessionSigs"] = session_sigs + + return self._post("/litNodeClient/decryptString", payload) + def __del__(self): """Cleanup: Stop the Node.js server when the client is destroyed""" if hasattr(self, 'server') and self.server is not None: diff --git a/python/tests/test_client.py b/python/tests/test_client.py index 40af9df..e514474 100644 --- a/python/tests/test_client.py +++ b/python/tests/test_client.py @@ -148,4 +148,62 @@ def test_contracts_and_auth(): to_sign=to_sign, session_sigs=session_sigs ) - assert "signature" in sign_result, "Expected signature in response" \ No newline at end of file + assert "signature" in sign_result, "Expected signature in response" + +def test_encrypt_decrypt_string(): + # First get session sigs + expiration = (datetime.now(timezone.utc) + timedelta(minutes=10)).strftime("%Y-%m-%dT%H:%M:%SZ") + session_sigs_result = client.get_session_sigs( + chain="ethereum", + expiration=expiration, + resource_ability_requests=[{ + "resource": { + "resource": "*", + "resourcePrefix": "lit-litaction", + }, + "ability": "lit-action-execution", + }, { + "resource": { + "resource": "*", + "resourcePrefix": "lit-pkp", + }, + "ability": "pkp-signing", + }] + ) + assert "sessionSigs" in session_sigs_result, "Expected sessionSigs in response" + session_sigs = session_sigs_result["sessionSigs"] + + # Test string to encrypt + test_string = "Hello, World!" + + # Set up access control conditions + access_control_conditions = [{ + "contractAddress": "", + "standardContractType": "", + "chain": "ethereum", + "method": "", + "parameters": [":userAddress"], + "returnValueTest": { + "comparator": "=", + "value": wallet_address + } + }] + + # Test encryption + encrypt_result = client.encrypt_string( + data_to_encrypt=test_string, + access_control_conditions=access_control_conditions + ) + assert "ciphertext" in encrypt_result, "Expected ciphertext in response" + assert "dataToEncryptHash" in encrypt_result, "Expected dataToEncryptHash in response" + + # Test decryption + decrypt_result = client.decrypt_string( + ciphertext=encrypt_result["ciphertext"], + data_to_encrypt_hash=encrypt_result["dataToEncryptHash"], + chain="ethereum", + access_control_conditions=access_control_conditions, + session_sigs=session_sigs + ) + assert "decryptedString" in decrypt_result, "Expected decryptedString in response" + assert decrypt_result["decryptedString"] == test_string, "Decrypted string should match original" \ No newline at end of file