chore(security): enhance environment configuration, CI workflows, and wallet daemon with security improvements
- Restructure .env.example with security-focused documentation, service-specific environment file references, and AWS Secrets Manager integration - Update CLI tests workflow to single Python 3.13 version, add pytest-mock dependency, and consolidate test execution with coverage - Add comprehensive security validation to package publishing workflow with manual approval gates, secret scanning, and release
This commit is contained in:
oib
@@ -0,0 +1,584 @@
|
||||
"""
|
||||
AITBC Agent Wallet Security Implementation
|
||||
|
||||
This module implements the security layer for autonomous agent wallets,
|
||||
integrating the guardian contract to prevent unlimited spending in case
|
||||
of agent compromise.
|
||||
"""
|
||||
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta
|
||||
import json
|
||||
from eth_account import Account
|
||||
from eth_utils import to_checksum_address
|
||||
|
||||
from .guardian_contract import (
|
||||
GuardianContract,
|
||||
SpendingLimit,
|
||||
TimeLockConfig,
|
||||
GuardianConfig,
|
||||
create_guardian_contract,
|
||||
CONSERVATIVE_CONFIG,
|
||||
AGGRESSIVE_CONFIG,
|
||||
HIGH_SECURITY_CONFIG
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentSecurityProfile:
|
||||
"""Security profile for an agent"""
|
||||
agent_address: str
|
||||
security_level: str # "conservative", "aggressive", "high_security"
|
||||
guardian_addresses: List[str]
|
||||
custom_limits: Optional[Dict] = None
|
||||
enabled: bool = True
|
||||
created_at: datetime = None
|
||||
|
||||
def __post_init__(self):
|
||||
if self.created_at is None:
|
||||
self.created_at = datetime.utcnow()
|
||||
|
||||
|
||||
class AgentWalletSecurity:
|
||||
"""
|
||||
Security manager for autonomous agent wallets
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.agent_profiles: Dict[str, AgentSecurityProfile] = {}
|
||||
self.guardian_contracts: Dict[str, GuardianContract] = {}
|
||||
self.security_events: List[Dict] = []
|
||||
|
||||
# Default configurations
|
||||
self.configurations = {
|
||||
"conservative": CONSERVATIVE_CONFIG,
|
||||
"aggressive": AGGRESSIVE_CONFIG,
|
||||
"high_security": HIGH_SECURITY_CONFIG
|
||||
}
|
||||
|
||||
def register_agent(self,
|
||||
agent_address: str,
|
||||
security_level: str = "conservative",
|
||||
guardian_addresses: List[str] = None,
|
||||
custom_limits: Dict = None) -> Dict:
|
||||
"""
|
||||
Register an agent for security protection
|
||||
|
||||
Args:
|
||||
agent_address: Agent wallet address
|
||||
security_level: Security level (conservative, aggressive, high_security)
|
||||
guardian_addresses: List of guardian addresses for recovery
|
||||
custom_limits: Custom spending limits (overrides security_level)
|
||||
|
||||
Returns:
|
||||
Registration result
|
||||
"""
|
||||
try:
|
||||
agent_address = to_checksum_address(agent_address)
|
||||
|
||||
if agent_address in self.agent_profiles:
|
||||
return {
|
||||
"status": "error",
|
||||
"reason": "Agent already registered"
|
||||
}
|
||||
|
||||
# Validate security level
|
||||
if security_level not in self.configurations:
|
||||
return {
|
||||
"status": "error",
|
||||
"reason": f"Invalid security level: {security_level}"
|
||||
}
|
||||
|
||||
# Default guardians if none provided
|
||||
if guardian_addresses is None:
|
||||
guardian_addresses = [agent_address] # Self-guardian (should be overridden)
|
||||
|
||||
# Validate guardian addresses
|
||||
guardian_addresses = [to_checksum_address(addr) for addr in guardian_addresses]
|
||||
|
||||
# Create security profile
|
||||
profile = AgentSecurityProfile(
|
||||
agent_address=agent_address,
|
||||
security_level=security_level,
|
||||
guardian_addresses=guardian_addresses,
|
||||
custom_limits=custom_limits
|
||||
)
|
||||
|
||||
# Create guardian contract
|
||||
config = self.configurations[security_level]
|
||||
if custom_limits:
|
||||
config.update(custom_limits)
|
||||
|
||||
guardian_contract = create_guardian_contract(
|
||||
agent_address=agent_address,
|
||||
guardians=guardian_addresses,
|
||||
**config
|
||||
)
|
||||
|
||||
# Store profile and contract
|
||||
self.agent_profiles[agent_address] = profile
|
||||
self.guardian_contracts[agent_address] = guardian_contract
|
||||
|
||||
# Log security event
|
||||
self._log_security_event(
|
||||
event_type="agent_registered",
|
||||
agent_address=agent_address,
|
||||
security_level=security_level,
|
||||
guardian_count=len(guardian_addresses)
|
||||
)
|
||||
|
||||
return {
|
||||
"status": "registered",
|
||||
"agent_address": agent_address,
|
||||
"security_level": security_level,
|
||||
"guardian_addresses": guardian_addresses,
|
||||
"limits": guardian_contract.config.limits,
|
||||
"time_lock_threshold": guardian_contract.config.time_lock.threshold,
|
||||
"registered_at": profile.created_at.isoformat()
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
return {
|
||||
"status": "error",
|
||||
"reason": f"Registration failed: {str(e)}"
|
||||
}
|
||||
|
||||
def protect_transaction(self,
|
||||
agent_address: str,
|
||||
to_address: str,
|
||||
amount: int,
|
||||
data: str = "") -> Dict:
|
||||
"""
|
||||
Protect a transaction with guardian contract
|
||||
|
||||
Args:
|
||||
agent_address: Agent wallet address
|
||||
to_address: Recipient address
|
||||
amount: Amount to transfer
|
||||
data: Transaction data
|
||||
|
||||
Returns:
|
||||
Protection result
|
||||
"""
|
||||
try:
|
||||
agent_address = to_checksum_address(agent_address)
|
||||
|
||||
# Check if agent is registered
|
||||
if agent_address not in self.agent_profiles:
|
||||
return {
|
||||
"status": "unprotected",
|
||||
"reason": "Agent not registered for security protection",
|
||||
"suggestion": "Register agent with register_agent() first"
|
||||
}
|
||||
|
||||
# Check if protection is enabled
|
||||
profile = self.agent_profiles[agent_address]
|
||||
if not profile.enabled:
|
||||
return {
|
||||
"status": "unprotected",
|
||||
"reason": "Security protection disabled for this agent"
|
||||
}
|
||||
|
||||
# Get guardian contract
|
||||
guardian_contract = self.guardian_contracts[agent_address]
|
||||
|
||||
# Initiate transaction protection
|
||||
result = guardian_contract.initiate_transaction(to_address, amount, data)
|
||||
|
||||
# Log security event
|
||||
self._log_security_event(
|
||||
event_type="transaction_protected",
|
||||
agent_address=agent_address,
|
||||
to_address=to_address,
|
||||
amount=amount,
|
||||
protection_status=result["status"]
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
return {
|
||||
"status": "error",
|
||||
"reason": f"Transaction protection failed: {str(e)}"
|
||||
}
|
||||
|
||||
def execute_protected_transaction(self,
|
||||
agent_address: str,
|
||||
operation_id: str,
|
||||
signature: str) -> Dict:
|
||||
"""
|
||||
Execute a previously protected transaction
|
||||
|
||||
Args:
|
||||
agent_address: Agent wallet address
|
||||
operation_id: Operation ID from protection
|
||||
signature: Transaction signature
|
||||
|
||||
Returns:
|
||||
Execution result
|
||||
"""
|
||||
try:
|
||||
agent_address = to_checksum_address(agent_address)
|
||||
|
||||
if agent_address not in self.guardian_contracts:
|
||||
return {
|
||||
"status": "error",
|
||||
"reason": "Agent not registered"
|
||||
}
|
||||
|
||||
guardian_contract = self.guardian_contracts[agent_address]
|
||||
result = guardian_contract.execute_transaction(operation_id, signature)
|
||||
|
||||
# Log security event
|
||||
if result["status"] == "executed":
|
||||
self._log_security_event(
|
||||
event_type="transaction_executed",
|
||||
agent_address=agent_address,
|
||||
operation_id=operation_id,
|
||||
transaction_hash=result.get("transaction_hash")
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
return {
|
||||
"status": "error",
|
||||
"reason": f"Transaction execution failed: {str(e)}"
|
||||
}
|
||||
|
||||
def emergency_pause_agent(self, agent_address: str, guardian_address: str) -> Dict:
|
||||
"""
|
||||
Emergency pause an agent's operations
|
||||
|
||||
Args:
|
||||
agent_address: Agent wallet address
|
||||
guardian_address: Guardian address initiating pause
|
||||
|
||||
Returns:
|
||||
Pause result
|
||||
"""
|
||||
try:
|
||||
agent_address = to_checksum_address(agent_address)
|
||||
guardian_address = to_checksum_address(guardian_address)
|
||||
|
||||
if agent_address not in self.guardian_contracts:
|
||||
return {
|
||||
"status": "error",
|
||||
"reason": "Agent not registered"
|
||||
}
|
||||
|
||||
guardian_contract = self.guardian_contracts[agent_address]
|
||||
result = guardian_contract.emergency_pause(guardian_address)
|
||||
|
||||
# Log security event
|
||||
if result["status"] == "paused":
|
||||
self._log_security_event(
|
||||
event_type="emergency_pause",
|
||||
agent_address=agent_address,
|
||||
guardian_address=guardian_address
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
return {
|
||||
"status": "error",
|
||||
"reason": f"Emergency pause failed: {str(e)}"
|
||||
}
|
||||
|
||||
def update_agent_security(self,
|
||||
agent_address: str,
|
||||
new_limits: Dict,
|
||||
guardian_address: str) -> Dict:
|
||||
"""
|
||||
Update security limits for an agent
|
||||
|
||||
Args:
|
||||
agent_address: Agent wallet address
|
||||
new_limits: New spending limits
|
||||
guardian_address: Guardian address making the change
|
||||
|
||||
Returns:
|
||||
Update result
|
||||
"""
|
||||
try:
|
||||
agent_address = to_checksum_address(agent_address)
|
||||
guardian_address = to_checksum_address(guardian_address)
|
||||
|
||||
if agent_address not in self.guardian_contracts:
|
||||
return {
|
||||
"status": "error",
|
||||
"reason": "Agent not registered"
|
||||
}
|
||||
|
||||
guardian_contract = self.guardian_contracts[agent_address]
|
||||
|
||||
# Create new spending limits
|
||||
limits = SpendingLimit(
|
||||
per_transaction=new_limits.get("per_transaction", 1000),
|
||||
per_hour=new_limits.get("per_hour", 5000),
|
||||
per_day=new_limits.get("per_day", 20000),
|
||||
per_week=new_limits.get("per_week", 100000)
|
||||
)
|
||||
|
||||
result = guardian_contract.update_limits(limits, guardian_address)
|
||||
|
||||
# Log security event
|
||||
if result["status"] == "updated":
|
||||
self._log_security_event(
|
||||
event_type="security_limits_updated",
|
||||
agent_address=agent_address,
|
||||
guardian_address=guardian_address,
|
||||
new_limits=new_limits
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
return {
|
||||
"status": "error",
|
||||
"reason": f"Security update failed: {str(e)}"
|
||||
}
|
||||
|
||||
def get_agent_security_status(self, agent_address: str) -> Dict:
|
||||
"""
|
||||
Get security status for an agent
|
||||
|
||||
Args:
|
||||
agent_address: Agent wallet address
|
||||
|
||||
Returns:
|
||||
Security status
|
||||
"""
|
||||
try:
|
||||
agent_address = to_checksum_address(agent_address)
|
||||
|
||||
if agent_address not in self.agent_profiles:
|
||||
return {
|
||||
"status": "not_registered",
|
||||
"message": "Agent not registered for security protection"
|
||||
}
|
||||
|
||||
profile = self.agent_profiles[agent_address]
|
||||
guardian_contract = self.guardian_contracts[agent_address]
|
||||
|
||||
return {
|
||||
"status": "protected",
|
||||
"agent_address": agent_address,
|
||||
"security_level": profile.security_level,
|
||||
"enabled": profile.enabled,
|
||||
"guardian_addresses": profile.guardian_addresses,
|
||||
"registered_at": profile.created_at.isoformat(),
|
||||
"spending_status": guardian_contract.get_spending_status(),
|
||||
"pending_operations": guardian_contract.get_pending_operations(),
|
||||
"recent_activity": guardian_contract.get_operation_history(10)
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
return {
|
||||
"status": "error",
|
||||
"reason": f"Status check failed: {str(e)}"
|
||||
}
|
||||
|
||||
def list_protected_agents(self) -> List[Dict]:
|
||||
"""List all protected agents"""
|
||||
agents = []
|
||||
|
||||
for agent_address, profile in self.agent_profiles.items():
|
||||
guardian_contract = self.guardian_contracts[agent_address]
|
||||
|
||||
agents.append({
|
||||
"agent_address": agent_address,
|
||||
"security_level": profile.security_level,
|
||||
"enabled": profile.enabled,
|
||||
"guardian_count": len(profile.guardian_addresses),
|
||||
"pending_operations": len(guardian_contract.pending_operations),
|
||||
"paused": guardian_contract.paused,
|
||||
"emergency_mode": guardian_contract.emergency_mode,
|
||||
"registered_at": profile.created_at.isoformat()
|
||||
})
|
||||
|
||||
return sorted(agents, key=lambda x: x["registered_at"], reverse=True)
|
||||
|
||||
def get_security_events(self, agent_address: str = None, limit: int = 50) -> List[Dict]:
|
||||
"""
|
||||
Get security events
|
||||
|
||||
Args:
|
||||
agent_address: Filter by agent address (optional)
|
||||
limit: Maximum number of events
|
||||
|
||||
Returns:
|
||||
Security events
|
||||
"""
|
||||
events = self.security_events
|
||||
|
||||
if agent_address:
|
||||
agent_address = to_checksum_address(agent_address)
|
||||
events = [e for e in events if e.get("agent_address") == agent_address]
|
||||
|
||||
return sorted(events, key=lambda x: x["timestamp"], reverse=True)[:limit]
|
||||
|
||||
def _log_security_event(self, **kwargs):
|
||||
"""Log a security event"""
|
||||
event = {
|
||||
"timestamp": datetime.utcnow().isoformat(),
|
||||
**kwargs
|
||||
}
|
||||
self.security_events.append(event)
|
||||
|
||||
def disable_agent_protection(self, agent_address: str, guardian_address: str) -> Dict:
|
||||
"""
|
||||
Disable protection for an agent (guardian only)
|
||||
|
||||
Args:
|
||||
agent_address: Agent wallet address
|
||||
guardian_address: Guardian address
|
||||
|
||||
Returns:
|
||||
Disable result
|
||||
"""
|
||||
try:
|
||||
agent_address = to_checksum_address(agent_address)
|
||||
guardian_address = to_checksum_address(guardian_address)
|
||||
|
||||
if agent_address not in self.agent_profiles:
|
||||
return {
|
||||
"status": "error",
|
||||
"reason": "Agent not registered"
|
||||
}
|
||||
|
||||
profile = self.agent_profiles[agent_address]
|
||||
|
||||
if guardian_address not in profile.guardian_addresses:
|
||||
return {
|
||||
"status": "error",
|
||||
"reason": "Not authorized: not a guardian"
|
||||
}
|
||||
|
||||
profile.enabled = False
|
||||
|
||||
# Log security event
|
||||
self._log_security_event(
|
||||
event_type="protection_disabled",
|
||||
agent_address=agent_address,
|
||||
guardian_address=guardian_address
|
||||
)
|
||||
|
||||
return {
|
||||
"status": "disabled",
|
||||
"agent_address": agent_address,
|
||||
"disabled_at": datetime.utcnow().isoformat(),
|
||||
"guardian": guardian_address
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
return {
|
||||
"status": "error",
|
||||
"reason": f"Disable protection failed: {str(e)}"
|
||||
}
|
||||
|
||||
|
||||
# Global security manager instance
|
||||
agent_wallet_security = AgentWalletSecurity()
|
||||
|
||||
|
||||
# Convenience functions for common operations
|
||||
def register_agent_for_protection(agent_address: str,
|
||||
security_level: str = "conservative",
|
||||
guardians: List[str] = None) -> Dict:
|
||||
"""Register an agent for security protection"""
|
||||
return agent_wallet_security.register_agent(
|
||||
agent_address=agent_address,
|
||||
security_level=security_level,
|
||||
guardian_addresses=guardians
|
||||
)
|
||||
|
||||
|
||||
def protect_agent_transaction(agent_address: str,
|
||||
to_address: str,
|
||||
amount: int,
|
||||
data: str = "") -> Dict:
|
||||
"""Protect a transaction for an agent"""
|
||||
return agent_wallet_security.protect_transaction(
|
||||
agent_address=agent_address,
|
||||
to_address=to_address,
|
||||
amount=amount,
|
||||
data=data
|
||||
)
|
||||
|
||||
|
||||
def get_agent_security_summary(agent_address: str) -> Dict:
|
||||
"""Get security summary for an agent"""
|
||||
return agent_wallet_security.get_agent_security_status(agent_address)
|
||||
|
||||
|
||||
# Security audit and monitoring functions
|
||||
def generate_security_report() -> Dict:
|
||||
"""Generate comprehensive security report"""
|
||||
protected_agents = agent_wallet_security.list_protected_agents()
|
||||
|
||||
total_agents = len(protected_agents)
|
||||
active_agents = len([a for a in protected_agents if a["enabled"]])
|
||||
paused_agents = len([a for a in protected_agents if a["paused"]])
|
||||
emergency_agents = len([a for a in protected_agents if a["emergency_mode"]])
|
||||
|
||||
recent_events = agent_wallet_security.get_security_events(limit=20)
|
||||
|
||||
return {
|
||||
"generated_at": datetime.utcnow().isoformat(),
|
||||
"summary": {
|
||||
"total_protected_agents": total_agents,
|
||||
"active_agents": active_agents,
|
||||
"paused_agents": paused_agents,
|
||||
"emergency_mode_agents": emergency_agents,
|
||||
"protection_coverage": f"{(active_agents / total_agents * 100):.1f}%" if total_agents > 0 else "0%"
|
||||
},
|
||||
"agents": protected_agents,
|
||||
"recent_security_events": recent_events,
|
||||
"security_levels": {
|
||||
level: len([a for a in protected_agents if a["security_level"] == level])
|
||||
for level in ["conservative", "aggressive", "high_security"]
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def detect_suspicious_activity(agent_address: str, hours: int = 24) -> Dict:
|
||||
"""Detect suspicious activity for an agent"""
|
||||
status = agent_wallet_security.get_agent_security_status(agent_address)
|
||||
|
||||
if status["status"] != "protected":
|
||||
return {
|
||||
"status": "not_protected",
|
||||
"suspicious_activity": False
|
||||
}
|
||||
|
||||
spending_status = status["spending_status"]
|
||||
recent_events = agent_wallet_security.get_security_events(agent_address, limit=50)
|
||||
|
||||
# Suspicious patterns
|
||||
suspicious_patterns = []
|
||||
|
||||
# Check for rapid spending
|
||||
if spending_status["spent"]["current_hour"] > spending_status["current_limits"]["per_hour"] * 0.8:
|
||||
suspicious_patterns.append("High hourly spending rate")
|
||||
|
||||
# Check for many small transactions (potential dust attack)
|
||||
recent_tx_count = len([e for e in recent_events if e["event_type"] == "transaction_executed"])
|
||||
if recent_tx_count > 20:
|
||||
suspicious_patterns.append("High transaction frequency")
|
||||
|
||||
# Check for emergency pauses
|
||||
recent_pauses = len([e for e in recent_events if e["event_type"] == "emergency_pause"])
|
||||
if recent_pauses > 0:
|
||||
suspicious_patterns.append("Recent emergency pauses detected")
|
||||
|
||||
return {
|
||||
"status": "analyzed",
|
||||
"agent_address": agent_address,
|
||||
"suspicious_activity": len(suspicious_patterns) > 0,
|
||||
"suspicious_patterns": suspicious_patterns,
|
||||
"analysis_period_hours": hours,
|
||||
"analyzed_at": datetime.utcnow().isoformat()
|
||||
}
|
||||
@@ -0,0 +1,405 @@
|
||||
"""
|
||||
Fixed Guardian Configuration with Proper Guardian Setup
|
||||
Addresses the critical vulnerability where guardian lists were empty
|
||||
"""
|
||||
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta
|
||||
import json
|
||||
from eth_account import Account
|
||||
from eth_utils import to_checksum_address, keccak
|
||||
|
||||
from .guardian_contract import (
|
||||
SpendingLimit,
|
||||
TimeLockConfig,
|
||||
GuardianConfig,
|
||||
GuardianContract
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class GuardianSetup:
|
||||
"""Guardian setup configuration"""
|
||||
primary_guardian: str # Main guardian address
|
||||
backup_guardians: List[str] # Backup guardian addresses
|
||||
multisig_threshold: int # Number of signatures required
|
||||
emergency_contacts: List[str] # Additional emergency contacts
|
||||
|
||||
|
||||
class SecureGuardianManager:
|
||||
"""
|
||||
Secure guardian management with proper initialization
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.guardian_registrations: Dict[str, GuardianSetup] = {}
|
||||
self.guardian_contracts: Dict[str, GuardianContract] = {}
|
||||
|
||||
def create_guardian_setup(
|
||||
self,
|
||||
agent_address: str,
|
||||
owner_address: str,
|
||||
security_level: str = "conservative",
|
||||
custom_guardians: Optional[List[str]] = None
|
||||
) -> GuardianSetup:
|
||||
"""
|
||||
Create a proper guardian setup for an agent
|
||||
|
||||
Args:
|
||||
agent_address: Agent wallet address
|
||||
owner_address: Owner of the agent
|
||||
security_level: Security level (conservative, aggressive, high_security)
|
||||
custom_guardians: Optional custom guardian addresses
|
||||
|
||||
Returns:
|
||||
Guardian setup configuration
|
||||
"""
|
||||
agent_address = to_checksum_address(agent_address)
|
||||
owner_address = to_checksum_address(owner_address)
|
||||
|
||||
# Determine guardian requirements based on security level
|
||||
if security_level == "conservative":
|
||||
required_guardians = 3
|
||||
multisig_threshold = 2
|
||||
elif security_level == "aggressive":
|
||||
required_guardians = 2
|
||||
multisig_threshold = 2
|
||||
elif security_level == "high_security":
|
||||
required_guardians = 5
|
||||
multisig_threshold = 3
|
||||
else:
|
||||
raise ValueError(f"Invalid security level: {security_level}")
|
||||
|
||||
# Build guardian list
|
||||
guardians = []
|
||||
|
||||
# Always include the owner as primary guardian
|
||||
guardians.append(owner_address)
|
||||
|
||||
# Add custom guardians if provided
|
||||
if custom_guardians:
|
||||
for guardian in custom_guardians:
|
||||
guardian = to_checksum_address(guardian)
|
||||
if guardian not in guardians:
|
||||
guardians.append(guardian)
|
||||
|
||||
# Generate backup guardians if needed
|
||||
while len(guardians) < required_guardians:
|
||||
# Generate a deterministic backup guardian based on agent address
|
||||
# In production, these would be trusted service addresses
|
||||
backup_index = len(guardians) - 1 # -1 because owner is already included
|
||||
backup_guardian = self._generate_backup_guardian(agent_address, backup_index)
|
||||
|
||||
if backup_guardian not in guardians:
|
||||
guardians.append(backup_guardian)
|
||||
|
||||
# Create setup
|
||||
setup = GuardianSetup(
|
||||
primary_guardian=owner_address,
|
||||
backup_guardians=[g for g in guardians if g != owner_address],
|
||||
multisig_threshold=multisig_threshold,
|
||||
emergency_contacts=guardians.copy()
|
||||
)
|
||||
|
||||
self.guardian_registrations[agent_address] = setup
|
||||
|
||||
return setup
|
||||
|
||||
def _generate_backup_guardian(self, agent_address: str, index: int) -> str:
|
||||
"""
|
||||
Generate deterministic backup guardian address
|
||||
|
||||
In production, these would be pre-registered trusted guardian addresses
|
||||
"""
|
||||
# Create a deterministic address based on agent address and index
|
||||
seed = f"{agent_address}_{index}_backup_guardian"
|
||||
hash_result = keccak(seed.encode())
|
||||
|
||||
# Use the hash to generate a valid address
|
||||
address_bytes = hash_result[-20:] # Take last 20 bytes
|
||||
address = "0x" + address_bytes.hex()
|
||||
|
||||
return to_checksum_address(address)
|
||||
|
||||
def create_secure_guardian_contract(
|
||||
self,
|
||||
agent_address: str,
|
||||
security_level: str = "conservative",
|
||||
custom_guardians: Optional[List[str]] = None
|
||||
) -> GuardianContract:
|
||||
"""
|
||||
Create a guardian contract with proper guardian configuration
|
||||
|
||||
Args:
|
||||
agent_address: Agent wallet address
|
||||
security_level: Security level
|
||||
custom_guardians: Optional custom guardian addresses
|
||||
|
||||
Returns:
|
||||
Configured guardian contract
|
||||
"""
|
||||
# Create guardian setup
|
||||
setup = self.create_guardian_setup(
|
||||
agent_address=agent_address,
|
||||
owner_address=agent_address, # Agent is its own owner initially
|
||||
security_level=security_level,
|
||||
custom_guardians=custom_guardians
|
||||
)
|
||||
|
||||
# Get security configuration
|
||||
config = self._get_security_config(security_level, setup)
|
||||
|
||||
# Create contract
|
||||
contract = GuardianContract(agent_address, config)
|
||||
|
||||
# Store contract
|
||||
self.guardian_contracts[agent_address] = contract
|
||||
|
||||
return contract
|
||||
|
||||
def _get_security_config(self, security_level: str, setup: GuardianSetup) -> GuardianConfig:
|
||||
"""Get security configuration with proper guardian list"""
|
||||
|
||||
# Build guardian list
|
||||
all_guardians = [setup.primary_guardian] + setup.backup_guardians
|
||||
|
||||
if security_level == "conservative":
|
||||
return GuardianConfig(
|
||||
limits=SpendingLimit(
|
||||
per_transaction=1000,
|
||||
per_hour=5000,
|
||||
per_day=20000,
|
||||
per_week=100000
|
||||
),
|
||||
time_lock=TimeLockConfig(
|
||||
threshold=5000,
|
||||
delay_hours=24,
|
||||
max_delay_hours=168
|
||||
),
|
||||
guardians=all_guardians,
|
||||
pause_enabled=True,
|
||||
emergency_mode=False,
|
||||
multisig_threshold=setup.multisig_threshold
|
||||
)
|
||||
|
||||
elif security_level == "aggressive":
|
||||
return GuardianConfig(
|
||||
limits=SpendingLimit(
|
||||
per_transaction=5000,
|
||||
per_hour=25000,
|
||||
per_day=100000,
|
||||
per_week=500000
|
||||
),
|
||||
time_lock=TimeLockConfig(
|
||||
threshold=20000,
|
||||
delay_hours=12,
|
||||
max_delay_hours=72
|
||||
),
|
||||
guardians=all_guardians,
|
||||
pause_enabled=True,
|
||||
emergency_mode=False,
|
||||
multisig_threshold=setup.multisig_threshold
|
||||
)
|
||||
|
||||
elif security_level == "high_security":
|
||||
return GuardianConfig(
|
||||
limits=SpendingLimit(
|
||||
per_transaction=500,
|
||||
per_hour=2000,
|
||||
per_day=8000,
|
||||
per_week=40000
|
||||
),
|
||||
time_lock=TimeLockConfig(
|
||||
threshold=2000,
|
||||
delay_hours=48,
|
||||
max_delay_hours=168
|
||||
),
|
||||
guardians=all_guardians,
|
||||
pause_enabled=True,
|
||||
emergency_mode=False,
|
||||
multisig_threshold=setup.multisig_threshold
|
||||
)
|
||||
|
||||
else:
|
||||
raise ValueError(f"Invalid security level: {security_level}")
|
||||
|
||||
def test_emergency_pause(self, agent_address: str, guardian_address: str) -> Dict:
|
||||
"""
|
||||
Test emergency pause functionality
|
||||
|
||||
Args:
|
||||
agent_address: Agent address
|
||||
guardian_address: Guardian attempting pause
|
||||
|
||||
Returns:
|
||||
Test result
|
||||
"""
|
||||
if agent_address not in self.guardian_contracts:
|
||||
return {
|
||||
"status": "error",
|
||||
"reason": "Agent not registered"
|
||||
}
|
||||
|
||||
contract = self.guardian_contracts[agent_address]
|
||||
return contract.emergency_pause(guardian_address)
|
||||
|
||||
def verify_guardian_authorization(self, agent_address: str, guardian_address: str) -> bool:
|
||||
"""
|
||||
Verify if a guardian is authorized for an agent
|
||||
|
||||
Args:
|
||||
agent_address: Agent address
|
||||
guardian_address: Guardian address to verify
|
||||
|
||||
Returns:
|
||||
True if guardian is authorized
|
||||
"""
|
||||
if agent_address not in self.guardian_registrations:
|
||||
return False
|
||||
|
||||
setup = self.guardian_registrations[agent_address]
|
||||
all_guardians = [setup.primary_guardian] + setup.backup_guardians
|
||||
|
||||
return to_checksum_address(guardian_address) in [
|
||||
to_checksum_address(g) for g in all_guardians
|
||||
]
|
||||
|
||||
def get_guardian_summary(self, agent_address: str) -> Dict:
|
||||
"""
|
||||
Get guardian setup summary for an agent
|
||||
|
||||
Args:
|
||||
agent_address: Agent address
|
||||
|
||||
Returns:
|
||||
Guardian summary
|
||||
"""
|
||||
if agent_address not in self.guardian_registrations:
|
||||
return {"error": "Agent not registered"}
|
||||
|
||||
setup = self.guardian_registrations[agent_address]
|
||||
contract = self.guardian_contracts.get(agent_address)
|
||||
|
||||
return {
|
||||
"agent_address": agent_address,
|
||||
"primary_guardian": setup.primary_guardian,
|
||||
"backup_guardians": setup.backup_guardians,
|
||||
"total_guardians": len(setup.backup_guardians) + 1,
|
||||
"multisig_threshold": setup.multisig_threshold,
|
||||
"emergency_contacts": setup.emergency_contacts,
|
||||
"contract_status": contract.get_spending_status() if contract else None,
|
||||
"pause_functional": contract is not None and len(setup.backup_guardians) > 0
|
||||
}
|
||||
|
||||
|
||||
# Fixed security configurations with proper guardians
|
||||
def get_fixed_conservative_config(agent_address: str, owner_address: str) -> GuardianConfig:
|
||||
"""Get fixed conservative configuration with proper guardians"""
|
||||
return GuardianConfig(
|
||||
limits=SpendingLimit(
|
||||
per_transaction=1000,
|
||||
per_hour=5000,
|
||||
per_day=20000,
|
||||
per_week=100000
|
||||
),
|
||||
time_lock=TimeLockConfig(
|
||||
threshold=5000,
|
||||
delay_hours=24,
|
||||
max_delay_hours=168
|
||||
),
|
||||
guardians=[owner_address], # At least the owner
|
||||
pause_enabled=True,
|
||||
emergency_mode=False
|
||||
)
|
||||
|
||||
|
||||
def get_fixed_aggressive_config(agent_address: str, owner_address: str) -> GuardianConfig:
|
||||
"""Get fixed aggressive configuration with proper guardians"""
|
||||
return GuardianConfig(
|
||||
limits=SpendingLimit(
|
||||
per_transaction=5000,
|
||||
per_hour=25000,
|
||||
per_day=100000,
|
||||
per_week=500000
|
||||
),
|
||||
time_lock=TimeLockConfig(
|
||||
threshold=20000,
|
||||
delay_hours=12,
|
||||
max_delay_hours=72
|
||||
),
|
||||
guardians=[owner_address], # At least the owner
|
||||
pause_enabled=True,
|
||||
emergency_mode=False
|
||||
)
|
||||
|
||||
|
||||
def get_fixed_high_security_config(agent_address: str, owner_address: str) -> GuardianConfig:
|
||||
"""Get fixed high security configuration with proper guardians"""
|
||||
return GuardianConfig(
|
||||
limits=SpendingLimit(
|
||||
per_transaction=500,
|
||||
per_hour=2000,
|
||||
per_day=8000,
|
||||
per_week=40000
|
||||
),
|
||||
time_lock=TimeLockConfig(
|
||||
threshold=2000,
|
||||
delay_hours=48,
|
||||
max_delay_hours=168
|
||||
),
|
||||
guardians=[owner_address], # At least the owner
|
||||
pause_enabled=True,
|
||||
emergency_mode=False
|
||||
)
|
||||
|
||||
|
||||
# Global secure guardian manager
|
||||
secure_guardian_manager = SecureGuardianManager()
|
||||
|
||||
|
||||
# Convenience function for secure agent registration
|
||||
def register_agent_with_guardians(
|
||||
agent_address: str,
|
||||
owner_address: str,
|
||||
security_level: str = "conservative",
|
||||
custom_guardians: Optional[List[str]] = None
|
||||
) -> Dict:
|
||||
"""
|
||||
Register an agent with proper guardian configuration
|
||||
|
||||
Args:
|
||||
agent_address: Agent wallet address
|
||||
owner_address: Owner address
|
||||
security_level: Security level
|
||||
custom_guardians: Optional custom guardians
|
||||
|
||||
Returns:
|
||||
Registration result
|
||||
"""
|
||||
try:
|
||||
# Create secure guardian contract
|
||||
contract = secure_guardian_manager.create_secure_guardian_contract(
|
||||
agent_address=agent_address,
|
||||
security_level=security_level,
|
||||
custom_guardians=custom_guardians
|
||||
)
|
||||
|
||||
# Get guardian summary
|
||||
summary = secure_guardian_manager.get_guardian_summary(agent_address)
|
||||
|
||||
return {
|
||||
"status": "registered",
|
||||
"agent_address": agent_address,
|
||||
"security_level": security_level,
|
||||
"guardian_count": summary["total_guardians"],
|
||||
"multisig_threshold": summary["multisig_threshold"],
|
||||
"pause_functional": summary["pause_functional"],
|
||||
"registered_at": datetime.utcnow().isoformat()
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
return {
|
||||
"status": "error",
|
||||
"reason": f"Registration failed: {str(e)}"
|
||||
}
|
||||
@@ -0,0 +1,477 @@
|
||||
"""
|
||||
AITBC Guardian Contract - Spending Limit Protection for Agent Wallets
|
||||
|
||||
This contract implements a spending limit guardian that protects autonomous agent
|
||||
wallets from unlimited spending in case of compromise. It provides:
|
||||
- Per-transaction spending limits
|
||||
- Per-period (daily/hourly) spending caps
|
||||
- Time-lock for large withdrawals
|
||||
- Emergency pause functionality
|
||||
- Multi-signature recovery for critical operations
|
||||
"""
|
||||
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta
|
||||
import json
|
||||
from eth_account import Account
|
||||
from eth_utils import to_checksum_address, keccak
|
||||
|
||||
|
||||
@dataclass
|
||||
class SpendingLimit:
|
||||
"""Spending limit configuration"""
|
||||
per_transaction: int # Maximum per transaction
|
||||
per_hour: int # Maximum per hour
|
||||
per_day: int # Maximum per day
|
||||
per_week: int # Maximum per week
|
||||
|
||||
@dataclass
|
||||
class TimeLockConfig:
|
||||
"""Time lock configuration for large withdrawals"""
|
||||
threshold: int # Amount that triggers time lock
|
||||
delay_hours: int # Delay period in hours
|
||||
max_delay_hours: int # Maximum delay period
|
||||
|
||||
|
||||
@dataclass
|
||||
class GuardianConfig:
|
||||
"""Complete guardian configuration"""
|
||||
limits: SpendingLimit
|
||||
time_lock: TimeLockConfig
|
||||
guardians: List[str] # Guardian addresses for recovery
|
||||
pause_enabled: bool = True
|
||||
emergency_mode: bool = False
|
||||
|
||||
|
||||
class GuardianContract:
|
||||
"""
|
||||
Guardian contract implementation for agent wallet protection
|
||||
"""
|
||||
|
||||
def __init__(self, agent_address: str, config: GuardianConfig):
|
||||
self.agent_address = to_checksum_address(agent_address)
|
||||
self.config = config
|
||||
self.spending_history: List[Dict] = []
|
||||
self.pending_operations: Dict[str, Dict] = {}
|
||||
self.paused = False
|
||||
self.emergency_mode = False
|
||||
|
||||
# Contract state
|
||||
self.nonce = 0
|
||||
self.guardian_approvals: Dict[str, bool] = {}
|
||||
|
||||
def _get_period_key(self, timestamp: datetime, period: str) -> str:
|
||||
"""Generate period key for spending tracking"""
|
||||
if period == "hour":
|
||||
return timestamp.strftime("%Y-%m-%d-%H")
|
||||
elif period == "day":
|
||||
return timestamp.strftime("%Y-%m-%d")
|
||||
elif period == "week":
|
||||
# Get week number (Monday as first day)
|
||||
week_num = timestamp.isocalendar()[1]
|
||||
return f"{timestamp.year}-W{week_num:02d}"
|
||||
else:
|
||||
raise ValueError(f"Invalid period: {period}")
|
||||
|
||||
def _get_spent_in_period(self, period: str, timestamp: datetime = None) -> int:
|
||||
"""Calculate total spent in given period"""
|
||||
if timestamp is None:
|
||||
timestamp = datetime.utcnow()
|
||||
|
||||
period_key = self._get_period_key(timestamp, period)
|
||||
|
||||
total = 0
|
||||
for record in self.spending_history:
|
||||
record_time = datetime.fromisoformat(record["timestamp"])
|
||||
record_period = self._get_period_key(record_time, period)
|
||||
|
||||
if record_period == period_key and record["status"] == "completed":
|
||||
total += record["amount"]
|
||||
|
||||
return total
|
||||
|
||||
def _check_spending_limits(self, amount: int, timestamp: datetime = None) -> Tuple[bool, str]:
|
||||
"""Check if amount exceeds spending limits"""
|
||||
if timestamp is None:
|
||||
timestamp = datetime.utcnow()
|
||||
|
||||
# Check per-transaction limit
|
||||
if amount > self.config.limits.per_transaction:
|
||||
return False, f"Amount {amount} exceeds per-transaction limit {self.config.limits.per_transaction}"
|
||||
|
||||
# Check per-hour limit
|
||||
spent_hour = self._get_spent_in_period("hour", timestamp)
|
||||
if spent_hour + amount > self.config.limits.per_hour:
|
||||
return False, f"Hourly spending {spent_hour + amount} would exceed limit {self.config.limits.per_hour}"
|
||||
|
||||
# Check per-day limit
|
||||
spent_day = self._get_spent_in_period("day", timestamp)
|
||||
if spent_day + amount > self.config.limits.per_day:
|
||||
return False, f"Daily spending {spent_day + amount} would exceed limit {self.config.limits.per_day}"
|
||||
|
||||
# Check per-week limit
|
||||
spent_week = self._get_spent_in_period("week", timestamp)
|
||||
if spent_week + amount > self.config.limits.per_week:
|
||||
return False, f"Weekly spending {spent_week + amount} would exceed limit {self.config.limits.per_week}"
|
||||
|
||||
return True, "Spending limits check passed"
|
||||
|
||||
def _requires_time_lock(self, amount: int) -> bool:
|
||||
"""Check if amount requires time lock"""
|
||||
return amount >= self.config.time_lock.threshold
|
||||
|
||||
def _create_operation_hash(self, operation: Dict) -> str:
|
||||
"""Create hash for operation identification"""
|
||||
operation_str = json.dumps(operation, sort_keys=True)
|
||||
return keccak(operation_str.encode()).hex()
|
||||
|
||||
def initiate_transaction(self, to_address: str, amount: int, data: str = "") -> Dict:
|
||||
"""
|
||||
Initiate a transaction with guardian protection
|
||||
|
||||
Args:
|
||||
to_address: Recipient address
|
||||
amount: Amount to transfer
|
||||
data: Transaction data (optional)
|
||||
|
||||
Returns:
|
||||
Operation result with status and details
|
||||
"""
|
||||
# Check if paused
|
||||
if self.paused:
|
||||
return {
|
||||
"status": "rejected",
|
||||
"reason": "Guardian contract is paused",
|
||||
"operation_id": None
|
||||
}
|
||||
|
||||
# Check emergency mode
|
||||
if self.emergency_mode:
|
||||
return {
|
||||
"status": "rejected",
|
||||
"reason": "Emergency mode activated",
|
||||
"operation_id": None
|
||||
}
|
||||
|
||||
# Validate address
|
||||
try:
|
||||
to_address = to_checksum_address(to_address)
|
||||
except:
|
||||
return {
|
||||
"status": "rejected",
|
||||
"reason": "Invalid recipient address",
|
||||
"operation_id": None
|
||||
}
|
||||
|
||||
# Check spending limits
|
||||
limits_ok, limits_reason = self._check_spending_limits(amount)
|
||||
if not limits_ok:
|
||||
return {
|
||||
"status": "rejected",
|
||||
"reason": limits_reason,
|
||||
"operation_id": None
|
||||
}
|
||||
|
||||
# Create operation
|
||||
operation = {
|
||||
"type": "transaction",
|
||||
"to": to_address,
|
||||
"amount": amount,
|
||||
"data": data,
|
||||
"timestamp": datetime.utcnow().isoformat(),
|
||||
"nonce": self.nonce,
|
||||
"status": "pending"
|
||||
}
|
||||
|
||||
operation_id = self._create_operation_hash(operation)
|
||||
operation["operation_id"] = operation_id
|
||||
|
||||
# Check if time lock is required
|
||||
if self._requires_time_lock(amount):
|
||||
unlock_time = datetime.utcnow() + timedelta(hours=self.config.time_lock.delay_hours)
|
||||
operation["unlock_time"] = unlock_time.isoformat()
|
||||
operation["status"] = "time_locked"
|
||||
|
||||
# Store for later execution
|
||||
self.pending_operations[operation_id] = operation
|
||||
|
||||
return {
|
||||
"status": "time_locked",
|
||||
"operation_id": operation_id,
|
||||
"unlock_time": unlock_time.isoformat(),
|
||||
"delay_hours": self.config.time_lock.delay_hours,
|
||||
"message": f"Transaction requires {self.config.time_lock.delay_hours}h time lock"
|
||||
}
|
||||
|
||||
# Immediate execution for smaller amounts
|
||||
self.pending_operations[operation_id] = operation
|
||||
|
||||
return {
|
||||
"status": "approved",
|
||||
"operation_id": operation_id,
|
||||
"message": "Transaction approved for execution"
|
||||
}
|
||||
|
||||
def execute_transaction(self, operation_id: str, signature: str) -> Dict:
|
||||
"""
|
||||
Execute a previously approved transaction
|
||||
|
||||
Args:
|
||||
operation_id: Operation ID from initiate_transaction
|
||||
signature: Transaction signature from agent
|
||||
|
||||
Returns:
|
||||
Execution result
|
||||
"""
|
||||
if operation_id not in self.pending_operations:
|
||||
return {
|
||||
"status": "error",
|
||||
"reason": "Operation not found"
|
||||
}
|
||||
|
||||
operation = self.pending_operations[operation_id]
|
||||
|
||||
# Check if operation is time locked
|
||||
if operation["status"] == "time_locked":
|
||||
unlock_time = datetime.fromisoformat(operation["unlock_time"])
|
||||
if datetime.utcnow() < unlock_time:
|
||||
return {
|
||||
"status": "error",
|
||||
"reason": f"Operation locked until {unlock_time.isoformat()}"
|
||||
}
|
||||
|
||||
operation["status"] = "ready"
|
||||
|
||||
# Verify signature (simplified - in production, use proper verification)
|
||||
try:
|
||||
# In production, verify the signature matches the agent address
|
||||
# For now, we'll assume signature is valid
|
||||
pass
|
||||
except Exception as e:
|
||||
return {
|
||||
"status": "error",
|
||||
"reason": f"Invalid signature: {str(e)}"
|
||||
}
|
||||
|
||||
# Record the transaction
|
||||
record = {
|
||||
"operation_id": operation_id,
|
||||
"to": operation["to"],
|
||||
"amount": operation["amount"],
|
||||
"data": operation.get("data", ""),
|
||||
"timestamp": operation["timestamp"],
|
||||
"executed_at": datetime.utcnow().isoformat(),
|
||||
"status": "completed",
|
||||
"nonce": operation["nonce"]
|
||||
}
|
||||
|
||||
self.spending_history.append(record)
|
||||
self.nonce += 1
|
||||
|
||||
# Remove from pending
|
||||
del self.pending_operations[operation_id]
|
||||
|
||||
return {
|
||||
"status": "executed",
|
||||
"operation_id": operation_id,
|
||||
"transaction_hash": f"0x{keccak(f'{operation_id}{signature}'.encode()).hex()}",
|
||||
"executed_at": record["executed_at"]
|
||||
}
|
||||
|
||||
def emergency_pause(self, guardian_address: str) -> Dict:
|
||||
"""
|
||||
Emergency pause function (guardian only)
|
||||
|
||||
Args:
|
||||
guardian_address: Address of guardian initiating pause
|
||||
|
||||
Returns:
|
||||
Pause result
|
||||
"""
|
||||
if guardian_address not in self.config.guardians:
|
||||
return {
|
||||
"status": "rejected",
|
||||
"reason": "Not authorized: guardian address not recognized"
|
||||
}
|
||||
|
||||
self.paused = True
|
||||
self.emergency_mode = True
|
||||
|
||||
return {
|
||||
"status": "paused",
|
||||
"paused_at": datetime.utcnow().isoformat(),
|
||||
"guardian": guardian_address,
|
||||
"message": "Emergency pause activated - all operations halted"
|
||||
}
|
||||
|
||||
def emergency_unpause(self, guardian_signatures: List[str]) -> Dict:
|
||||
"""
|
||||
Emergency unpause function (requires multiple guardian signatures)
|
||||
|
||||
Args:
|
||||
guardian_signatures: Signatures from required guardians
|
||||
|
||||
Returns:
|
||||
Unpause result
|
||||
"""
|
||||
# In production, verify all guardian signatures
|
||||
required_signatures = len(self.config.guardians)
|
||||
if len(guardian_signatures) < required_signatures:
|
||||
return {
|
||||
"status": "rejected",
|
||||
"reason": f"Requires {required_signatures} guardian signatures, got {len(guardian_signatures)}"
|
||||
}
|
||||
|
||||
# Verify signatures (simplified)
|
||||
# In production, verify each signature matches a guardian address
|
||||
|
||||
self.paused = False
|
||||
self.emergency_mode = False
|
||||
|
||||
return {
|
||||
"status": "unpaused",
|
||||
"unpaused_at": datetime.utcnow().isoformat(),
|
||||
"message": "Emergency pause lifted - operations resumed"
|
||||
}
|
||||
|
||||
def update_limits(self, new_limits: SpendingLimit, guardian_address: str) -> Dict:
|
||||
"""
|
||||
Update spending limits (guardian only)
|
||||
|
||||
Args:
|
||||
new_limits: New spending limits
|
||||
guardian_address: Address of guardian making the change
|
||||
|
||||
Returns:
|
||||
Update result
|
||||
"""
|
||||
if guardian_address not in self.config.guardians:
|
||||
return {
|
||||
"status": "rejected",
|
||||
"reason": "Not authorized: guardian address not recognized"
|
||||
}
|
||||
|
||||
old_limits = self.config.limits
|
||||
self.config.limits = new_limits
|
||||
|
||||
return {
|
||||
"status": "updated",
|
||||
"old_limits": old_limits,
|
||||
"new_limits": new_limits,
|
||||
"updated_at": datetime.utcnow().isoformat(),
|
||||
"guardian": guardian_address
|
||||
}
|
||||
|
||||
def get_spending_status(self) -> Dict:
|
||||
"""Get current spending status and limits"""
|
||||
now = datetime.utcnow()
|
||||
|
||||
return {
|
||||
"agent_address": self.agent_address,
|
||||
"current_limits": self.config.limits,
|
||||
"spent": {
|
||||
"current_hour": self._get_spent_in_period("hour", now),
|
||||
"current_day": self._get_spent_in_period("day", now),
|
||||
"current_week": self._get_spent_in_period("week", now)
|
||||
},
|
||||
"remaining": {
|
||||
"current_hour": self.config.limits.per_hour - self._get_spent_in_period("hour", now),
|
||||
"current_day": self.config.limits.per_day - self._get_spent_in_period("day", now),
|
||||
"current_week": self.config.limits.per_week - self._get_spent_in_period("week", now)
|
||||
},
|
||||
"pending_operations": len(self.pending_operations),
|
||||
"paused": self.paused,
|
||||
"emergency_mode": self.emergency_mode,
|
||||
"nonce": self.nonce
|
||||
}
|
||||
|
||||
def get_operation_history(self, limit: int = 50) -> List[Dict]:
|
||||
"""Get operation history"""
|
||||
return sorted(self.spending_history, key=lambda x: x["timestamp"], reverse=True)[:limit]
|
||||
|
||||
def get_pending_operations(self) -> List[Dict]:
|
||||
"""Get all pending operations"""
|
||||
return list(self.pending_operations.values())
|
||||
|
||||
|
||||
# Factory function for creating guardian contracts
|
||||
def create_guardian_contract(
|
||||
agent_address: str,
|
||||
per_transaction: int = 1000,
|
||||
per_hour: int = 5000,
|
||||
per_day: int = 20000,
|
||||
per_week: int = 100000,
|
||||
time_lock_threshold: int = 10000,
|
||||
time_lock_delay: int = 24,
|
||||
guardians: List[str] = None
|
||||
) -> GuardianContract:
|
||||
"""
|
||||
Create a guardian contract with default security parameters
|
||||
|
||||
Args:
|
||||
agent_address: The agent wallet address to protect
|
||||
per_transaction: Maximum amount per transaction
|
||||
per_hour: Maximum amount per hour
|
||||
per_day: Maximum amount per day
|
||||
per_week: Maximum amount per week
|
||||
time_lock_threshold: Amount that triggers time lock
|
||||
time_lock_delay: Time lock delay in hours
|
||||
guardians: List of guardian addresses
|
||||
|
||||
Returns:
|
||||
Configured GuardianContract instance
|
||||
"""
|
||||
if guardians is None:
|
||||
# Default to using the agent address as guardian (should be overridden)
|
||||
guardians = [agent_address]
|
||||
|
||||
limits = SpendingLimit(
|
||||
per_transaction=per_transaction,
|
||||
per_hour=per_hour,
|
||||
per_day=per_day,
|
||||
per_week=per_week
|
||||
)
|
||||
|
||||
time_lock = TimeLockConfig(
|
||||
threshold=time_lock_threshold,
|
||||
delay_hours=time_lock_delay,
|
||||
max_delay_hours=168 # 1 week max
|
||||
)
|
||||
|
||||
config = GuardianConfig(
|
||||
limits=limits,
|
||||
time_lock=time_lock,
|
||||
guardians=[to_checksum_address(g) for g in guardians]
|
||||
)
|
||||
|
||||
return GuardianContract(agent_address, config)
|
||||
|
||||
|
||||
# Example usage and security configurations
|
||||
CONSERVATIVE_CONFIG = {
|
||||
"per_transaction": 100, # $100 per transaction
|
||||
"per_hour": 500, # $500 per hour
|
||||
"per_day": 2000, # $2,000 per day
|
||||
"per_week": 10000, # $10,000 per week
|
||||
"time_lock_threshold": 1000, # Time lock over $1,000
|
||||
"time_lock_delay": 24 # 24 hour delay
|
||||
}
|
||||
|
||||
AGGRESSIVE_CONFIG = {
|
||||
"per_transaction": 1000, # $1,000 per transaction
|
||||
"per_hour": 5000, # $5,000 per hour
|
||||
"per_day": 20000, # $20,000 per day
|
||||
"per_week": 100000, # $100,000 per week
|
||||
"time_lock_threshold": 10000, # Time lock over $10,000
|
||||
"time_lock_delay": 12 # 12 hour delay
|
||||
}
|
||||
|
||||
HIGH_SECURITY_CONFIG = {
|
||||
"per_transaction": 50, # $50 per transaction
|
||||
"per_hour": 200, # $200 per hour
|
||||
"per_day": 1000, # $1,000 per day
|
||||
"per_week": 5000, # $5,000 per week
|
||||
"time_lock_threshold": 500, # Time lock over $500
|
||||
"time_lock_delay": 48 # 48 hour delay
|
||||
}
|
||||
@@ -0,0 +1,470 @@
|
||||
"""
|
||||
Persistent Spending Tracker - Database-Backed Security
|
||||
Fixes the critical vulnerability where spending limits were lost on restart
|
||||
"""
|
||||
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta
|
||||
from sqlalchemy import create_engine, Column, String, Integer, Float, DateTime, Index
|
||||
from sqlalchemy.ext.declarative import declarative_base
|
||||
from sqlalchemy.orm import sessionmaker, Session
|
||||
from eth_utils import to_checksum_address
|
||||
import json
|
||||
|
||||
Base = declarative_base()
|
||||
|
||||
|
||||
class SpendingRecord(Base):
|
||||
"""Database model for spending tracking"""
|
||||
__tablename__ = "spending_records"
|
||||
|
||||
id = Column(String, primary_key=True)
|
||||
agent_address = Column(String, index=True)
|
||||
period_type = Column(String, index=True) # hour, day, week
|
||||
period_key = Column(String, index=True)
|
||||
amount = Column(Float)
|
||||
transaction_hash = Column(String)
|
||||
timestamp = Column(DateTime, default=datetime.utcnow)
|
||||
|
||||
# Composite indexes for performance
|
||||
__table_args__ = (
|
||||
Index('idx_agent_period', 'agent_address', 'period_type', 'period_key'),
|
||||
Index('idx_timestamp', 'timestamp'),
|
||||
)
|
||||
|
||||
|
||||
class SpendingLimit(Base):
|
||||
"""Database model for spending limits"""
|
||||
__tablename__ = "spending_limits"
|
||||
|
||||
agent_address = Column(String, primary_key=True)
|
||||
per_transaction = Column(Float)
|
||||
per_hour = Column(Float)
|
||||
per_day = Column(Float)
|
||||
per_week = Column(Float)
|
||||
time_lock_threshold = Column(Float)
|
||||
time_lock_delay_hours = Column(Integer)
|
||||
updated_at = Column(DateTime, default=datetime.utcnow)
|
||||
updated_by = Column(String) # Guardian who updated
|
||||
|
||||
|
||||
class GuardianAuthorization(Base):
|
||||
"""Database model for guardian authorizations"""
|
||||
__tablename__ = "guardian_authorizations"
|
||||
|
||||
id = Column(String, primary_key=True)
|
||||
agent_address = Column(String, index=True)
|
||||
guardian_address = Column(String, index=True)
|
||||
is_active = Column(Boolean, default=True)
|
||||
added_at = Column(DateTime, default=datetime.utcnow)
|
||||
added_by = Column(String)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SpendingCheckResult:
|
||||
"""Result of spending limit check"""
|
||||
allowed: bool
|
||||
reason: str
|
||||
current_spent: Dict[str, float]
|
||||
remaining: Dict[str, float]
|
||||
requires_time_lock: bool
|
||||
time_lock_until: Optional[datetime] = None
|
||||
|
||||
|
||||
class PersistentSpendingTracker:
|
||||
"""
|
||||
Database-backed spending tracker that survives restarts
|
||||
"""
|
||||
|
||||
def __init__(self, database_url: str = "sqlite:///spending_tracker.db"):
|
||||
self.engine = create_engine(database_url)
|
||||
Base.metadata.create_all(self.engine)
|
||||
self.SessionLocal = sessionmaker(bind=self.engine)
|
||||
|
||||
def get_session(self) -> Session:
|
||||
"""Get database session"""
|
||||
return self.SessionLocal()
|
||||
|
||||
def _get_period_key(self, timestamp: datetime, period: str) -> str:
|
||||
"""Generate period key for spending tracking"""
|
||||
if period == "hour":
|
||||
return timestamp.strftime("%Y-%m-%d-%H")
|
||||
elif period == "day":
|
||||
return timestamp.strftime("%Y-%m-%d")
|
||||
elif period == "week":
|
||||
# Get week number (Monday as first day)
|
||||
week_num = timestamp.isocalendar()[1]
|
||||
return f"{timestamp.year}-W{week_num:02d}"
|
||||
else:
|
||||
raise ValueError(f"Invalid period: {period}")
|
||||
|
||||
def get_spent_in_period(self, agent_address: str, period: str, timestamp: datetime = None) -> float:
|
||||
"""
|
||||
Get total spent in given period from database
|
||||
|
||||
Args:
|
||||
agent_address: Agent wallet address
|
||||
period: Period type (hour, day, week)
|
||||
timestamp: Timestamp to check (default: now)
|
||||
|
||||
Returns:
|
||||
Total amount spent in period
|
||||
"""
|
||||
if timestamp is None:
|
||||
timestamp = datetime.utcnow()
|
||||
|
||||
period_key = self._get_period_key(timestamp, period)
|
||||
agent_address = to_checksum_address(agent_address)
|
||||
|
||||
with self.get_session() as session:
|
||||
total = session.query(SpendingRecord).filter(
|
||||
SpendingRecord.agent_address == agent_address,
|
||||
SpendingRecord.period_type == period,
|
||||
SpendingRecord.period_key == period_key
|
||||
).with_entities(SpendingRecord.amount).all()
|
||||
|
||||
return sum(record.amount for record in total)
|
||||
|
||||
def record_spending(self, agent_address: str, amount: float, transaction_hash: str, timestamp: datetime = None) -> bool:
|
||||
"""
|
||||
Record a spending transaction in the database
|
||||
|
||||
Args:
|
||||
agent_address: Agent wallet address
|
||||
amount: Amount spent
|
||||
transaction_hash: Transaction hash
|
||||
timestamp: Transaction timestamp (default: now)
|
||||
|
||||
Returns:
|
||||
True if recorded successfully
|
||||
"""
|
||||
if timestamp is None:
|
||||
timestamp = datetime.utcnow()
|
||||
|
||||
agent_address = to_checksum_address(agent_address)
|
||||
|
||||
try:
|
||||
with self.get_session() as session:
|
||||
# Record for all periods
|
||||
periods = ["hour", "day", "week"]
|
||||
|
||||
for period in periods:
|
||||
period_key = self._get_period_key(timestamp, period)
|
||||
|
||||
record = SpendingRecord(
|
||||
id=f"{transaction_hash}_{period}",
|
||||
agent_address=agent_address,
|
||||
period_type=period,
|
||||
period_key=period_key,
|
||||
amount=amount,
|
||||
transaction_hash=transaction_hash,
|
||||
timestamp=timestamp
|
||||
)
|
||||
|
||||
session.add(record)
|
||||
|
||||
session.commit()
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"Failed to record spending: {e}")
|
||||
return False
|
||||
|
||||
def check_spending_limits(self, agent_address: str, amount: float, timestamp: datetime = None) -> SpendingCheckResult:
|
||||
"""
|
||||
Check if amount exceeds spending limits using persistent data
|
||||
|
||||
Args:
|
||||
agent_address: Agent wallet address
|
||||
amount: Amount to check
|
||||
timestamp: Timestamp for check (default: now)
|
||||
|
||||
Returns:
|
||||
Spending check result
|
||||
"""
|
||||
if timestamp is None:
|
||||
timestamp = datetime.utcnow()
|
||||
|
||||
agent_address = to_checksum_address(agent_address)
|
||||
|
||||
# Get spending limits from database
|
||||
with self.get_session() as session:
|
||||
limits = session.query(SpendingLimit).filter(
|
||||
SpendingLimit.agent_address == agent_address
|
||||
).first()
|
||||
|
||||
if not limits:
|
||||
# Default limits if not set
|
||||
limits = SpendingLimit(
|
||||
agent_address=agent_address,
|
||||
per_transaction=1000.0,
|
||||
per_hour=5000.0,
|
||||
per_day=20000.0,
|
||||
per_week=100000.0,
|
||||
time_lock_threshold=5000.0,
|
||||
time_lock_delay_hours=24
|
||||
)
|
||||
session.add(limits)
|
||||
session.commit()
|
||||
|
||||
# Check each limit
|
||||
current_spent = {}
|
||||
remaining = {}
|
||||
|
||||
# Per-transaction limit
|
||||
if amount > limits.per_transaction:
|
||||
return SpendingCheckResult(
|
||||
allowed=False,
|
||||
reason=f"Amount {amount} exceeds per-transaction limit {limits.per_transaction}",
|
||||
current_spent=current_spent,
|
||||
remaining=remaining,
|
||||
requires_time_lock=False
|
||||
)
|
||||
|
||||
# Per-hour limit
|
||||
spent_hour = self.get_spent_in_period(agent_address, "hour", timestamp)
|
||||
current_spent["hour"] = spent_hour
|
||||
remaining["hour"] = limits.per_hour - spent_hour
|
||||
|
||||
if spent_hour + amount > limits.per_hour:
|
||||
return SpendingCheckResult(
|
||||
allowed=False,
|
||||
reason=f"Hourly spending {spent_hour + amount} would exceed limit {limits.per_hour}",
|
||||
current_spent=current_spent,
|
||||
remaining=remaining,
|
||||
requires_time_lock=False
|
||||
)
|
||||
|
||||
# Per-day limit
|
||||
spent_day = self.get_spent_in_period(agent_address, "day", timestamp)
|
||||
current_spent["day"] = spent_day
|
||||
remaining["day"] = limits.per_day - spent_day
|
||||
|
||||
if spent_day + amount > limits.per_day:
|
||||
return SpendingCheckResult(
|
||||
allowed=False,
|
||||
reason=f"Daily spending {spent_day + amount} would exceed limit {limits.per_day}",
|
||||
current_spent=current_spent,
|
||||
remaining=remaining,
|
||||
requires_time_lock=False
|
||||
)
|
||||
|
||||
# Per-week limit
|
||||
spent_week = self.get_spent_in_period(agent_address, "week", timestamp)
|
||||
current_spent["week"] = spent_week
|
||||
remaining["week"] = limits.per_week - spent_week
|
||||
|
||||
if spent_week + amount > limits.per_week:
|
||||
return SpendingCheckResult(
|
||||
allowed=False,
|
||||
reason=f"Weekly spending {spent_week + amount} would exceed limit {limits.per_week}",
|
||||
current_spent=current_spent,
|
||||
remaining=remaining,
|
||||
requires_time_lock=False
|
||||
)
|
||||
|
||||
# Check time lock requirement
|
||||
requires_time_lock = amount >= limits.time_lock_threshold
|
||||
time_lock_until = None
|
||||
|
||||
if requires_time_lock:
|
||||
time_lock_until = timestamp + timedelta(hours=limits.time_lock_delay_hours)
|
||||
|
||||
return SpendingCheckResult(
|
||||
allowed=True,
|
||||
reason="Spending limits check passed",
|
||||
current_spent=current_spent,
|
||||
remaining=remaining,
|
||||
requires_time_lock=requires_time_lock,
|
||||
time_lock_until=time_lock_until
|
||||
)
|
||||
|
||||
def update_spending_limits(self, agent_address: str, new_limits: Dict, guardian_address: str) -> bool:
|
||||
"""
|
||||
Update spending limits for an agent
|
||||
|
||||
Args:
|
||||
agent_address: Agent wallet address
|
||||
new_limits: New spending limits
|
||||
guardian_address: Guardian making the change
|
||||
|
||||
Returns:
|
||||
True if updated successfully
|
||||
"""
|
||||
agent_address = to_checksum_address(agent_address)
|
||||
guardian_address = to_checksum_address(guardian_address)
|
||||
|
||||
# Verify guardian authorization
|
||||
if not self.is_guardian_authorized(agent_address, guardian_address):
|
||||
return False
|
||||
|
||||
try:
|
||||
with self.get_session() as session:
|
||||
limits = session.query(SpendingLimit).filter(
|
||||
SpendingLimit.agent_address == agent_address
|
||||
).first()
|
||||
|
||||
if limits:
|
||||
limits.per_transaction = new_limits.get("per_transaction", limits.per_transaction)
|
||||
limits.per_hour = new_limits.get("per_hour", limits.per_hour)
|
||||
limits.per_day = new_limits.get("per_day", limits.per_day)
|
||||
limits.per_week = new_limits.get("per_week", limits.per_week)
|
||||
limits.time_lock_threshold = new_limits.get("time_lock_threshold", limits.time_lock_threshold)
|
||||
limits.time_lock_delay_hours = new_limits.get("time_lock_delay_hours", limits.time_lock_delay_hours)
|
||||
limits.updated_at = datetime.utcnow()
|
||||
limits.updated_by = guardian_address
|
||||
else:
|
||||
limits = SpendingLimit(
|
||||
agent_address=agent_address,
|
||||
per_transaction=new_limits.get("per_transaction", 1000.0),
|
||||
per_hour=new_limits.get("per_hour", 5000.0),
|
||||
per_day=new_limits.get("per_day", 20000.0),
|
||||
per_week=new_limits.get("per_week", 100000.0),
|
||||
time_lock_threshold=new_limits.get("time_lock_threshold", 5000.0),
|
||||
time_lock_delay_hours=new_limits.get("time_lock_delay_hours", 24),
|
||||
updated_at=datetime.utcnow(),
|
||||
updated_by=guardian_address
|
||||
)
|
||||
session.add(limits)
|
||||
|
||||
session.commit()
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"Failed to update spending limits: {e}")
|
||||
return False
|
||||
|
||||
def add_guardian(self, agent_address: str, guardian_address: str, added_by: str) -> bool:
|
||||
"""
|
||||
Add a guardian for an agent
|
||||
|
||||
Args:
|
||||
agent_address: Agent wallet address
|
||||
guardian_address: Guardian address
|
||||
added_by: Who added this guardian
|
||||
|
||||
Returns:
|
||||
True if added successfully
|
||||
"""
|
||||
agent_address = to_checksum_address(agent_address)
|
||||
guardian_address = to_checksum_address(guardian_address)
|
||||
added_by = to_checksum_address(added_by)
|
||||
|
||||
try:
|
||||
with self.get_session() as session:
|
||||
# Check if already exists
|
||||
existing = session.query(GuardianAuthorization).filter(
|
||||
GuardianAuthorization.agent_address == agent_address,
|
||||
GuardianAuthorization.guardian_address == guardian_address
|
||||
).first()
|
||||
|
||||
if existing:
|
||||
existing.is_active = True
|
||||
existing.added_at = datetime.utcnow()
|
||||
existing.added_by = added_by
|
||||
else:
|
||||
auth = GuardianAuthorization(
|
||||
id=f"{agent_address}_{guardian_address}",
|
||||
agent_address=agent_address,
|
||||
guardian_address=guardian_address,
|
||||
is_active=True,
|
||||
added_at=datetime.utcnow(),
|
||||
added_by=added_by
|
||||
)
|
||||
session.add(auth)
|
||||
|
||||
session.commit()
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"Failed to add guardian: {e}")
|
||||
return False
|
||||
|
||||
def is_guardian_authorized(self, agent_address: str, guardian_address: str) -> bool:
|
||||
"""
|
||||
Check if a guardian is authorized for an agent
|
||||
|
||||
Args:
|
||||
agent_address: Agent wallet address
|
||||
guardian_address: Guardian address
|
||||
|
||||
Returns:
|
||||
True if authorized
|
||||
"""
|
||||
agent_address = to_checksum_address(agent_address)
|
||||
guardian_address = to_checksum_address(guardian_address)
|
||||
|
||||
with self.get_session() as session:
|
||||
auth = session.query(GuardianAuthorization).filter(
|
||||
GuardianAuthorization.agent_address == agent_address,
|
||||
GuardianAuthorization.guardian_address == guardian_address,
|
||||
GuardianAuthorization.is_active == True
|
||||
).first()
|
||||
|
||||
return auth is not None
|
||||
|
||||
def get_spending_summary(self, agent_address: str) -> Dict:
|
||||
"""
|
||||
Get comprehensive spending summary for an agent
|
||||
|
||||
Args:
|
||||
agent_address: Agent wallet address
|
||||
|
||||
Returns:
|
||||
Spending summary
|
||||
"""
|
||||
agent_address = to_checksum_address(agent_address)
|
||||
now = datetime.utcnow()
|
||||
|
||||
# Get current spending
|
||||
current_spent = {
|
||||
"hour": self.get_spent_in_period(agent_address, "hour", now),
|
||||
"day": self.get_spent_in_period(agent_address, "day", now),
|
||||
"week": self.get_spent_in_period(agent_address, "week", now)
|
||||
}
|
||||
|
||||
# Get limits
|
||||
with self.get_session() as session:
|
||||
limits = session.query(SpendingLimit).filter(
|
||||
SpendingLimit.agent_address == agent_address
|
||||
).first()
|
||||
|
||||
if not limits:
|
||||
return {"error": "No spending limits set"}
|
||||
|
||||
# Calculate remaining
|
||||
remaining = {
|
||||
"hour": limits.per_hour - current_spent["hour"],
|
||||
"day": limits.per_day - current_spent["day"],
|
||||
"week": limits.per_week - current_spent["week"]
|
||||
}
|
||||
|
||||
# Get authorized guardians
|
||||
with self.get_session() as session:
|
||||
guardians = session.query(GuardianAuthorization).filter(
|
||||
GuardianAuthorization.agent_address == agent_address,
|
||||
GuardianAuthorization.is_active == True
|
||||
).all()
|
||||
|
||||
return {
|
||||
"agent_address": agent_address,
|
||||
"current_spending": current_spent,
|
||||
"remaining_spending": remaining,
|
||||
"limits": {
|
||||
"per_transaction": limits.per_transaction,
|
||||
"per_hour": limits.per_hour,
|
||||
"per_day": limits.per_day,
|
||||
"per_week": limits.per_week
|
||||
},
|
||||
"time_lock": {
|
||||
"threshold": limits.time_lock_threshold,
|
||||
"delay_hours": limits.time_lock_delay_hours
|
||||
},
|
||||
"authorized_guardians": [g.guardian_address for g in guardians],
|
||||
"last_updated": limits.updated_at.isoformat() if limits.updated_at else None
|
||||
}
|
||||
|
||||
|
||||
# Global persistent tracker instance
|
||||
persistent_tracker = PersistentSpendingTracker()
|
||||
Reference in New Issue
Block a user