#!/usr/bin/env python3
from __future__ import annotations

import argparse
import struct
from dataclasses import dataclass
from pathlib import Path

from capstone import Cs, CS_ARCH_X86, CS_MODE_64
from capstone.x86 import X86_OP_IMM, X86_OP_MEM, X86_OP_REG
from pysat.solvers import Solver


CODE = Path("myfavoriteinstructions").read_bytes()
TEXT_START = 0x11A0
INPUT = 0x100000
STACK = 0x800000
MASK64 = (1 << 64) - 1

MD = Cs(CS_ARCH_X86, CS_MODE_64)
MD.detail = True
INSNS = {ins.address: ins for ins in MD.disasm(CODE[TEXT_START:0x4D990], TEXT_START)}

REG64 = {}
for names, full in [
    (["rax", "eax", "ax", "al", "ah"], "rax"),
    (["rbx", "ebx", "bx", "bl", "bh"], "rbx"),
    (["rcx", "ecx", "cx", "cl", "ch"], "rcx"),
    (["rdx", "edx", "dx", "dl", "dh"], "rdx"),
    (["rsi", "esi", "si", "sil"], "rsi"),
    (["rdi", "edi", "di", "dil"], "rdi"),
    (["rbp", "ebp", "bp", "bpl"], "rbp"),
    (["rsp", "esp", "sp", "spl"], "rsp"),
]:
    for name in names:
        REG64[name] = full
for i in range(8, 16):
    for suffix in ("", "d", "w", "b"):
        REG64[f"r{i}{suffix}"] = f"r{i}"


def trits_to_bytes_prefix(trits: list[int]) -> bytes:
    out = bytearray()
    pos = 0
    for _ in range(4):
        x = 0
        p = 1
        for t in trits[pos : pos + 5]:
            x += t * p
            p *= 3
        out.append(x & 0xff)
        pos += 5
    for ntrits, size in [(82, 16), (82, 16), (41, 8), (41, 8), (82, 16)]:
        x = 0
        p = 1
        for t in trits[pos : pos + ntrits]:
            x += t * p
            p *= 3
        out += x.to_bytes(size, "little")
        pos += ntrits
    return bytes(out)


def flag_to_trits(s: bytes) -> list[int]:
    assert len(s) >= 68
    out: list[int] = []
    for b in s[:4]:
        x = b
        for _ in range(5):
            out.append(x % 3)
            x //= 3
    for off, ntrits in [(4, 82), (20, 82), (36, 41), (44, 41), (52, 82)]:
        size = 16 if ntrits == 82 else 8
        x = int.from_bytes(s[off : off + size], "little")
        for _ in range(ntrits):
            out.append(x % 3)
            x //= 3
    assert len(out) == 348
    return out


def canon(name: str) -> str:
    return REG64.get(name, name)


@dataclass(frozen=True)
class TVal:
    const: int | None = None
    lits: tuple[int, int, int] | None = None

    def is_const(self) -> bool:
        return self.const is not None


class CNF:
    def __init__(self):
        self.next_var = 1
        self.clauses: list[list[int]] = []
        self.cache: dict[tuple, TVal] = {}

    def new_lit(self) -> int:
        lit = self.next_var
        self.next_var += 1
        return lit

    def add(self, clause: list[int]) -> None:
        self.clauses.append(clause)

    def new_val(self) -> TVal:
        lits = (self.new_lit(), self.new_lit(), self.new_lit())
        self.add([*lits])
        self.add([-lits[0], -lits[1]])
        self.add([-lits[0], -lits[2]])
        self.add([-lits[1], -lits[2]])
        return TVal(lits=lits)

    def const(self, value: int) -> TVal | int:
        if 0 <= value <= 2:
            return TVal(const=value)
        return value

    def force(self, val: TVal, value: int) -> None:
        if val.is_const():
            if val.const != value:
                self.add([])
            return
        assert val.lits is not None
        self.add([val.lits[value]])

    def gate2(self, op: str, a: TVal, b: TVal, table) -> TVal:
        if a.is_const() and b.is_const():
            return TVal(const=table[a.const][b.const])
        key = (op, a, b)
        if key in self.cache:
            return self.cache[key]
        out = self.new_val()
        assert out.lits is not None
        avals = [a.const] if a.is_const() else [0, 1, 2]
        bvals = [b.const] if b.is_const() else [0, 1, 2]
        for av in avals:
            for bv in bvals:
                cond = []
                if not a.is_const():
                    assert a.lits is not None
                    cond.append(-a.lits[av])
                if not b.is_const():
                    assert b.lits is not None
                    cond.append(-b.lits[bv])
                self.add([*cond, out.lits[table[av][bv]]])
        self.cache[key] = out
        return out

    def bsr(self, old: TVal | int, src: TVal | int) -> TVal | int:
        if isinstance(src, int):
            if src == 0:
                return old
            return src.bit_length() - 1
        if isinstance(old, int):
            if 0 <= old <= 2:
                old = TVal(const=old)
            else:
                raise NotImplementedError(("nontrit old bsr", old, src))
        table = [[oldv, 0, 1] for oldv in range(3)]
        return self.gate2("bsr", old, src, table)

    def bzhi(self, src: TVal | int, idx: TVal | int) -> TVal | int:
        if isinstance(idx, int):
            if idx == 0:
                return 0
            if idx == 1:
                if isinstance(src, int):
                    return src & 1
                return self.gate2("mod2", src, TVal(const=0), [[0], [1], [0]])  # unreachable
            return src
        if isinstance(src, int):
            if 0 <= src <= 2:
                src = TVal(const=src)
            else:
                raise NotImplementedError(("nontrit src bzhi", src, idx))
        table = [
            [0, 0, 0],
            [0, 1, 1],
            [0, 0, 2],
        ]
        return self.gate2("bzhi", src, idx, table)


def is_conc(x) -> bool:
    return isinstance(x, int)


class SatEmu:
    def __init__(self, fixed: list[int], symbolic: list[int]):
        self.cnf = CNF()
        self.regs = {r: 0 for r in ["rax", "rbx", "rcx", "rdx", "rsi", "rdi", "rbp", "rsp", "r8", "r9", "r10", "r11", "r12", "r13", "r14", "r15"]}
        self.regs.update({
            "rbx": 349,
            "r9": 266,
            "r12": INPUT,
            "r13": 0x200000,
            "r14": 0xAAAAAAAAAAAAAAAB,
            "r15": 0xAAAAAAAAAAAAAAAA,
            "rsp": STACK + 0x100008,
        })
        self.mem: dict[int, TVal | int] = {}
        self.xmm: dict[str, tuple[TVal | int, TVal | int]] = {}
        self.inputs: dict[int, TVal] = {}
        symbolic_set = set(symbolic)
        for i in range(348):
            if i in symbolic_set:
                val = self.cnf.new_val()
                self.inputs[i] = val
                self.mem[INPUT + i * 8] = val
            else:
                self.mem[INPUT + i * 8] = fixed[i]
        self.last_cmp = (0, 0)
        self.steps = 0

    def reg(self, name: str):
        return self.regs.get(canon(name), 0)

    def write_reg(self, name: str, val):
        self.regs[canon(name)] = val

    def addr(self, ins, op) -> int:
        m = op.mem
        base = 0
        if m.base:
            bname = ins.reg_name(m.base)
            base = ins.address + ins.size if bname == "rip" else self.reg(bname)
            if not is_conc(base):
                raise NotImplementedError(("symbolic base", hex(ins.address), ins.op_str))
        idx = 0
        if m.index:
            idx = self.reg(ins.reg_name(m.index))
            if not is_conc(idx):
                raise NotImplementedError(("symbolic index", hex(ins.address), ins.op_str))
        return (base + idx * m.scale + m.disp) & MASK64

    def qload(self, addr: int):
        if addr in self.mem:
            return self.mem[addr]
        if 0 <= addr <= len(CODE) - 8:
            return int.from_bytes(CODE[addr:addr + 8], "little")
        return 0

    def qstore(self, addr: int, val):
        self.mem[addr] = val

    def read_op(self, ins, op):
        if op.type == X86_OP_IMM:
            return op.imm & MASK64
        if op.type == X86_OP_REG:
            return self.reg(ins.reg_name(op.reg))
        if op.type == X86_OP_MEM:
            return self.qload(self.addr(ins, op))
        raise NotImplementedError(op.type)

    def write_op(self, ins, op, val):
        if op.type == X86_OP_REG:
            self.write_reg(ins.reg_name(op.reg), val)
            return
        if op.type == X86_OP_MEM:
            self.qstore(self.addr(ins, op), val)
            return
        raise NotImplementedError(op.type)

    def copy(self, dst: int, src: int, nbytes: int):
        vals = [self.qload(src + i * 8) for i in range(nbytes // 8)]
        for i, val in enumerate(vals):
            self.qstore(dst + i * 8, val)

    def zero(self, dst: int, nbytes: int):
        for i in range(nbytes // 8):
            self.qstore(dst + i * 8, 0)

    def run(self, stop: int, start: int = TEXT_START):
        rip = start
        while rip != stop:
            ins = INSNS[rip]
            nrip = rip + ins.size
            self.steps += 1
            if self.steps % 250000 == 0:
                print("steps", self.steps, "rip", hex(rip), "vars", self.cnf.next_var - 1, "clauses", len(self.cnf.clauses), flush=True)
            m = ins.mnemonic
            ops = ins.operands
            if m in {"nop", "endbr64"}:
                pass
            elif m == "push":
                self.regs["rsp"] -= 8
                self.qstore(self.regs["rsp"], self.read_op(ins, ops[0]))
            elif m == "pop":
                self.write_op(ins, ops[0], self.qload(self.regs["rsp"]))
                self.regs["rsp"] += 8
            elif m == "mov":
                self.write_op(ins, ops[0], self.read_op(ins, ops[1]))
            elif m in {"movups", "movaps"}:
                dst, src = ops
                if dst.type == X86_OP_REG:
                    addr = self.addr(ins, src)
                    self.xmm[ins.reg_name(dst.reg)] = (self.qload(addr), self.qload(addr + 8))
                elif dst.type == X86_OP_MEM:
                    addr = self.addr(ins, dst)
                    lo, hi = self.xmm[ins.reg_name(src.reg)]
                    self.qstore(addr, lo)
                    self.qstore(addr + 8, hi)
                else:
                    raise NotImplementedError((hex(rip), m, ins.op_str))
            elif m == "xorps":
                self.xmm[ins.reg_name(ops[0].reg)] = (0, 0)
            elif m == "xor":
                if ops[0].type == X86_OP_REG and ops[1].type == X86_OP_REG and ops[0].reg == ops[1].reg:
                    self.write_op(ins, ops[0], 0)
                else:
                    a, b = self.read_op(ins, ops[0]), self.read_op(ins, ops[1])
                    if not (is_conc(a) and is_conc(b)):
                        raise NotImplementedError(("symbolic xor", hex(rip), ins.op_str))
                    self.write_op(ins, ops[0], a ^ b)
            elif m == "add":
                a, b = self.read_op(ins, ops[0]), self.read_op(ins, ops[1])
                if not (is_conc(a) and is_conc(b)):
                    raise NotImplementedError(("symbolic add", hex(rip), ins.op_str))
                self.write_op(ins, ops[0], (a + b) & MASK64)
            elif m == "sub":
                a, b = self.read_op(ins, ops[0]), self.read_op(ins, ops[1])
                if not (is_conc(a) and is_conc(b)):
                    raise NotImplementedError(("symbolic sub", hex(rip), ins.op_str))
                self.write_op(ins, ops[0], (a - b) & MASK64)
            elif m == "inc":
                a = self.read_op(ins, ops[0])
                if not is_conc(a):
                    raise NotImplementedError(("symbolic inc", hex(rip), ins.op_str))
                self.write_op(ins, ops[0], a + 1)
            elif m == "shr":
                a, b = self.read_op(ins, ops[0]), self.read_op(ins, ops[1])
                if not (is_conc(a) and is_conc(b)):
                    raise NotImplementedError(("symbolic shr", hex(rip), ins.op_str))
                self.write_op(ins, ops[0], (a >> (b & 0x3f)) & MASK64)
            elif m == "shl":
                a, b = self.read_op(ins, ops[0]), self.read_op(ins, ops[1])
                if not (is_conc(a) and is_conc(b)):
                    raise NotImplementedError(("symbolic shl", hex(rip), ins.op_str))
                self.write_op(ins, ops[0], (a << (b & 0x3f)) & MASK64)
            elif m == "addl":
                a, b = self.read_op(ins, ops[0]), self.read_op(ins, ops[1])
                if not (is_conc(a) and is_conc(b)):
                    raise NotImplementedError(("symbolic addl", hex(rip), ins.op_str))
                self.write_op(ins, ops[0], (a + b) & MASK64)
            elif m == "lea":
                self.write_op(ins, ops[0], self.addr(ins, ops[1]))
            elif m == "imul":
                if len(ops) == 3:
                    a, b = self.read_op(ins, ops[1]), self.read_op(ins, ops[2])
                    self.write_op(ins, ops[0], a * b)
                elif len(ops) == 2:
                    a, b = self.read_op(ins, ops[0]), self.read_op(ins, ops[1])
                    self.write_op(ins, ops[0], a * b)
                else:
                    raise NotImplementedError((hex(rip), m, ins.op_str))
            elif m == "bsr":
                self.write_op(ins, ops[0], self.cnf.bsr(self.read_op(ins, ops[0]), self.read_op(ins, ops[1])))
            elif m == "bzhi":
                self.write_op(ins, ops[0], self.cnf.bzhi(self.read_op(ins, ops[1]), self.read_op(ins, ops[2])))
            elif m == "cmp":
                self.last_cmp = (self.read_op(ins, ops[0]), self.read_op(ins, ops[1]))
            elif m in {"jne", "je", "jb"}:
                a, b = self.last_cmp
                if not (is_conc(a) and is_conc(b)):
                    raise NotImplementedError(("symbolic branch", hex(rip), ins.op_str))
                take = (a != b) if m == "jne" else (a == b) if m == "je" else (a < b)
                if take:
                    nrip = int(ops[0].imm)
            elif m == "jmp":
                nrip = int(ops[0].imm)
            elif m == "cmovae":
                a, b = self.last_cmp
                if not (is_conc(a) and is_conc(b)):
                    raise NotImplementedError(("symbolic cmovae", hex(rip), ins.op_str))
                if a >= b:
                    self.write_op(ins, ops[0], self.read_op(ins, ops[1]))
            elif m == "call":
                target = int(ops[0].imm)
                if target == 0x1050:
                    self.zero(self.regs["rdi"], self.regs["rdx"])
                    self.regs["rax"] = self.regs["rdi"]
                elif target in {0x1060, 0x1090}:
                    self.copy(self.regs["rdi"], self.regs["rsi"], self.regs["rdx"])
                    self.regs["rax"] = self.regs["rdi"]
                else:
                    raise NotImplementedError(("call", hex(rip), hex(target)))
            else:
                raise NotImplementedError((hex(rip), m, ins.op_str))
            rip = nrip


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--stage", choices=["f10"], default="f10")
    ap.add_argument("--solver", default="kissat")
    args = ap.parse_args()
    fixed = flag_to_trits((b"bbb{" + b"A" * 64))
    symbolic = list(range(20, 184))
    emu = SatEmu(fixed, symbolic)
    emu.run(0x17CB6)
    target = struct.unpack("<168Q", CODE[0x4E050:0x4E050 + 168 * 8])
    base = emu.regs["rsp"] + 0x1500
    for i, value in enumerate(target):
        val = emu.qload(base + i * 8)
        if isinstance(val, TVal):
            emu.cnf.force(val, value)
        else:
            if val != value:
                emu.cnf.add([])
    print("built", "steps", emu.steps, "vars", emu.cnf.next_var - 1, "clauses", len(emu.cnf.clauses), flush=True)
    with Solver(name=args.solver, bootstrap_with=emu.cnf.clauses) as solver:
        print("solving", flush=True)
        ok = solver.solve()
        print("sat", ok, flush=True)
        if not ok:
            return
        model = set(solver.get_model())
    trits = fixed[:]
    for i, val in emu.inputs.items():
        assert val.lits is not None
        for digit, lit in enumerate(val.lits):
            if lit in model:
                trits[i] = digit
                break
    raw = trits_to_bytes_prefix(trits)
    print(raw[:68])
    print(trits[20:184])


if __name__ == "__main__":
    main()