// Unit tests for synthesizeChain — the semantics-driven ROP chain planner. Feeds
// synthetic GadgetEffects (no r2/Ghidra) and checks the composed payload layout.
#include <catch2/catch_test_macros.hpp>
#include <optional>
#include <vector>

#include "fold/binary/gadget_semantics.hpp"
#include "fold/binary/rop_chain.hpp"

namespace {
using namespace fold::binary;

// x86-64 register model (matches the Ghidra x86 register file offsets).
GadgetRegModel x64Model() {
    GadgetRegModel m;
    m.ptr_bits = 64;
    m.sp_offset = 0x20;
    m.args = {{"rdi", 0x38}, {"rsi", 0x30}, {"rdx", 0x10}, {"rcx", 0x8}};
    m.syscall_reg = "rax";
    m.syscall_off = 0x0;
    return m;
}

// A `pop <reg> ; ret` effect: reg from stack slot 0, sp_delta 16, ret.
GadgetEffect popRet(std::uint64_t addr, const std::string& reg) {
    GadgetEffect g;
    g.addr = addr;
    g.reg_out[reg] = ValRef{ValSrc::StackSlot, 0, 0, ""};
    g.sp_delta = 16;
    g.term = GadgetEffect::Term::Ret;
    return g;
}

// A `mov <dst>, <src> ; ret` effect: dst sourced from another register.
GadgetEffect movRet(std::uint64_t addr, const std::string& dst, const std::string& src) {
    GadgetEffect g;
    g.addr = addr;
    g.reg_out[dst] = ValRef{ValSrc::Reg, 0, 0, src, 0};
    g.sp_delta = 8;
    g.term = GadgetEffect::Term::Ret;
    return g;
}

// A constant-load (e.g. `xor reg,reg ; ret` for value 0) effect.
GadgetEffect constRet(std::uint64_t addr, const std::string& reg, std::uint64_t value) {
    GadgetEffect g;
    g.addr = addr;
    g.reg_out[reg] = ValRef{ValSrc::Const, 0, value, "", 0};
    g.sp_delta = 8;
    g.term = GadgetEffect::Term::Ret;
    return g;
}

// A `pop <reg> ; add <reg>, <addend> ; ret` effect: stack slot 0 + addend.
GadgetEffect addPopRet(std::uint64_t addr, const std::string& reg, std::int64_t addend) {
    GadgetEffect g;
    g.addr = addr;
    g.reg_out[reg] = ValRef{ValSrc::StackSlot, 0, 0, "", addend};
    g.sp_delta = 16;
    g.term = GadgetEffect::Term::Ret;
    return g;
}

GadgetEffect syscallGadget(std::uint64_t addr) {
    GadgetEffect g;
    g.addr = addr;
    g.term = GadgetEffect::Term::Syscall;
    return g;
}

const ChainStep* step(const RopChain& c, const std::string& role, std::uint64_t addr) {
    for (const auto& s : c.steps)
        if (s.role == role && s.addr == addr)
            return &s;
    return nullptr;
}

const ChainStep* valueStep(const RopChain& c, std::uint64_t value) {
    for (const auto& s : c.steps)
        if (s.role == "value" && s.value == value)
            return &s;
    return nullptr;
}
} // namespace

TEST_CASE("chain: win target with a pop-rdi gadget sets arg0 then calls win", "[chain]") {
    std::vector<GadgetEffect> effects = {popRet(0x1000, "rdi")};
    std::vector<std::pair<std::uint64_t, std::string>> wins = {{0x401000, "system"}};
    auto c = synthesizeChain("x86", 256, x64Model(), effects, wins, std::uint64_t{0x600000});
    REQUIRE(c.has_value());
    REQUIRE(c->target == "win:system");
    REQUIRE(c->padding_len == 256);
    // pop-rdi gadget at payload offset 0 (overwrites RA), rdi value at offset 8,
    // win target at offset 16 (= sp_delta).
    const ChainStep* pop = step(*c, "gadget", 0x1000);
    REQUIRE(pop != nullptr);
    REQUIRE(pop->stack_offset == 0);
    const ChainStep* tgt = step(*c, "target", 0x401000);
    REQUIRE(tgt != nullptr);
    REQUIRE(tgt->stack_offset == 16);
    bool has_binsh_value = false;
    for (const auto& s : c->steps)
        if (s.role == "value" && s.value == 0x600000 && s.stack_offset == 8)
            has_binsh_value = true;
    REQUIRE(has_binsh_value);
}

TEST_CASE("chain: execve from pop gadgets + syscall solves with exact offsets", "[chain]") {
    GadgetEffect sys;
    sys.addr = 0x2000;
    sys.term = GadgetEffect::Term::Syscall;
    std::vector<GadgetEffect> effects = {popRet(0x1000, "rdi"), popRet(0x1008, "rsi"),
                                         popRet(0x1010, "rdx"), popRet(0x1018, "rax"), sys};
    auto c = synthesizeChain("x86", 64, x64Model(), effects, {}, std::uint64_t{0x600000});
    REQUIRE(c.has_value());
    REQUIRE(c->target == "execve");
    // Four pop gadgets (each sp_delta 16) then the syscall gadget. Gadget address
    // slots at 0, 16, 32, 48; syscall target at 64.
    REQUIRE(step(*c, "target", 0x2000) != nullptr);
    REQUIRE(step(*c, "target", 0x2000)->stack_offset == 64);
    // rax = execve number (59 on x86-64) placed as a value somewhere.
    bool has_nr = false;
    for (const auto& s : c->steps)
        if (s.role == "value" && s.value == 59)
            has_nr = true;
    REQUIRE(has_nr);
}

TEST_CASE("chain: no /bin/sh plants it via a write-what-where gadget", "[chain]") {
    GadgetEffect wgt; // mov [rdi], rsi ; ret
    wgt.addr = 0x4000;
    wgt.mem_writes.push_back({ValRef{ValSrc::Reg, 0, 0, "rdi"}, ValRef{ValSrc::Reg, 0, 0, "rsi"}});
    wgt.sp_delta = 8;
    wgt.term = GadgetEffect::Term::Ret;
    GadgetEffect sys;
    sys.addr = 0x2000;
    sys.term = GadgetEffect::Term::Syscall;
    std::vector<GadgetEffect> effects = {popRet(0x1000, "rdi"),
                                         popRet(0x1008, "rsi"),
                                         popRet(0x1010, "rdx"),
                                         popRet(0x1018, "rax"),
                                         wgt,
                                         sys};
    // No /bin/sh in the binary, but scratch memory is available.
    auto c =
        synthesizeChain("x86", 64, x64Model(), effects, {}, std::nullopt, std::uint64_t{0x404000});
    REQUIRE(c.has_value());
    REQUIRE(c->target == "execve");
    REQUIRE(step(*c, "gadget", 0x4000) != nullptr); // the write gadget is used
    bool planted = false;
    for (const auto& cav : c->caveats)
        if (cav.find("write-what-where") != std::string::npos)
            planted = true;
    REQUIRE(planted);
}

TEST_CASE("chain: no syscall and no win target yields no chain", "[chain]") {
    std::vector<GadgetEffect> effects = {popRet(0x1000, "rdi")};
    REQUIRE_FALSE(synthesizeChain("x86", 64, x64Model(), effects, {}, std::nullopt).has_value());
}

TEST_CASE("chain: mov-chain controls an arg with no direct pop (mov rdi,rax ← pop rax)",
          "[chain]") {
    // No `pop rdi`; rdi is only reachable via `mov rdi, rax` fed by `pop rax`.
    std::vector<GadgetEffect> effects = {movRet(0x1100, "rdi", "rax"), popRet(0x1000, "rax")};
    std::vector<std::pair<std::uint64_t, std::string>> wins = {{0x401000, "system"}};
    auto c = synthesizeChain("x86", 256, x64Model(), effects, wins, std::uint64_t{0x600000});
    REQUIRE(c.has_value());
    REQUIRE(c->target == "win:system");
    const ChainStep* pop = step(*c, "gadget", 0x1000); // pop rax (sets the source)
    const ChainStep* mov = step(*c, "gadget", 0x1100); // mov rdi, rax (the mover)
    REQUIRE(pop != nullptr);
    REQUIRE(mov != nullptr);
    REQUIRE(pop->stack_offset < mov->stack_offset); // source set before the move runs
    const ChainStep* v = valueStep(*c, 0x600000);   // binsh placed at rax's slot
    REQUIRE(v != nullptr);
    REQUIRE(v->stack_offset > pop->stack_offset);
    REQUIRE(v->stack_offset < mov->stack_offset);
}

TEST_CASE("chain: const/zero gadgets set NULL execve args without stack values", "[chain]") {
    // rsi/rdx zeroed via `xor reg,reg` (Const 0); rdi/rax via pops; then syscall.
    std::vector<GadgetEffect> effects = {popRet(0x1000, "rdi"), popRet(0x1018, "rax"),
                                         constRet(0x1200, "rsi", 0), constRet(0x1208, "rdx", 0),
                                         syscallGadget(0x2000)};
    auto c = synthesizeChain("x86", 64, x64Model(), effects, {}, std::uint64_t{0x600000});
    REQUIRE(c.has_value());
    REQUIRE(c->target == "execve");
    REQUIRE(step(*c, "gadget", 0x1200) != nullptr); // xor rsi,rsi used
    REQUIRE(step(*c, "gadget", 0x1208) != nullptr); // xor rdx,rdx used
    // Every argument register was controlled — no "incomplete" caveat.
    for (const auto& cav : c->caveats)
        REQUIRE(cav.find("incomplete") == std::string::npos);
}

TEST_CASE("chain: additive pop;add gadget places value minus the addend", "[chain]") {
    // `pop rdi ; add rdi, 0x10 ; ret` reaches binsh when the stacked value is
    // binsh - 0x10.
    std::vector<GadgetEffect> effects = {addPopRet(0x1300, "rdi", 0x10)};
    std::vector<std::pair<std::uint64_t, std::string>> wins = {{0x401000, "system"}};
    auto c = synthesizeChain("x86", 256, x64Model(), effects, wins, std::uint64_t{0x600000});
    REQUIRE(c.has_value());
    REQUIRE(step(*c, "gadget", 0x1300) != nullptr);
    REQUIRE(valueStep(*c, 0x600000 - 0x10) != nullptr); // inverted addend
}

TEST_CASE("chain: a register-move cycle terminates and controls nothing", "[chain]") {
    // rdi←r8, r8←r9, r9←r8 with no pop: the recursion must hit the visited guard
    // and roll back, leaving no gadget steps (honest: rdi uncontrolled).
    std::vector<GadgetEffect> effects = {movRet(0x1000, "rdi", "r8"), movRet(0x1008, "r8", "r9"),
                                         movRet(0x1010, "r9", "r8")};
    std::vector<std::pair<std::uint64_t, std::string>> wins = {{0x401000, "system"}};
    auto c = synthesizeChain("x86", 256, x64Model(), effects, wins, std::uint64_t{0x600000});
    REQUIRE(c.has_value());
    REQUIRE(c->target == "win:system");
    for (const auto& s : c->steps)
        REQUIRE(s.role != "gadget"); // nothing was chainable — no partial litter
}

TEST_CASE("chain: a gadget with non-positive sp_delta is rejected (can't drive a linear chain)",
          "[chain]") {
    GadgetEffect bad = popRet(0x1000, "rdi"); // pop rdi but rsp rewinds
    bad.sp_delta = -8;
    std::vector<GadgetEffect> effects = {bad};
    std::vector<std::pair<std::uint64_t, std::string>> wins = {{0x401000, "system"}};
    auto c = synthesizeChain("x86", 256, x64Model(), effects, wins, std::uint64_t{0x600000});
    REQUIRE(c.has_value());
    REQUIRE(c->target == "win:system");
    for (const auto& s : c->steps)
        REQUIRE(s.role != "gadget"); // the rewinding gadget must not be used
}

TEST_CASE("chain: among equal const gadgets the smallest sp_delta is chosen", "[chain]") {
    GadgetEffect big = constRet(0x5000, "rdi", 0x600000); // larger stack adjustment
    big.sp_delta = 64;
    GadgetEffect small = constRet(0x5008, "rdi", 0x600000); // compact: sp_delta 8
    std::vector<GadgetEffect> effects = {big, small};
    std::vector<std::pair<std::uint64_t, std::string>> wins = {{0x401000, "system"}};
    auto c = synthesizeChain("x86", 256, x64Model(), effects, wins, std::uint64_t{0x600000});
    REQUIRE(c.has_value());
    REQUIRE(step(*c, "gadget", 0x5008) != nullptr); // compact gadget chosen
    REQUIRE(step(*c, "gadget", 0x5000) == nullptr); // wasteful one avoided
}

TEST_CASE("chain: execve goal is best-effort — keeps controllable args when one is unreachable",
          "[chain]") {
    // rdi/rsi controllable; rdx/rax have no setter. The chain must retain the
    // rdi/rsi gadgets (not roll back to empty) and honestly caveat the rest.
    std::vector<GadgetEffect> effects = {popRet(0x1000, "rdi"), popRet(0x1008, "rsi"),
                                         syscallGadget(0x2000)};
    auto c = synthesizeChain("x86", 64, x64Model(), effects, {}, std::uint64_t{0x600000});
    REQUIRE(c.has_value());
    REQUIRE(c->target == "execve");
    REQUIRE(step(*c, "gadget", 0x1000) != nullptr); // rdi controlled and kept
    REQUIRE(step(*c, "gadget", 0x1008) != nullptr); // rsi controlled and kept
    REQUIRE(step(*c, "target", 0x2000) != nullptr); // still lands on the syscall
    bool incomplete = false;
    for (const auto& cav : c->caveats)
        if (cav.find("incomplete") != std::string::npos)
            incomplete = true;
    REQUIRE(incomplete);
}

TEST_CASE("chain: a multi-pop gadget fills several argument registers at once", "[chain]") {
    GadgetEffect multi; // pop rdi ; pop rsi ; ret
    multi.addr = 0x3000;
    multi.reg_out["rdi"] = ValRef{ValSrc::StackSlot, 0, 0, ""};
    multi.reg_out["rsi"] = ValRef{ValSrc::StackSlot, 8, 0, ""};
    multi.sp_delta = 24;
    multi.term = GadgetEffect::Term::Ret;
    GadgetEffect popdx = popRet(0x1010, "rdx");
    GadgetEffect popax = popRet(0x1018, "rax");
    GadgetEffect sys;
    sys.addr = 0x2000;
    sys.term = GadgetEffect::Term::Syscall;
    std::vector<GadgetEffect> effects = {multi, popdx, popax, sys};
    auto c = synthesizeChain("x86", 64, x64Model(), effects, {}, std::uint64_t{0x600000});
    REQUIRE(c.has_value());
    REQUIRE(c->target == "execve");
    // The multi-pop gadget appears once and controls both rdi and rsi.
    REQUIRE(step(*c, "gadget", 0x3000) != nullptr);
}

TEST_CASE("chain: later pop gadget must not clobber an already satisfied register", "[chain]") {
    GadgetEffect clobber; // pop rdi ; pop rsi ; ret, unsafe after rdi is already set
    clobber.addr = 0x3000;
    clobber.reg_out["rdi"] = ValRef{ValSrc::StackSlot, 0, 0, ""};
    clobber.reg_out["rsi"] = ValRef{ValSrc::StackSlot, 8, 0, ""};
    clobber.sp_delta = 24;
    clobber.term = GadgetEffect::Term::Ret;
    std::vector<GadgetEffect> effects = {popRet(0x1000, "rdi"), clobber,
                                         popRet(0x1008, "rsi"), popRet(0x1010, "rdx"),
                                         popRet(0x1018, "rax"), syscallGadget(0x2000)};
    auto c = synthesizeChain("x86", 64, x64Model(), effects, {}, std::uint64_t{0x600000});
    REQUIRE(c.has_value());
    REQUIRE(step(*c, "gadget", 0x1000) != nullptr);
    REQUIRE(step(*c, "gadget", 0x1008) != nullptr);
    REQUIRE(step(*c, "gadget", 0x3000) == nullptr);
}

TEST_CASE("chain: helper registers do not make goal satisfaction stop early", "[chain]") {
    std::vector<GadgetEffect> effects = {movRet(0x1100, "rdi", "r10"), popRet(0x1000, "r10"),
                                         movRet(0x1110, "rsi", "r11"), popRet(0x1008, "r11"),
                                         popRet(0x1010, "rdx"),        popRet(0x1018, "rax"),
                                         syscallGadget(0x2000)};
    auto c = synthesizeChain("x86", 64, x64Model(), effects, {}, std::uint64_t{0x600000});
    REQUIRE(c.has_value());
    REQUIRE(step(*c, "gadget", 0x1010) != nullptr); // rdx goal still attempted
    REQUIRE(step(*c, "gadget", 0x1018) != nullptr); // rax goal still attempted
    for (const auto& cav : c->caveats)
        REQUIRE(cav.find("incomplete") == std::string::npos);
}

TEST_CASE("chain: shell win without an argument falls through to plantable execve", "[chain]") {
    GadgetEffect wgt; // mov [rdi], rsi ; ret
    wgt.addr = 0x4000;
    wgt.mem_writes.push_back({ValRef{ValSrc::Reg, 0, 0, "rdi"}, ValRef{ValSrc::Reg, 0, 0, "rsi"}});
    wgt.sp_delta = 8;
    wgt.term = GadgetEffect::Term::Ret;
    std::vector<GadgetEffect> effects = {popRet(0x1000, "rdi"),
                                         popRet(0x1008, "rsi"),
                                         popRet(0x1010, "rdx"),
                                         popRet(0x1018, "rax"),
                                         wgt,
                                         syscallGadget(0x2000)};
    std::vector<std::pair<std::uint64_t, std::string>> wins = {{0x401000, "system"}};
    auto c = synthesizeChain("x86", 64, x64Model(), effects, wins, std::nullopt,
                             std::uint64_t{0x404000});
    REQUIRE(c.has_value());
    REQUIRE(c->target == "execve");
    REQUIRE(step(*c, "gadget", 0x4000) != nullptr);
}

TEST_CASE("chain: non-chainable write-what-where gadget is not used for planting", "[chain]") {
    GadgetEffect bad_write;
    bad_write.addr = 0x4000;
    bad_write.mem_writes.push_back(
        {ValRef{ValSrc::Reg, 0, 0, "rdi"}, ValRef{ValSrc::Reg, 0, 0, "rsi"}});
    bad_write.sp_delta = INT64_MIN;
    bad_write.term = GadgetEffect::Term::Ret;
    std::vector<GadgetEffect> effects = {popRet(0x1000, "rdi"),
                                         popRet(0x1008, "rsi"),
                                         popRet(0x1010, "rdx"),
                                         popRet(0x1018, "rax"),
                                         bad_write,
                                         syscallGadget(0x2000)};
    auto c =
        synthesizeChain("x86", 64, x64Model(), effects, {}, std::nullopt, std::uint64_t{0x404000});
    REQUIRE(c.has_value());
    REQUIRE(step(*c, "gadget", 0x4000) == nullptr);
    bool missing_binsh = false;
    for (const auto& cav : c->caveats)
        if (cav.find("no \"/bin/sh\"") != std::string::npos)
            missing_binsh = true;
    REQUIRE(missing_binsh);
}