#!/usr/bin/env python3
# -----------------------------------------------------------------------------
# spokepool_smt.py
#
# SMT proof (z3) of NO-THEFT-OF-PRINCIPAL for AereSpokePool (the corrected V2
# with the ROUND-3 `totalLocked` accounting fix), a NEGATIVE CONTROL showing the
# pre-V2 sweepResidual bug is a real theft vector, and one honestly-surfaced
# COMPOSITION FINDING (sweepResidual does not reserve solver bonds).
#
# Contract:  contracts/contracts/intents/AereSpokePool.sol
#
# We model the per-input-token accounting as an inductive transition system and
# prove the safety invariant is INDUCTIVE (base case + every operation preserves
# it).  Each Solidity `require`/revert is a guard; Solidity 0.8 checked
# arithmetic is modeled as "underflow reverts (no state change)".
#
#   Per input token T:
#     bal  = IERC20(T).balanceOf(address(this))     -- real token balance
#     lock = totalLocked[T]                          -- sum of unsettled solverPortions (USER PRINCIPAL)
#
#   PRINCIPAL SAFETY INVARIANT (the task's target -- no theft of principal):
#     INV_P := bal >= lock  AND  lock >= 0
#
#   INV_P => every open user order's locked principal is ALWAYS fully backed by
#   the contract balance: no settle / sweepResidual can drive balance below what
#   users are collectively owed, so settle() can always pay out.
#
# Method: for each op OP, check  INV_P(pre) AND guards AND post=OP(pre) AND
# NOT INV_P(post)  is UNSAT.  UNSAT => OP cannot break the invariant.
# Unbounded Ints = the accounting/design abstraction (matches solc SMTChecker's
# default int model). This checks the DESIGN math, NOT the EVM bytecode.
# -----------------------------------------------------------------------------
from z3 import Int, Solver, And, Or, Not, If, sat, unsat

FEE_BPS = 50          # PROTOCOL_FEE_BPS (immutable)
BPS_DEN = 10_000

def INV_P(bal, lock):
    return And(bal >= lock, lock >= 0)

results = []
def check(name, s, expect_unsat=True, kind="PROOF"):
    r = s.check()
    if expect_unsat:
        ok = (r == unsat); tag = "PROVED" if ok else "FAILED"
    else:
        ok = (r == sat);   tag = "CEX-FOUND" if ok else "FAILED"
    results.append((name, tag, ok, kind))
    print(f"[{tag:9}] ({kind}) {name}: z3={r} (expected {'unsat' if expect_unsat else 'sat'})")
    if r == sat:
        m = s.model()
        print("             witness:", {str(d): m[d].as_long() for d in m.decls()})
    return ok

print("### AereSpokePool  --  no-theft-of-PRINCIPAL   INV_P: bal >= totalLocked\n")

# ---- BASE CASE: fresh contract ------------------------------------------------
s = Solver(); s.add(Not(INV_P(0, 0)))
check("base case (empty contract) satisfies INV_P", s)

# ---- open(inputAmount): pull inputAmount, route fee, lock solverPortion --------
# solverPortion sp = inputAmount - fee. Balance rises by AT LEAST sp (fee either
# leaves to sink OR stays as residual >= 0). lock rises by sp.
s = Solver()
bal, lock, inp, feeStays = Int('bal'), Int('lock'), Int('inputAmount'), Int('feeStays')
fee = (inp * FEE_BPS) / BPS_DEN
sp  = inp - fee
s.add(INV_P(bal, lock), inp >= 1, Or(feeStays == 0, feeStays == fee))
s.add(Not(INV_P(bal + sp + feeStays, lock + sp)))
check("open() preserves INV_P", s)

# ---- settle(order a): pays a = o.inputAmount to solver, exactly once -----------
# The order's principal a was added to lock at open and is still counted (a<=lock).
# settle: lock -= a ; bal -= a  (checked-sub => a<=lock).
s = Solver()
bal, lock, a = Int('bal'), Int('lock'), Int('orderAmount')
s.add(INV_P(bal, lock), a >= 1, a <= lock)
s.add(Not(INV_P(bal - a, lock - a)))
check("settle() preserves INV_P (single payout)", s)

# ---- double-settle blocked: over-withdraw beyond locked amt is guard-infeasible-
s = Solver()
lock2, a = Int('lock2'), Int('a')
s.add(lock2 >= 0, a >= 1, a > lock2, a <= lock2)   # settled-flag/checked-sub guard
check("over-withdraw beyond locked amount is guard-infeasible", s)

# ---- sweepResidual(token) V2 (FIXED): reserved = max(lock, floor) -------------
s = Solver()
bal, lock, floor = Int('bal'), Int('lock'), Int('sweepFloor')
s.add(INV_P(bal, lock), floor >= 0)
reserved = If(lock > floor, lock, floor)           # max(lock,floor) >= lock
s.add(bal > reserved)                              # else revert ZeroAmount (no state change)
s.add(Not(INV_P(bal - (bal - reserved), lock)))    # bal2 == reserved
check("sweepResidual() V2 (max(lock,floor)) preserves INV_P", s)

# ---- NEGATIVE CONTROL: PRE-V2 sweepResidual reserved ONLY sweepFloor -----------
# Ignores totalLocked. If Foundation mis-sets floor < lock, an attacker sweeps
# bal-floor, draining locked user principal. z3 must find a THEFT state.
s = Solver()
bal, lock, floor = Int('bal'), Int('lock'), Int('sweepFloor')
s.add(INV_P(bal, lock), floor >= 0, bal > floor)   # buggy reserve = floor only
s.add(Not(INV_P(bal - (bal - floor), lock)))       # bal2 == floor, can be < lock
check("PRE-V2 sweepResidual (floor-only) CAN steal principal", s,
      expect_unsat=False, kind="NEG-CTRL")

# ---- NEGATIVE CONTROL: the OLD (pre-fix) composition WAS a real under-backing ---
# sweepResidual reserved totalLocked but NOT solver bonds. With WAERE as BOTH a
# bridge INPUT token and the bond token: bal_W = principal(lock) + bonds; the old
# sweep reserved only lock -> swept the bonds; a later slashSolverBond then drove
# bal_W = lock - amt < lock. z3 finds this theft state (proves the bug was real).
s = Solver()
lock, bonds, amt = Int('lock'), Int('bonds'), Int('slashAmt')
s.add(lock >= 1, bonds >= 1, amt >= 1, amt <= bonds)
bal_after_sweep_OLD = If(lock > 0, lock, 0)            # OLD sweep: bonds gone
s.add(Not(INV_P(bal_after_sweep_OLD - amt, lock)))     # principal under-backed
check("PRE-FIX composition (sweep ignores bonds)+slash CAN under-back principal", s,
      expect_unsat=False, kind="NEG-CTRL")

# ---- BOND-ACCOUNTING FIX: sweepResidual reserves totalBond ----------------------
# Fix: track totalBond[T], reserve max(lock+bond, floor) in sweepResidual, and
# keep it in lockstep (depositSolverBond +=, slashSolverBond -=). Prove the
# STRONGER invariant INV_PB := bal >= lock + bond is INDUCTIVE (subsumes INV_P and
# closes the finding). Each op: INV_PB(pre) AND guards AND post => INV_PB(post).
def INV_PB(bal, lock, bond):
    return And(bal >= lock + bond, lock >= 0, bond >= 0)

s = Solver(); bal, lock, bond, amt = Int('bal'), Int('lock'), Int('bond'), Int('amt')
s.add(INV_PB(bal, lock, bond), amt >= 1); s.add(Not(INV_PB(bal + amt, lock, bond + amt)))
check("FIX depositSolverBond preserves INV_PB (bal>=lock+bond)", s)

s = Solver(); bal, lock, bond, amt = Int('bal'), Int('lock'), Int('bond'), Int('amt')
s.add(INV_PB(bal, lock, bond), amt >= 1, amt <= bond); s.add(Not(INV_PB(bal - amt, lock, bond - amt)))
check("FIX slashSolverBond (bond-=amt, bal-=amt) preserves INV_PB", s)

s = Solver(); bal, lock, bond, floor = Int('bal'), Int('lock'), Int('bond'), Int('floor')
s.add(INV_PB(bal, lock, bond), floor >= 0)
reservedF = If(lock + bond > floor, lock + bond, floor)
s.add(bal > reservedF); s.add(Not(INV_PB(bal - (bal - reservedF), lock, bond)))
check("FIX sweepResidual (reserve max(lock+bond,floor)) preserves INV_PB", s)

s = Solver(); bal, lock, bond, inp, feeStays = Int('bal'), Int('lock'), Int('bond'), Int('inputAmount'), Int('feeStays')
feeO = (inp * FEE_BPS) / BPS_DEN; spO = inp - feeO
s.add(INV_PB(bal, lock, bond), inp >= 1, Or(feeStays == 0, feeStays == feeO))
s.add(Not(INV_PB(bal + spO + feeStays, lock + spO, bond)))
check("FIX open preserves INV_PB (bond unchanged)", s)

s = Solver(); bal, lock, bond, a = Int('bal'), Int('lock'), Int('bond'), Int('a')
s.add(INV_PB(bal, lock, bond), a >= 1, a <= lock); s.add(Not(INV_PB(bal - a, lock - a, bond)))
check("FIX settle preserves INV_PB (bond unchanged)", s)

# ---- REFUND FIX (adversarial-review-found HIGH, 2026-07-12): cancel(order a) -----
# returns the original user's principal a after fillDeadline when the order is
# unclaimed. Same accounting delta as settle: lock -= a ; bal -= a (Solidity
# checked-sub => a <= lock; the o.settled flag blocks double-cancel exactly as it
# blocks double-settle). Prove it preserves the stronger invariant INV_PB, i.e. the
# refund path cannot under-back any other user's principal or a solver bond.
s = Solver(); bal, lock, bond, a = Int('bal'), Int('lock'), Int('bond'), Int('a')
s.add(INV_PB(bal, lock, bond), a >= 1, a <= lock); s.add(Not(INV_PB(bal - a, lock - a, bond)))
check("FIX cancel/refund preserves INV_PB (lock-=a, bal-=a, bond unchanged)", s)

# The exact WAERE-as-input composition is now SAFE: sweep reserves lock+bond (sweeps
# nothing), then slash bond-=amt, bal-=amt leaves bal = lock+bond-amt >= lock.
s = Solver(); lock, bonds, amt = Int('lock'), Int('bonds'), Int('slashAmt')
s.add(lock >= 1, bonds >= 1, amt >= 1, amt <= bonds)
bal_fixed = (lock + bonds) - amt                       # FIX: sweep took nothing, slash -amt
s.add(Not(INV_P(bal_fixed, lock)))                     # want bal>=lock -> negation UNSAT
check("FIXED: reserve-bond sweep + slash keeps WAERE-input principal backed", s)

# ------------------------------------------------------------------------------
print("\n=== SUMMARY ===")
allok = True
for name, tag, ok, kind in results:
    print(f"  {tag:9} [{kind}]  {name}")
    allok = allok and ok
print()
print("NO-THEFT-OF-PRINCIPAL (INV_P: bal >= totalLocked) for corrected V2:",
      "\n  PROVED inductive -- base case + open/settle/sweepResidual-V2 all preserve it;"
      "\n  pre-V2 floor-only sweep reproduces the theft (control)." if allok else
      "\n  NOT fully established (see FAILED).")
print("BOND-ACCOUNTING FIX PROVED: totalBond[T] is now reserved by sweepResidual"
      "\n  (reserve = max(totalLocked+totalBond, floor)) and kept in lockstep by"
      "\n  depositSolverBond/slashSolverBond. The stronger invariant bal>=lock+bond is"
      "\n  inductive, and the exact WAERE-as-input composition is now UNSAT (safe)."
      "\n  The pre-fix composition is retained above as a NEG-CTRL (it WAS a real bug).")
import sys
sys.exit(0 if allok else 1)