/*
 * BastionGuard™
 * Copyright (C) 2025–2026 Calogero Scarnà
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, version 3.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 *
 * BastionGuard™ is a trademark of Calogero Scarnà.
 * The BastionGuard™ name and branding are not licensed under the GPL.
 */

#include "ArchiveSandbox.hpp"

#include <algorithm>
#include <array>
#include <cerrno>
#include <chrono>
#include <cctype>
#include <csignal>
#include <cstring>
#include <filesystem>
#include <string>
#include <thread>
#include <vector>

#include <fcntl.h>
#include <poll.h>
#include <sys/resource.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>

namespace fs = std::filesystem;

namespace {

// Legge tutto quello che è disponibile subito su fd (non-blocking).
// Ritorna false solo su errore reale (non su EAGAIN/EWOULDBLOCK/EINTR).
// Ritorna true anche in caso di EOF (n == 0), segnalando chiusura via got_eof.
bool drain_fd_nonblocking(int fd, std::string& out, bool& got_eof, bool& truncated) {
    constexpr std::size_t max_output_bytes = 1024 * 1024;
    got_eof = false;
    std::array<char, 4096> buf{};
    while (true) {
        const ssize_t n = ::read(fd, buf.data(), buf.size());
        if (n == 0) {
            got_eof = true;
            return true;
        }
        if (n < 0) {
            if (errno == EINTR) continue;
            if (errno == EAGAIN || errno == EWOULDBLOCK) return true;
            return false;
        }
        const auto length = static_cast<std::size_t>(n);
        if (out.size() < max_output_bytes) {
            const auto copy = std::min(length, max_output_bytes - out.size());
            out.append(buf.data(), copy);
            if (copy != length) truncated = true;
        } else {
            truncated = true;
        }
    }
}

bool set_nonblocking(int fd) {
    const int flags = ::fcntl(fd, F_GETFL, 0);
    if (flags < 0) return false;
    return ::fcntl(fd, F_SETFL, flags | O_NONBLOCK) == 0;
}

void set_worker_rlimits() {
    struct rlimit cpu_lim {};
    cpu_lim.rlim_cur = 10;
    cpu_lim.rlim_max = 10;
    ::setrlimit(RLIMIT_CPU, &cpu_lim);

    struct rlimit as_lim {};
    as_lim.rlim_cur = 512ULL * 1024ULL * 1024ULL;
    as_lim.rlim_max = 512ULL * 1024ULL * 1024ULL;
    ::setrlimit(RLIMIT_AS, &as_lim);

    struct rlimit fsize_lim {};
    fsize_lim.rlim_cur = 128ULL * 1024ULL * 1024ULL;
    fsize_lim.rlim_max = 128ULL * 1024ULL * 1024ULL;
    ::setrlimit(RLIMIT_FSIZE, &fsize_lim);

    struct rlimit nofile_lim {};
    nofile_lim.rlim_cur = 64;
    nofile_lim.rlim_max = 64;
    ::setrlimit(RLIMIT_NOFILE, &nofile_lim);
}

bool json_has_true(const std::string& json, const std::string& key) {
    const std::string pat = "\"" + key + "\":true";
    return json.find(pat) != std::string::npos;
}

int json_read_int(const std::string& json, const std::string& key, int fallback = 0) {
    const std::string pat = "\"" + key + "\":";
    std::size_t pos = json.find(pat);
    if (pos == std::string::npos) {
        return fallback;
    }

    pos += pat.size();
    while (pos < json.size() && std::isspace(static_cast<unsigned char>(json[pos]))) {
        ++pos;
    }

    bool neg = false;
    if (pos < json.size() && json[pos] == '-') {
        neg = true;
        ++pos;
    }

    int value = 0;
    bool seen = false;
    while (pos < json.size() && std::isdigit(static_cast<unsigned char>(json[pos]))) {
        seen = true;
        value = value * 10 + (json[pos] - '0');
        ++pos;
    }

    if (!seen) {
        return fallback;
    }

    return neg ? -value : value;
}

std::string json_read_string(const std::string& json, const std::string& key) {
    const std::string pat = "\"" + key + "\":\"";
    std::size_t pos = json.find(pat);
    if (pos == std::string::npos) {
        return {};
    }

    pos += pat.size();

    std::string out;
    while (pos < json.size()) {
        char c = json[pos++];
        if (c == '\\') {
            if (pos < json.size()) {
                out.push_back(json[pos++]);
            }
            continue;
        }
        if (c == '"') {
            break;
        }
        out.push_back(c);
    }

    return out;
}

} // namespace

SandboxExecResult run_archive_worker_bwrap(
    const std::string& worker_path,
    const std::string& input_path,
    int timeout_seconds)
{
    SandboxExecResult result;

    if (!fs::exists(worker_path) || !fs::is_regular_file(worker_path)) {
        result.stderr_text = "invalid worker_path";
        return result;
    }

    if (!fs::exists(input_path) || !fs::is_regular_file(input_path)) {
        result.stderr_text = "invalid input_path";
        return result;
    }

    char tmp_template[] = "/tmp/bastionguard-archive-XXXXXX";
    char* tmp_dir = ::mkdtemp(tmp_template);
    if (!tmp_dir) {
        result.stderr_text = "mkdtemp failed";
        return result;
    }

    const fs::path sandbox_root(tmp_dir);
    const fs::path work_dir = sandbox_root / "work";
    std::error_code ec;
    fs::create_directories(work_dir, ec);
    if (ec) {
        result.stderr_text = "create_directories failed";
        fs::remove_all(sandbox_root, ec);
        return result;
    }

    int out_pipe[2] = {-1, -1};
    int err_pipe[2] = {-1, -1};

    if (::pipe(out_pipe) != 0 || ::pipe(err_pipe) != 0) {
        result.stderr_text = "pipe failed";
        for (int fd : {out_pipe[0], out_pipe[1], err_pipe[0], err_pipe[1]}) {
            if (fd >= 0) ::close(fd);
        }
        fs::remove_all(sandbox_root, ec);
        return result;
    }

    const pid_t pid = ::fork();
    if (pid < 0) {
        result.stderr_text = "fork failed";
        ::close(out_pipe[0]);
        ::close(out_pipe[1]);
        ::close(err_pipe[0]);
        ::close(err_pipe[1]);
        fs::remove_all(sandbox_root, ec);
        return result;
    }

    if (pid == 0) {
        ::close(out_pipe[0]);
        ::close(err_pipe[0]);

        ::dup2(out_pipe[1], STDOUT_FILENO);
        ::dup2(err_pipe[1], STDERR_FILENO);

        ::close(out_pipe[1]);
        ::close(err_pipe[1]);

        set_worker_rlimits();

        std::vector<std::string> args = {
        "bwrap",
        "--die-with-parent",
        "--new-session",
        "--unshare-ipc",
        "--unshare-pid",
        "--unshare-uts",
        "--unshare-cgroup-try",
        "--proc", "/proc",
        "--dev", "/dev",
        "--ro-bind", "/usr", "/usr",
        "--ro-bind", "/bin", "/bin",
        "--ro-bind", "/lib", "/lib",
        "--ro-bind", "/lib64", "/lib64",
        "--tmpfs", "/tmp",
        "--dir", "/work",
        "--ro-bind", input_path, "/work/input.zip",
        "--ro-bind", worker_path, "/worker",
        "--chdir", "/work",
        "--setenv", "HOME", "/tmp",
        "--setenv", "TMPDIR", "/tmp",
        "/worker",
        "/work/input.zip"
    };

        std::vector<char*> argv;
        argv.reserve(args.size() + 1);
        for (auto& s : args) {
            argv.push_back(s.data());
        }
        argv.push_back(nullptr);

        ::execvp("bwrap", argv.data());
        _exit(127);
    }

    result.launched = true;

    ::close(out_pipe[1]);
    ::close(err_pipe[1]);

    // Mettiamo entrambe le pipe in non-blocking così possiamo drenarle
    // mentre il child sta ancora girando. Se non lo facessimo, un child
    // che scrive più di PIPE_BUF si bloccherebbe su write() e non uscirebbe
    // mai -> deadlock fino al timeout.
    set_nonblocking(out_pipe[0]);
    set_nonblocking(err_pipe[0]);

    const auto start = std::chrono::steady_clock::now();
    bool child_done = false;
    int status = 0;

    bool out_eof = false;
    bool err_eof = false;

    const int effective_timeout = std::max(1, timeout_seconds);
    while (true) {
        // Se il child è già uscito E entrambe le pipe sono in EOF,
        // abbiamo finito.
        if (!child_done) {
            const pid_t w = ::waitpid(pid, &status, WNOHANG);
            if (w == pid) {
                child_done = true;
            } else if (w < 0 && errno != EINTR) {
                result.stderr_text = "waitpid failed";
                // Do not exit leaving an untracked child behind.
                ::kill(pid, SIGKILL);
                ::waitpid(pid, &status, 0);
                break;
            }
        }

        // Timeout: SIGKILL e reap sincrono.
        const auto now = std::chrono::steady_clock::now();
        const auto elapsed =
            std::chrono::duration_cast<std::chrono::seconds>(now - start).count();
        if (elapsed >= effective_timeout && (!child_done || !out_eof || !err_eof)) {
            // A forked grandchild may still own an output pipe after its
            // direct parent exits. Bound the entire operation, not only waitpid.
            result.timed_out = true;
            if (!child_done) {
                ::kill(pid, SIGKILL);
                ::waitpid(pid, &status, 0);
                child_done = true;
            }
            break;
        }

        // Drena le pipe disponibili (non-blocking).
        struct pollfd fds[2];
        int nfds = 0;
        if (!out_eof) { fds[nfds].fd = out_pipe[0]; fds[nfds].events = POLLIN; ++nfds; }
        if (!err_eof) { fds[nfds].fd = err_pipe[0]; fds[nfds].events = POLLIN; ++nfds; }

        if (nfds > 0) {
            // Poll breve: se il child non è ancora morto, controlliamo poi
            // waitpid al prossimo giro; se è morto, drenaiamo finché non c'è
            // più nulla da leggere.
            const int poll_timeout_ms = child_done ? 100 : 50;
            const int pr = ::poll(fds, nfds, poll_timeout_ms);

            if (pr > 0) {
                for (int i = 0; i < nfds; ++i) {
                    if (!(fds[i].revents & (POLLIN | POLLHUP | POLLERR))) continue;

                    bool eof_here = false;
                    if (fds[i].fd == out_pipe[0]) {
                        drain_fd_nonblocking(out_pipe[0], result.stdout_text, eof_here, result.output_truncated);
                        if (eof_here) out_eof = true;
                    } else if (fds[i].fd == err_pipe[0]) {
                        drain_fd_nonblocking(err_pipe[0], result.stderr_text, eof_here, result.output_truncated);
                        if (eof_here) err_eof = true;
                    }
                }
            }
            // pr == 0 (timeout poll) o pr < 0 con EINTR: proseguiamo il loop.
        }

        // Condizione di uscita: child morto e nessun ulteriore dato in arrivo.
        if (child_done && (out_eof || nfds == 0) && (err_eof || nfds == 0)) {
            break;
        }

        // Sanity: se non c'è più nulla da drenare ma il child non è morto,
        // aspettiamo un pochino per non bruciare CPU.
        if (!child_done && nfds == 0) {
            std::this_thread::sleep_for(std::chrono::milliseconds(50));
        }
    }

    // Drena residui (le EOF dovrebbero essere già arrivate, ma per sicurezza).
    if (!out_eof) { bool e; drain_fd_nonblocking(out_pipe[0], result.stdout_text, e, result.output_truncated); }
    if (!err_eof) { bool e; drain_fd_nonblocking(err_pipe[0], result.stderr_text, e, result.output_truncated); }

    ::close(out_pipe[0]);
    ::close(err_pipe[0]);

    if (child_done) {
        result.exited = true;
        if (WIFEXITED(status)) {
            result.exit_code = WEXITSTATUS(status);
        } else if (WIFSIGNALED(status)) {
            result.exit_code = 128 + WTERMSIG(status);
        }
    }

    fs::remove_all(sandbox_root, ec);
    return result;
}

ArchiveInspectionResult analyze_archive_in_sandbox(
    const std::string& worker_path,
    const std::string& input_path,
    int timeout_seconds)
{
    ArchiveInspectionResult r;

    const auto exec = run_archive_worker_bwrap(worker_path, input_path, timeout_seconds);

    if (!exec.launched) {
        r.success = false;
        r.sandbox_failed = true;
        r.reason_code = "ARCHIVE_SANDBOX_LAUNCH_FAILED";
        r.detail = exec.stderr_text;
        return r;
    }

    if (exec.output_truncated) {
        r.success = false;
        r.sandbox_failed = true;
        r.reason_code = "ARCHIVE_SANDBOX_OUTPUT_LIMIT";
        r.detail = "archive worker exceeded output limit";
        return r;
    }

    if (exec.timed_out) {
        r.success = false;
        r.sandbox_timeout = true;
        r.reason_code = "ARCHIVE_SANDBOX_TIMEOUT";
        r.detail = "archive worker timeout";
        r.raw_json = exec.stdout_text;
        return r;
    }

    if (!exec.exited || exec.exit_code != 0) {
        r.success = false;
        r.sandbox_failed = true;
        r.reason_code = "ARCHIVE_SANDBOX_ANALYSIS_FAILED";
        r.detail = exec.stderr_text.empty() ? "archive worker failed" : exec.stderr_text;
        r.raw_json = exec.stdout_text;
        return r;
    }

    r.success = true;
    r.raw_json = exec.stdout_text;
    r.is_zip = json_has_true(exec.stdout_text, "is_zip");
    r.stored_size_mismatch = json_has_true(exec.stdout_text, "stored_size_mismatch");
    r.header_mismatch = json_has_true(exec.stdout_text, "header_mismatch");
    r.invalid_offset = json_has_true(exec.stdout_text, "invalid_offset");
    r.overlapping_entries = json_has_true(exec.stdout_text, "overlapping_entries");
    r.nested_archive = json_has_true(exec.stdout_text, "nested_archive");
    r.risk_score = json_read_int(exec.stdout_text, "risk_score", 0);
    r.reason_code = json_read_string(exec.stdout_text, "reason_code");
    r.detail = json_read_string(exec.stdout_text, "detail");

    const std::string risk = json_read_string(exec.stdout_text, "risk");
    if (risk == "EVASIVE") {
        r.risk = ArchiveRisk::EVASIVE;
    } else if (risk == "MALFORMED") {
        r.risk = ArchiveRisk::MALFORMED;
    } else if (risk == "SUSPICIOUS") {
        r.risk = ArchiveRisk::SUSPICIOUS;
    } else {
        r.risk = ArchiveRisk::CLEAN;
    }

    return r;
}