#include <clang/Frontend/ASTUnit.h>
#include <clang/Tooling/Tooling.h>

#include <catch2/catch_test_macros.hpp>
#include <memory>
#include <string>
#include <vector>

#include "fold/detect/cross_tu_cache.hpp"
#include "fold/detect/taint_track.hpp"
#include "fold/report/finding.hpp"
#include "fold/symbolic/budget.hpp"

namespace {

struct TestResult {
    std::vector<fold::Finding> findings;
    fold::SolverStats stats;
};

std::unique_ptr<clang::ASTUnit> parse(const std::string& code) {
    return clang::tooling::buildASTFromCodeWithArgs(code, {"-std=c++20"}, "input.cpp");
}

TestResult run(const std::string& code) {
    fold::detect::CrossTUCache::instance().clear();
    auto unit = parse(code);
    REQUIRE(unit);
    fold::SolverBudget b;
    fold::SolverStats s;
    auto findings = fold::detectTaintTrackInContext(unit->getASTContext(), b, s);
    return {std::move(findings), s};
}

} // namespace

TEST_CASE("taint: system(getenv(...)) is flagged", "[taint][getenv]") {
    auto [findings, stats] = run("extern \"C\" char* getenv(const char*);\n"
                                 "extern \"C\" int system(const char*);\n"
                                 "int f() {\n"
                                 "  char* p = getenv(\"PATH\");\n"
                                 "  return system(p);\n"
                                 "}");
    REQUIRE(findings.size() == 1);
    REQUIRE(findings[0].kind == fold::FindingKind::TaintTrack);
    REQUIRE(findings[0].message.find("tainted source") != std::string::npos);
    REQUIRE(findings[0].message.find("system") != std::string::npos);
}

TEST_CASE("taint: hardcoded string to system is safe", "[taint][safe]") {
    auto [findings, stats] = run("extern \"C\" int system(const char*);\n"
                                 "int f() {\n"
                                 "  return system(\"ls\");\n"
                                 "}");
    REQUIRE(findings.empty());
}

TEST_CASE("taint: move propagation", "[taint][move]") {
    auto [findings, stats] = run("namespace std { template<typename T> T&& move(T& t); }\n"
                                 "extern \"C\" char* getenv(const char*);\n"
                                 "extern \"C\" int system(const char*);\n"
                                 "int f() {\n"
                                 "  char* p = getenv(\"PATH\");\n"
                                 "  char* q = std::move(p);\n"
                                 "  return system(q);\n"
                                 "}");
    REQUIRE(findings.size() == 1);
    REQUIRE(findings[0].kind == fold::FindingKind::TaintTrack);
}

TEST_CASE("taint: forward propagation", "[taint][forward]") {
    auto [findings, stats] = run("namespace std { template<typename T> T&& forward(typename "
                                 "std::remove_reference<T>::type& t); }\n"
                                 "extern \"C\" char* getenv(const char*);\n"
                                 "extern \"C\" int system(const char*);\n"
                                 "int f() {\n"
                                 "  char* p = getenv(\"PATH\");\n"
                                 "  char* q = std::forward<char*>(p);\n"
                                 "  return system(q);\n"
                                 "}");
    REQUIRE(findings.size() == 1);
    REQUIRE(findings[0].kind == fold::FindingKind::TaintTrack);
}

TEST_CASE("taint: BinaryDeserializer::readString source", "[taint][readString]") {
    auto [findings, stats] =
        run("namespace std { class string { public: const char* c_str() const; }; }\n"
            "namespace yams::daemon {\n"
            "  struct BinaryDeserializer {\n"
            "    std::string readString();\n"
            "  };\n"
            "}\n"
            "extern \"C\" int system(const char*);\n"
            "int f(yams::daemon::BinaryDeserializer& bd) {\n"
            "  std::string s = bd.readString();\n"
            "  return system(s.c_str());\n"
            "}");
    REQUIRE(findings.size() == 1);
    REQUIRE(findings[0].kind == fold::FindingKind::TaintTrack);
}

TEST_CASE("taint: Result::value() propagation", "[taint][result_value]") {
    auto [findings, stats] =
        run("namespace std { class string { public: const char* c_str() const; }; }\n"
            "template<typename T>\n"
            "struct Result {\n"
            "  T val;\n"
            "  T& value() { return val; }\n"
            "};\n"
            "extern \"C\" char* getenv(const char*);\n"
            "extern \"C\" int system(const char*);\n"
            "int f() {\n"
            "  Result<char*> r;\n"
            "  r.val = getenv(\"PATH\");\n"
            "  char* p = r.value();\n"
            "  return system(p);\n"
            "}");
    REQUIRE(findings.size() == 1);
    REQUIRE(findings[0].kind == fold::FindingKind::TaintTrack);
}