import sys
import re
from z3 import *

IPTABLES_FILE = "./service-target/iptables.save"


def solve_circuit(iptables_path: str) -> list[int]:
    rules_raw = []
    with open(iptables_path) as f:
        for line in f:
            line = line.strip()
            if line.startswith("-A INPUT"):
                rules_raw.append(line)

    input_knock: dict[str, int] = {}   # ipset_bit -> source_port
    eval_del: list[str] = []           # bits cleared unconditionally on sport=1338
    eval_gate: list[str] = []          # conditional gate rules on sport=1338

    for rule in rules_raw:
        # Input knock: --sport 2XXXX without any match-set condition
        m = re.search(r"--sport (2\d+) -j SET --add-set (i\d+) src", rule)
        if m and "match-set" not in rule:
            input_knock[m.group(2)] = int(m.group(1))
            continue
        # Eval del: unconditional clear on sport=1338
        if "--sport 1338" in rule and "--del-set" in rule and "match-set" not in rule:
            m = re.search(r"--del-set (i\d+)", rule)
            if m:
                eval_del.append(m.group(1))
            continue
        # Gate rule: conditional add on sport=1338
        if "--sport 1338" in rule and "--add-set" in rule:
            eval_gate.append(rule)

    print(f"[*] Circuit: {len(input_knock)} inputs, {len(eval_gate)} gates, {len(eval_del)} resets")

    # Build symbolic Z3 state
    bits: dict[str, object] = {}

    # Input bits = free boolean variables
    for bit in input_knock:
        bits[bit] = Bool(bit)

    # Cleared bits start as False (reset each eval cycle)
    for bit in eval_del:
        bits[bit] = BoolVal(False)

    def get_bit(b):
        return bits.get(b, BoolVal(False))

    # Process gate rules in file order (sequential iptables evaluation)
    for rule in eval_gate:
        # Double-condition: AND gate
        m2 = re.search(
            r"-m set (!?) ?--match-set (\S+) src -m set (!?) ?--match-set (\S+) src -j SET --add-set (\S+) src",
            rule,
        )
        if m2:
            neg1, a, neg2, b, out = m2.group(1) == "!", m2.group(2), m2.group(3) == "!", m2.group(4), m2.group(5)
            expr_a = Not(get_bit(a)) if neg1 else get_bit(a)
            expr_b = Not(get_bit(b)) if neg2 else get_bit(b)
            bits[out] = Or(get_bit(out), And(expr_a, expr_b))
            continue
        # Single-condition: OR feed
        m1 = re.search(
            r"-m set (!?) ?--match-set (\S+) src -j SET --add-set (\S+) src",
            rule,
        )
        if m1:
            neg, a, out = m1.group(1) == "!", m1.group(2), m1.group(3)
            expr_a = Not(get_bit(a)) if neg else get_bit(a)
            bits[out] = Or(get_bit(out), expr_a)

    # Solve: i00002 must be True (the ACCEPT gate)
    print("[*] Running Z3 solver...")
    s = Solver()
    s.add(get_bit("i00002") == True)

    if s.check() != sat:
        print("[!] UNSAT — circuit has no solution")
        sys.exit(1)

    model = s.model()
    knock_ports = []
    for bit, port in input_knock.items():
        if is_true(model.eval(Bool(bit))):
            knock_ports.append(port)

    knock_ports.sort()
    print(f"[+] Solution found: {len(knock_ports)} ports to knock")
    return knock_ports


if __name__ == "__main__":
    knock_ports = solve_circuit(IPTABLES_FILE)
    print(f"[*] Knock ports: {knock_ports}\n")