Agent Wallet Security & Spending Limits 2026: Complete Protection Guide

from turnkey_client import TurnkeyClient
from eth_account.messages import encode_defunct

# Initialize Turnkey client with API credentials
client = TurnkeyClient(
    api_public_key="YOUR_API_PUBLIC_KEY",
    api_private_key="YOUR_API_PRIVATE_KEY",
    organization_id="YOUR_ORG_ID",
    base_url="https://api.turnkey.com"
)

# Sign a transaction using MPC — key never leaves Turnkey infrastructure
def sign_with_mpc(wallet_id: str, transaction: dict) -> str:
    """
    Sign a transaction using Turnkey MPC.
    The private key is split across Turnkey's infrastructure.
    No single party ever holds the complete key.
    """
    # Turnkey validates the transaction against your policy engine
    # before producing the MPC signature
    result = client.sign_transaction(
        wallet_id=wallet_id,
        transaction=transaction,
        # Policy engine checks happen server-side:
        # - Per-transaction limit
        # - Daily spending cap
        # - Contract allowlist
    )
    return result.signed_transaction

# Create a policy that limits agent spending
policy = client.create_policy(
    name="prediction-market-agent",
    rules=[
        {
            "type": "max_transaction_amount",
            "amount": "100",       # $100 max per trade
            "currency": "USDC"
        },
        {
            "type": "daily_spending_limit",
            "amount": "1000",      # $1,000 daily cap
            "currency": "USDC"
        },
        {
            "type": "contract_allowlist",
            "addresses": [
                "0x4bFb41d5B3570DeFd03C39a9A4D8dE6Bd8B8982E",  # Polymarket CTF Exchange
                "0x2791Bca1f2de4661ED88A30C99A7a9449Aa84174",  # USDC on Polygon
            ]
        }
    ]
)

from safe_eth.safe import Safe
from safe_eth.safe.safe_tx import SafeTx
from web3 import Web3

w3 = Web3(Web3.HTTPProvider("https://polygon-rpc.com"))

# Load existing Safe
safe = Safe(
    address="0xYOUR_SAFE_ADDRESS",
    ethereum_client=w3
)

# Transaction guard contract enforces constraints on-chain
# Deploy this guard to restrict what the agent can do
TRANSACTION_GUARD_ABI = [
    # checkTransaction is called before every Safe transaction
    # Revert here = transaction is blocked at the contract level
    {
        "name": "checkTransaction",
        "inputs": [
            {"name": "to", "type": "address"},
            {"name": "value", "type": "uint256"},
            {"name": "data", "type": "bytes"},
            # ... additional parameters
        ]
    }
]

# Guard configuration (set via Safe's setGuard function)
GUARD_CONFIG = {
    "max_value_per_tx": Web3.to_wei(100, "ether"),   # 100 USDC max
    "daily_limit": Web3.to_wei(1000, "ether"),        # 1,000 USDC daily
    "allowed_targets": [
        "0x4bFb41d5B3570DeFd03C39a9A4D8dE6Bd8B8982E",  # Polymarket CTF
    ],
    "cooldown_seconds": 60,  # Minimum 60s between transactions
}

# Agent proposes a trade — needs co-signer approval
def propose_trade(safe: Safe, to: str, value: int, data: bytes):
    """
    Agent proposes a transaction to the Safe.
    The transaction guard validates constraints on-chain.
    Co-signer must approve before execution.
    """
    safe_tx = safe.build_multisig_tx(
        to=to,
        value=value,
        data=data,
        safe_nonce=safe.retrieve_nonce()
    )

    # Agent signs with its key
    safe_tx.sign(agent_private_key)

    # Submit to Safe Transaction Service for co-signer approval
    # Co-signer (backend or human) reviews and co-signs
    safe.post_transaction(safe_tx)

    return safe_tx.safe_tx_hash

from eth_account import Account
from web3 import Web3
import time
import json

# Session key configuration for a Polymarket trading agent
SESSION_KEY_CONFIG = {
    # Scoped permissions
    "allowed_contracts": [
        "0x4bFb41d5B3570DeFd03C39a9A4D8dE6Bd8B8982E",  # Polymarket CTF Exchange
        "0x2791Bca1f2de4661ED88A30C99A7a9449Aa84174",  # USDC on Polygon
    ],
    "allowed_functions": [
        "0x12345678",  # placeTrade function selector
        "0xabcdef01",  # cancelOrder function selector
        "0x095ea7b3",  # approve (USDC approval for CTF Exchange only)
    ],
    # Spending limits
    "max_spend_per_tx": Web3.to_wei(100, "mwei"),   # 100 USDC (6 decimals)
    "max_spend_session": Web3.to_wei(500, "mwei"),   # 500 USDC total
    # Time bounds
    "valid_after": int(time.time()),                  # Starts now
    "valid_until": int(time.time()) + (4 * 3600),     # Expires in 4 hours
}

def create_session_key(master_wallet, config: dict) -> dict:
    """
    Create a scoped session key for the agent.
    The smart contract wallet validates every UserOperation
    against these constraints before executing.
    """
    # Generate a fresh ephemeral key pair
    session_account = Account.create()

    # Register the session key with the smart contract wallet
    # This is a simplified representation — actual implementation
    # depends on your ERC-4337 wallet provider
    session_data = {
        "session_key": session_account.address,
        "permissions": {
            "allowed_contracts": config["allowed_contracts"],
            "allowed_functions": config["allowed_functions"],
            "max_spend_per_tx": config["max_spend_per_tx"],
            "max_spend_session": config["max_spend_session"],
            "valid_after": config["valid_after"],
            "valid_until": config["valid_until"],
        }
    }

    # The master key signs the session key registration
    # After this, the master key goes back to cold storage
    registration_tx = master_wallet.register_session_key(session_data)

    return {
        "session_private_key": session_account.key.hex(),
        "session_address": session_account.address,
        "permissions": session_data["permissions"],
        "expires_at": config["valid_until"],
    }

def validate_against_session(tx: dict, session: dict) -> bool:
    """
    Check if a transaction is within session key permissions.
    The smart contract does this on-chain, but we pre-validate
    in the agent to fail fast and avoid wasting gas.
    """
    perms = session["permissions"]

    if tx["to"] not in perms["allowed_contracts"]:
        raise ValueError(f"Contract {tx['to']} not in session allowlist")

    if tx["data"][:10] not in [f"0x{s}" for s in perms["allowed_functions"]]:
        raise ValueError(f"Function not permitted in this session")

    if tx.get("value", 0) > perms["max_spend_per_tx"]:
        raise ValueError(f"Transaction exceeds per-tx limit")

    if time.time() > perms["valid_until"]:
        raise ValueError("Session key expired")

    return True

from dataclasses import dataclass
from decimal import Decimal

@dataclass
class TransactionLimits:
    max_per_trade: Decimal       # Maximum USDC per single trade
    min_per_trade: Decimal       # Minimum to prevent dust attacks

def validate_trade_size(amount: Decimal, limits: TransactionLimits) -> bool:
    """
    Reject any trade above the per-transaction maximum.
    Called before every trade execution.
    """
    if amount > limits.max_per_trade:
        raise ValueError(
            f"Trade size ${amount} exceeds per-transaction limit "
            f"of ${limits.max_per_trade}"
        )
    if amount < limits.min_per_trade:
        raise ValueError(
            f"Trade size ${amount} below minimum of ${limits.min_per_trade}"
        )
    return True

# Usage
limits = TransactionLimits(
    max_per_trade=Decimal("100.00"),   # $100 max per trade
    min_per_trade=Decimal("1.00"),     # $1 minimum
)
validate_trade_size(Decimal("50.00"), limits)  # OK
validate_trade_size(Decimal("200.00"), limits) # Raises ValueError

from decimal import Decimal
from datetime import datetime, timezone

class SessionTracker:
    def __init__(self, session_cap: Decimal):
        self.session_cap = session_cap
        self.total_spent = Decimal("0.00")
        self.trade_count = 0
        self.started_at = datetime.now(timezone.utc)

    def record_and_validate(self, amount: Decimal) -> bool:
        """
        Check if this trade would exceed the session cap.
        Must be called before signing.
        """
        projected = self.total_spent + amount
        if projected > self.session_cap:
            raise ValueError(
                f"Session cap exceeded: ${self.total_spent} spent + "
                f"${amount} proposed = ${projected} > "
                f"cap of ${self.session_cap}"
            )
        self.total_spent = projected
        self.trade_count += 1
        return True

    @property
    def remaining(self) -> Decimal:
        return self.session_cap - self.total_spent

# Usage
session = SessionTracker(session_cap=Decimal("500.00"))
session.record_and_validate(Decimal("100.00"))  # OK, $400 remaining
session.record_and_validate(Decimal("350.00"))  # OK, $50 remaining
session.record_and_validate(Decimal("100.00"))  # Raises ValueError

import time
from decimal import Decimal
from collections import deque
from dataclasses import dataclass

@dataclass
class TradeRecord:
    amount: Decimal
    timestamp: float

class RollingLimitTracker:
    def __init__(self, daily_limit: Decimal, weekly_limit: Decimal):
        self.daily_limit = daily_limit
        self.weekly_limit = weekly_limit
        self.trades: deque[TradeRecord] = deque()

    def _prune_old_trades(self):
        """Remove trades older than 7 days."""
        cutoff = time.time() - (7 * 24 * 3600)
        while self.trades and self.trades[0].timestamp < cutoff:
            self.trades.popleft()

    def _sum_window(self, window_seconds: int) -> Decimal:
        """Sum trade amounts within a time window."""
        cutoff = time.time() - window_seconds
        return sum(
            (t.amount for t in self.trades if t.timestamp >= cutoff),
            Decimal("0.00")
        )

    def validate(self, amount: Decimal) -> bool:
        """Validate against daily and weekly rolling limits."""
        self._prune_old_trades()

        daily_spent = self._sum_window(24 * 3600)
        if daily_spent + amount > self.daily_limit:
            raise ValueError(
                f"Daily limit: ${daily_spent} spent in last 24h + "
                f"${amount} = ${daily_spent + amount} > "
                f"limit of ${self.daily_limit}"
            )

        weekly_spent = self._sum_window(7 * 24 * 3600)
        if weekly_spent + amount > self.weekly_limit:
            raise ValueError(
                f"Weekly limit: ${weekly_spent} spent in last 7d + "
                f"${amount} = ${weekly_spent + amount} > "
                f"limit of ${self.weekly_limit}"
            )

        self.trades.append(TradeRecord(amount=amount, timestamp=time.time()))
        return True

from web3 import Web3

class ContractAllowlist:
    """
    Only permit transactions to known, verified contracts.
    Prevents the agent from interacting with malicious contracts
    even if prompt injection directs it to do so.
    """

    POLYMARKET_CONTRACTS = {
        "0x4bFb41d5B3570DeFd03C39a9A4D8dE6Bd8B8982E": "CTF Exchange",
        "0x2791Bca1f2de4661ED88A30C99A7a9449Aa84174": "USDC (Polygon)",
        "0x4D97DCd97eC945f40cF65F87097ACe5EA0476045": "Neg Risk CTF Exchange",
        "0xd91E80cF2E7be2e162c6513ceD06f1dD0dA35296": "Neg Risk Adapter",
    }

    def __init__(self, additional_contracts: dict[str, str] | None = None):
        self.allowed = dict(self.POLYMARKET_CONTRACTS)
        if additional_contracts:
            self.allowed.update(additional_contracts)

    def validate(self, target_address: str) -> bool:
        """
        Reject transactions to any address not in the allowlist.
        This is a hard block — no override, no exceptions.
        """
        checksummed = Web3.to_checksum_address(target_address)
        if checksummed not in self.allowed:
            raise ValueError(
                f"Contract {checksummed} is not in the allowlist. "
                f"Allowed contracts: {list(self.allowed.values())}"
            )
        return True

import time
import json
import logging
from pathlib import Path

logger = logging.getLogger(__name__)

class KillSwitch:
    """
    Emergency halt mechanism for agent wallets.
    Supports multiple check sources for redundancy.
    """

    def __init__(
        self,
        local_flag_path: str = "/tmp/agent_kill_switch.json",
        redis_client=None,
        redis_key: str = "agent:kill_switch",
        heartbeat_timeout: int = 300,  # 5 minutes
    ):
        self.local_flag_path = Path(local_flag_path)
        self.redis_client = redis_client
        self.redis_key = redis_key
        self.heartbeat_timeout = heartbeat_timeout
        self.last_heartbeat = time.time()

    def is_killed(self) -> bool:
        """
        Check if the kill switch is active.
        ANY active signal means HALT.
        """
        # Check 1: Local file flag
        if self.local_flag_path.exists():
            try:
                data = json.loads(self.local_flag_path.read_text())
                if data.get("killed", False):
                    logger.critical(
                        f"KILL SWITCH ACTIVE (local): {data.get('reason', 'unknown')}"
                    )
                    return True
            except (json.JSONDecodeError, IOError):
                # If we can't read the flag file, assume killed (fail safe)
                logger.critical("KILL SWITCH: Cannot read flag file — halting")
                return True

        # Check 2: Redis flag (if configured)
        if self.redis_client:
            try:
                value = self.redis_client.get(self.redis_key)
                if value and json.loads(value).get("killed", False):
                    logger.critical("KILL SWITCH ACTIVE (Redis)")
                    return True
            except Exception as e:
                # Redis unreachable — fail safe
                logger.critical(f"KILL SWITCH: Redis unreachable ({e}) — halting")
                return True

        # Check 3: Dead-man's switch (heartbeat timeout)
        if time.time() - self.last_heartbeat > self.heartbeat_timeout:
            logger.critical(
                f"KILL SWITCH: No heartbeat for {self.heartbeat_timeout}s — halting"
            )
            return True

        return False

    def heartbeat(self):
        """Agent must call this regularly to prove it's operating normally."""
        self.last_heartbeat = time.time()

    def trigger(self, reason: str):
        """Manually trigger the kill switch."""
        data = {"killed": True, "reason": reason, "triggered_at": time.time()}
        self.local_flag_path.write_text(json.dumps(data))
        if self.redis_client:
            self.redis_client.set(self.redis_key, json.dumps(data))
        logger.critical(f"KILL SWITCH TRIGGERED: {reason}")

import time
import logging
from decimal import Decimal
from dataclasses import dataclass, field
from collections import deque

logger = logging.getLogger(__name__)

@dataclass
class TradeRecord:
    amount: Decimal
    timestamp: float
    market_id: str
    tx_hash: str = ""

@dataclass
class SpendingGuardConfig:
    max_per_trade: Decimal = Decimal("100.00")
    min_per_trade: Decimal = Decimal("1.00")
    session_cap: Decimal = Decimal("500.00")
    daily_limit: Decimal = Decimal("1000.00")
    weekly_limit: Decimal = Decimal("5000.00")
    allowed_contracts: dict = field(default_factory=lambda: {
        "0x4bFb41d5B3570DeFd03C39a9A4D8dE6Bd8B8982E": "CTF Exchange",
        "0x2791Bca1f2de4661ED88A30C99A7a9449Aa84174": "USDC (Polygon)",
    })
    max_open_positions: int = 10
    cooldown_seconds: int = 5

class SpendingGuard:
    """
    Unified spending controls for prediction market agents.
    Every trade must pass through validate() before signing.

    Implements five layers of protection:
    1. Per-transaction size limits
    2. Session spending cap
    3. Daily/weekly rolling limits
    4. Contract allowlist
    5. Emergency kill switch

    Usage:
        guard = SpendingGuard(config)
        try:
            guard.validate(
                amount=Decimal("50.00"),
                target_contract="0x4bFb41d5B3570DeFd03C39a9A4D8dE6Bd8B8982E",
                market_id="polymarket-election-2026"
            )
            # Trade is approved — proceed with signing
        except SpendingViolation as e:
            # Trade is blocked — log and skip
            logger.warning(f"Trade blocked: {e}")
    """

    def __init__(self, config: SpendingGuardConfig, kill_switch=None):
        self.config = config
        self.kill_switch = kill_switch
        self.session_spent = Decimal("0.00")
        self.session_trade_count = 0
        self.trades: deque[TradeRecord] = deque()
        self.last_trade_time = 0.0
        self.session_start = time.time()

    def validate(
        self,
        amount: Decimal,
        target_contract: str,
        market_id: str = "",
    ) -> bool:
        """
        Run all spending validations. Raises SpendingViolation
        if any check fails. Returns True if all checks pass.
        """
        # 0. Kill switch — checked first, overrides everything
        if self.kill_switch and self.kill_switch.is_killed():
            raise SpendingViolation("Kill switch is active — all trading halted")

        # 1. Contract allowlist
        from web3 import Web3
        checksummed = Web3.to_checksum_address(target_contract)
        if checksummed not in self.config.allowed_contracts:
            raise SpendingViolation(
                f"Contract {checksummed} not in allowlist"
            )

        # 2. Per-transaction limits
        if amount > self.config.max_per_trade:
            raise SpendingViolation(
                f"Trade ${amount} exceeds per-tx limit ${self.config.max_per_trade}"
            )
        if amount < self.config.min_per_trade:
            raise SpendingViolation(
                f"Trade ${amount} below minimum ${self.config.min_per_trade}"
            )

        # 3. Session cap
        if self.session_spent + amount > self.config.session_cap:
            raise SpendingViolation(
                f"Session cap: ${self.session_spent} + ${amount} "
                f"> ${self.config.session_cap}"
            )

        # 4. Daily rolling limit
        daily_spent = self._sum_window(24 * 3600)
        if daily_spent + amount > self.config.daily_limit:
            raise SpendingViolation(
                f"Daily limit: ${daily_spent} + ${amount} "
                f"> ${self.config.daily_limit}"
            )

        # 5. Weekly rolling limit
        weekly_spent = self._sum_window(7 * 24 * 3600)
        if weekly_spent + amount > self.config.weekly_limit:
            raise SpendingViolation(
                f"Weekly limit: ${weekly_spent} + ${amount} "
                f"> ${self.config.weekly_limit}"
            )

        # 6. Cooldown between trades
        elapsed = time.time() - self.last_trade_time
        if elapsed < self.config.cooldown_seconds and self.last_trade_time > 0:
            raise SpendingViolation(
                f"Cooldown: {self.config.cooldown_seconds - elapsed:.1f}s remaining"
            )

        # All checks passed — record the trade
        self.session_spent += amount
        self.session_trade_count += 1
        self.last_trade_time = time.time()
        self.trades.append(TradeRecord(
            amount=amount,
            timestamp=time.time(),
            market_id=market_id,
        ))
        self._prune_old_trades()

        logger.info(
            f"Trade approved: ${amount} on {market_id} | "
            f"Session: ${self.session_spent}/${self.config.session_cap} | "
            f"Daily: ${daily_spent + amount}/${self.config.daily_limit}"
        )
        return True

    def _sum_window(self, window_seconds: int) -> Decimal:
        cutoff = time.time() - window_seconds
        return sum(
            (t.amount for t in self.trades if t.timestamp >= cutoff),
            Decimal("0.00")
        )

    def _prune_old_trades(self):
        cutoff = time.time() - (7 * 24 * 3600)
        while self.trades and self.trades[0].timestamp < cutoff:
            self.trades.popleft()

    def get_status(self) -> dict:
        """Return current spending status for monitoring dashboards."""
        return {
            "session_spent": str(self.session_spent),
            "session_remaining": str(self.config.session_cap - self.session_spent),
            "session_trade_count": self.session_trade_count,
            "daily_spent": str(self._sum_window(24 * 3600)),
            "weekly_spent": str(self._sum_window(7 * 24 * 3600)),
            "kill_switch_active": (
                self.kill_switch.is_killed() if self.kill_switch else False
            ),
            "session_uptime_seconds": time.time() - self.session_start,
        }


class SpendingViolation(Exception):
    """Raised when a trade violates spending controls."""
    pass

from decimal import Decimal

# Initialize spending guard with conservative limits
config = SpendingGuardConfig(
    max_per_trade=Decimal("100.00"),     # 2% of $5K bankroll
    session_cap=Decimal("500.00"),       # 10% per session
    daily_limit=Decimal("1000.00"),      # 20% per day
    weekly_limit=Decimal("2500.00"),     # 50% per week
    cooldown_seconds=10,                 # 10s between trades
)

kill_switch = KillSwitch(redis_client=redis_conn)
guard = SpendingGuard(config=config, kill_switch=kill_switch)

# Trading loop
for signal in trading_signals:
    try:
        # Validate before signing
        guard.validate(
            amount=Decimal(str(signal.size)),
            target_contract=signal.contract_address,
            market_id=signal.market_id,
        )
        # All checks passed — execute trade
        execute_trade(signal)
        # Send heartbeat to keep dead-man's switch happy
        kill_switch.heartbeat()

    except SpendingViolation as e:
        logger.warning(f"Trade blocked by SpendingGuard: {e}")
        # Do NOT execute — move to next signal

    except Exception as e:
        logger.error(f"Unexpected error: {e}")
        # Trigger kill switch on unexpected errors
        kill_switch.trigger(reason=f"Unexpected error: {e}")
        break