Skip to content

Commit dab5983

Browse files
dineshpintodineshpinto
authored
Merge branch 'main' into dependabot/github_actions/actions-all-e1e20d6164
2 parents 71c6b7a + 424df15 commit dab5983

File tree

16 files changed

+3962
-4
lines changed

16 files changed

+3962
-4
lines changed

‎examples/03_turnkey_wallet_integration.py‎

Lines changed: 386 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,386 @@
1+
"""
2+
Turnkey Wallet Integration Example.
3+
4+
This example demonstrates how to use the Turnkey wallet integration
5+
for AI agents to securely manage blockchain transactions on Flare.
6+
7+
Features demonstrated:
8+
- Wallet creation and management
9+
- AI agent registration with policies
10+
- Transaction signing with permission enforcement
11+
- TEE attestation validation
12+
- Secure audit logging
13+
"""
14+
15+
import asyncio
16+
import os
17+
from decimal import Decimal
18+
from typing import Any
19+
20+
from flare_ai_kit.agents.turnkey import (
21+
AgentTransaction,
22+
AgentWalletConfig,
23+
TurnkeyAgentConnector,
24+
)
25+
from flare_ai_kit.tee.validation import VtpmValidation
26+
from flare_ai_kit.wallet import PermissionEngine, TransactionPolicy, TurnkeyWallet
27+
from flare_ai_kit.wallet.base import TransactionRequest
28+
from flare_ai_kit.wallet.permissions import TimeWindow
29+
from flare_ai_kit.wallet.tee_security import TEESecurityManager
30+
from flare_ai_kit.wallet.turnkey_wallet import TurnkeySettings
31+
32+
33+
async def setup_turnkey_wallet() -> TurnkeyWallet:
34+
"""Set up Turnkey wallet with proper configuration."""
35+
print("🔧 Setting up Turnkey wallet...")
36+
37+
# Configure Turnkey settings (in production, use environment variables)
38+
settings = TurnkeySettings(
39+
organization_id=os.getenv("TURNKEY_ORGANIZATION_ID", "demo_org_id"),
40+
api_public_key=os.getenv("TURNKEY_API_PUBLIC_KEY", "demo_public_key"),
41+
api_private_key=os.getenv("TURNKEY_API_PRIVATE_KEY", "demo_private_key"),
42+
)
43+
44+
# Create permission engine with default policies
45+
permission_engine = PermissionEngine()
46+
47+
# Add a restrictive default policy for security
48+
default_policy = TransactionPolicy(
49+
name="default_security_policy",
50+
description="Default restrictive policy for all transactions",
51+
max_transaction_value=Decimal("0.01"), # 0.01 ETH max per transaction
52+
daily_spending_limit=Decimal("0.1"), # 0.1 ETH max per day
53+
time_windows=[
54+
TimeWindow(
55+
duration_minutes=60, max_transactions=10, max_value=Decimal("0.05")
56+
),
57+
TimeWindow(
58+
duration_minutes=1440, max_transactions=100, max_value=Decimal("0.1")
59+
),
60+
],
61+
allowed_hours_utc=list(range(6, 22)), # 6 AM to 10 PM UTC only
62+
max_gas_price="50000000000", # 50 gwei max
63+
max_gas_limit="200000",
64+
)
65+
66+
permission_engine.add_policy(default_policy)
67+
68+
# Create TEE validator for secure operations
69+
tee_validator = VtpmValidation()
70+
71+
# Initialize wallet
72+
wallet = TurnkeyWallet(
73+
settings=settings,
74+
permission_engine=permission_engine,
75+
tee_validator=tee_validator,
76+
)
77+
78+
print("✅ Turnkey wallet configured successfully")
79+
return wallet
80+
81+
82+
async def create_demo_wallet(wallet: TurnkeyWallet) -> tuple[str, str]:
83+
"""Create a demo wallet for the AI agent."""
84+
print("\n💼 Creating demo wallet...")
85+
86+
try:
87+
wallet_id = await wallet.create_wallet("AI_Agent_Demo_Wallet")
88+
print(f"✅ Wallet created with ID: {wallet_id}")
89+
90+
# Get the wallet address
91+
address_info = await wallet.get_address(wallet_id)
92+
print(f"📍 Wallet address: {address_info.address}")
93+
print(f"📍 Derivation path: {address_info.derivation_path}")
94+
95+
return wallet_id, address_info.address
96+
97+
except Exception as e:
98+
print(f"❌ Failed to create wallet: {e}")
99+
# For demo purposes, return mock values
100+
return "demo_wallet_id", "0x742d35Cc6634C0532925a3b8D8C8EE7c9e92bb1b"
101+
102+
103+
async def register_ai_agent(
104+
agent_connector: TurnkeyAgentConnector, wallet_id: str
105+
) -> AgentWalletConfig:
106+
"""Register an AI agent with specific permissions."""
107+
print("\n🤖 Registering AI agent...")
108+
109+
# Configure agent with conservative limits
110+
agent_config = AgentWalletConfig(
111+
agent_id="yield_optimizer_v1",
112+
wallet_id=wallet_id,
113+
max_daily_spend=0.05, # 0.05 ETH per day
114+
max_transaction_value=0.01, # 0.01 ETH per transaction
115+
allowed_hours=list(range(9, 17)), # 9 AM to 5 PM UTC only
116+
require_tee_attestation=True,
117+
)
118+
119+
success = await agent_connector.register_agent(agent_config)
120+
121+
if success:
122+
print(f"✅ Agent '{agent_config.agent_id}' registered successfully")
123+
print(f" 📊 Max daily spend: {agent_config.max_daily_spend} ETH")
124+
print(f" 📊 Max transaction value: {agent_config.max_transaction_value} ETH")
125+
print(f" 🕒 Allowed hours: {agent_config.allowed_hours}")
126+
print(f" 🔒 TEE attestation required: {agent_config.require_tee_attestation}")
127+
else:
128+
print("❌ Failed to register AI agent")
129+
130+
return agent_config if success else None
131+
132+
133+
async def simulate_tee_operation(_tee_manager: TEESecurityManager) -> str:
134+
"""Simulate a TEE secure operation."""
135+
print("\n🔐 Simulating TEE secure operation...")
136+
137+
# Mock TEE attestation token (in production, this comes from the TEE)
138+
mock_attestation_token = (
139+
"eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9."
140+
"eyJzdWIiOiJkZW1vX3RlZSIsImlzcyI6Imh0dHBzOi8vY29uZmlkZW50aWFsY29tcHV0aW5nLmdvb2dsZWFwaXMuY29tIiwiaWF0IjoxNjk5OTk5OTk5fQ."
141+
"mock_signature"
142+
)
143+
144+
try:
145+
# This would normally validate against real TEE attestation
146+
print("⚠️ Note: Using mock TEE attestation for demo purposes")
147+
148+
# In production, this would create a secure operation
149+
150+
print("✅ TEE operation simulated (would validate in production)")
151+
return mock_attestation_token
152+
153+
except Exception as e:
154+
print(f"❌ TEE operation failed: {e}")
155+
print(" Using mock attestation for demo purposes")
156+
return mock_attestation_token
157+
158+
159+
async def execute_ai_transaction(
160+
agent_connector: TurnkeyAgentConnector,
161+
agent_config: AgentWalletConfig,
162+
attestation_token: str,
163+
) -> dict[str, Any]:
164+
"""Execute a transaction initiated by the AI agent."""
165+
print("\n💸 Executing AI agent transaction...")
166+
167+
# Create a sample transaction request
168+
transaction_request = TransactionRequest(
169+
to="0x742d35Cc6634C0532925a3b8D8C8EE7c9e92bb1b", # Demo recipient
170+
value="5000000000000000", # 0.005 ETH in wei
171+
gas_limit="21000",
172+
gas_price="20000000000", # 20 gwei
173+
chain_id=14, # Flare mainnet
174+
)
175+
176+
# Create agent transaction with justification and risk assessment
177+
agent_transaction = AgentTransaction(
178+
agent_id=agent_config.agent_id,
179+
transaction_request=transaction_request,
180+
justification=(
181+
"Automated yield optimization - moving funds to higher yield protocol"
182+
),
183+
confidence_score=0.85,
184+
risk_assessment="Low risk - established DeFi protocol with 99.9% uptime",
185+
)
186+
187+
print("📝 Transaction details:")
188+
print(f" To: {transaction_request.to}")
189+
print(f" Value: {float(transaction_request.value) / 10**18} ETH")
190+
print(f" Agent confidence: {agent_transaction.confidence_score}")
191+
print(f" Justification: {agent_transaction.justification}")
192+
193+
try:
194+
# Execute the transaction
195+
result = await agent_connector.execute_agent_transaction(
196+
agent_transaction, attestation_token=attestation_token
197+
)
198+
199+
if result["success"]:
200+
print("✅ Transaction executed successfully!")
201+
print(f" 📋 Transaction hash: {result['transaction_hash']}")
202+
print(f" 📋 Signed transaction: {result['signed_transaction'][:50]}...")
203+
else:
204+
print(f"❌ Transaction failed: {result['error']}")
205+
206+
return result
207+
208+
except Exception as e:
209+
print(f"❌ Transaction execution error: {e}")
210+
return {"success": False, "error": str(e)}
211+
212+
213+
async def demonstrate_policy_enforcement(
214+
agent_connector: TurnkeyAgentConnector,
215+
agent_config: AgentWalletConfig,
216+
attestation_token: str,
217+
) -> None:
218+
"""Demonstrate policy enforcement by attempting violating transactions."""
219+
print("\n🚫 Demonstrating policy enforcement...")
220+
221+
# Test 1: Transaction exceeding value limit
222+
print("\n📊 Test 1: Transaction exceeding value limit")
223+
large_transaction = AgentTransaction(
224+
agent_id=agent_config.agent_id,
225+
transaction_request=TransactionRequest(
226+
to="0x742d35Cc6634C0532925a3b8D8C8EE7c9e92bb1b",
227+
value="20000000000000000", # 0.02 ETH (exceeds 0.01 ETH limit)
228+
chain_id=14,
229+
),
230+
justification="Large yield optimization move",
231+
confidence_score=0.9,
232+
risk_assessment="Medium risk - large transaction",
233+
)
234+
235+
result = await agent_connector.execute_agent_transaction(
236+
large_transaction, attestation_token
237+
)
238+
if not result["success"]:
239+
print(f"✅ Policy correctly blocked large transaction: {result['error']}")
240+
else:
241+
print("❌ Policy should have blocked this transaction")
242+
243+
# Test 2: Low confidence transaction
244+
print("\n📊 Test 2: Low confidence transaction")
245+
low_confidence_tx = AgentTransaction(
246+
agent_id=agent_config.agent_id,
247+
transaction_request=TransactionRequest(
248+
to="0x742d35Cc6634C0532925a3b8D8C8EE7c9e92bb1b",
249+
value="5000000000000000", # 0.005 ETH
250+
chain_id=14,
251+
),
252+
justification="Uncertain yield opportunity",
253+
confidence_score=0.4, # Low confidence
254+
risk_assessment="High risk - unverified protocol",
255+
)
256+
257+
result = await agent_connector.execute_agent_transaction(
258+
low_confidence_tx, attestation_token
259+
)
260+
if not result["success"]:
261+
print(
262+
f"✅ Policy correctly blocked low confidence transaction: {result['error']}"
263+
)
264+
else:
265+
print("❌ Policy should have blocked this transaction")
266+
267+
# Test 3: Missing justification
268+
print("\n📊 Test 3: Missing justification")
269+
no_justification_tx = AgentTransaction(
270+
agent_id=agent_config.agent_id,
271+
transaction_request=TransactionRequest(
272+
to="0x742d35Cc6634C0532925a3b8D8C8EE7c9e92bb1b",
273+
value="5000000000000000",
274+
chain_id=14,
275+
),
276+
justification="", # Empty justification
277+
confidence_score=0.8,
278+
risk_assessment="Low risk",
279+
)
280+
281+
result = await agent_connector.execute_agent_transaction(
282+
no_justification_tx, attestation_token
283+
)
284+
if not result["success"]:
285+
print(
286+
"✅ Policy correctly blocked transaction without justification: "
287+
f"{result['error']}"
288+
)
289+
else:
290+
print("❌ Policy should have blocked this transaction")
291+
292+
293+
async def show_agent_status_and_history(
294+
agent_connector: TurnkeyAgentConnector, agent_id: str
295+
) -> None:
296+
"""Display agent status and transaction history."""
297+
print("\n📊 Agent Status and Transaction History")
298+
299+
# Get agent status
300+
status = await agent_connector.get_agent_status(agent_id)
301+
if status:
302+
print(f"\n🤖 Agent: {status['agent_id']}")
303+
print(f" 💼 Wallet ID: {status['wallet_id']}")
304+
print(f" 📊 Total transactions: {status['statistics']['total_transactions']}")
305+
print(f" 💰 Total value: {status['statistics']['total_value_eth']:.6f} ETH")
306+
avg_confidence = status["statistics"]["avg_confidence_score"]
307+
print(f" 📈 Average confidence: {avg_confidence:.2f}")
308+
recent_24h = status["statistics"]["recent_transactions_24h"]
309+
print(f" 🕒 Recent transactions (24h): {recent_24h}")
310+
311+
# Get transaction history
312+
history = await agent_connector.get_transaction_history(agent_id=agent_id, limit=10)
313+
314+
if history:
315+
print("\n📜 Recent Transaction History:")
316+
for i, tx in enumerate(history[:5], 1):
317+
print(f" {i}. Hash: {tx['transaction_hash']}")
318+
print(f" Value: {float(tx['value']):.6f} ETH")
319+
print(f" Confidence: {tx['confidence_score']}")
320+
print(f" TEE Attested: {tx['tee_attested']}")
321+
else:
322+
print(" No transaction history found")
323+
324+
325+
async def main() -> None:
326+
"""Main demonstration function."""
327+
print("🚀 Turnkey Wallet Integration Demo")
328+
print("=" * 50)
329+
330+
try:
331+
# 1. Set up Turnkey wallet
332+
wallet = await setup_turnkey_wallet()
333+
334+
# 2. Create demo wallet
335+
wallet_id, wallet_address = await create_demo_wallet(wallet)
336+
337+
# 3. Set up TEE security manager
338+
tee_manager = TEESecurityManager()
339+
340+
# 4. Create agent connector
341+
agent_connector = TurnkeyAgentConnector(wallet)
342+
343+
# 5. Register AI agent
344+
agent_config = await register_ai_agent(agent_connector, wallet_id)
345+
346+
if not agent_config:
347+
print("❌ Cannot continue without registered agent")
348+
return
349+
350+
# 6. Simulate TEE operation
351+
attestation_token = await simulate_tee_operation(tee_manager)
352+
353+
# 7. Execute valid AI transaction
354+
await execute_ai_transaction(agent_connector, agent_config, attestation_token)
355+
356+
# 8. Demonstrate policy enforcement
357+
await demonstrate_policy_enforcement(
358+
agent_connector, agent_config, attestation_token
359+
)
360+
361+
# 9. Show agent status and history
362+
await show_agent_status_and_history(agent_connector, agent_config.agent_id)
363+
364+
# 10. List all registered agents
365+
agents = await agent_connector.list_registered_agents()
366+
print(f"\n📋 All registered agents: {agents}")
367+
368+
# 11. Cleanup
369+
print("\n🧹 Cleaning up...")
370+
await agent_connector.unregister_agent(agent_config.agent_id)
371+
print("✅ Agent unregistered")
372+
373+
# Close wallet connection
374+
await wallet.__aexit__(None, None, None)
375+
print("✅ Wallet connection closed")
376+
377+
print("\n🎉 Demo completed successfully!")
378+
379+
except Exception as e:
380+
print(f"\n❌ Demo failed with error: {e}")
381+
raise
382+
383+
384+
if __name__ == "__main__":
385+
# Run the demo
386+
asyncio.run(main())

0 commit comments

Comments
 (0)