path: root/include/wireframe/capture_session.hpp

#pragma once

#include <pcap.h>

#include <atomic>
#include <csignal>
#include <cstdio>
#include <cstring>
#include <optional>
#include <string>
#include <thread>

#include "wireframe/capture_queue.hpp"
#include "wireframe/filter.hpp"
#include "wireframe/pcapng/reader.hpp"
#include "wireframe/pcapng/writer.hpp"
#include "wireframe/privileges.hpp"

// Device-open -> datalink-validate -> filter/pcapng-setup -> signal-hook
// pipeline, shared by every frontend (CLI, TUI, GUI). Centralized so a
// new frontend can't silently skip a step the others rely on - e.g.
// the DLT_RAW/DLT_EN10MB check that summarize_packet() depends on, or
// the pcap_breakloop() shutdown hook that keeps a -w pcapng file from
// being truncated on Ctrl-C (see main.cpp's history: both were real
// bugs before this was centralized).
//
// Also covers replay mode (-r <file>): reading a previously-saved
// pcapng file back through the exact same queue/render/search pipeline
// as a live capture, so every frontend gets it for free rather than
// needing a second code path. The render/consumer side only ever talks
// to a CaptureQueue - it has no way to tell whether packets are
// arriving from a live pcap_loop or being read back from disk.
namespace wireframe {

namespace detail {
inline pcap_t* g_capture_handle = nullptr;
inline std::atomic<bool>* g_replay_stop_flag = nullptr;
inline void handle_stop_signal(int) {
    if (g_capture_handle != nullptr) pcap_breakloop(g_capture_handle);
    if (g_replay_stop_flag != nullptr) g_replay_stop_flag->store(true);
}
}  // namespace detail

struct CaptureSessionOptions {
    std::string device;  // empty = pick the first device via pcap_findalldevs
    std::optional<std::string> filter_expr;
    std::optional<std::string> pcapng_output_path;
    std::optional<std::string> replay_input_path;  // -r: read from this pcapng file, not a live device
};

inline bool is_supported_datalink(int datalink) {
    return datalink == DLT_EN10MB || datalink == DLT_RAW;
}

// Kernel/NIC-level counters, distinct from CaptureQueue::dropped():
// the queue can only count packets libpcap already handed to our
// callback. A traffic spike can drop packets in the kernel's capture
// buffer before that ever happens - invisible without this. Not
// meaningful in replay mode (stats() returns nullopt there).
struct CaptureStats {
    unsigned int received;    // ps_recv
    unsigned int dropped;     // ps_drop: kernel buffer had no room
    unsigned int if_dropped;  // ps_ifdrop: dropped by the interface/driver
};

class CaptureSession {
public:
    ~CaptureSession() { close(); }

    CaptureSession() = default;
    CaptureSession(const CaptureSession&) = delete;
    CaptureSession& operator=(const CaptureSession&) = delete;

    // Returns an error message on failure. The session remains safe to
    // destroy (or close()) regardless of how far setup got.
    std::optional<std::string> open(const CaptureSessionOptions& options) {
        if (options.replay_input_path) {
            if (options.filter_expr) {
                return std::string(
                    "-f (capture filter) isn't supported with -r (replay); use -g to filter "
                    "what's displayed instead");
            }
            return open_replay(*options.replay_input_path, options.pcapng_output_path);
        }

        char errbuf[PCAP_ERRBUF_SIZE];

        if (options.device.empty()) {
            if (pcap_findalldevs(&all_devices_, errbuf) == -1 || all_devices_ == nullptr) {
                return std::string("no capture device found: ") + errbuf;
            }
            device_ = all_devices_->name;
        } else {
            device_ = options.device;
        }

        handle_ = pcap_open_live(device_.c_str(), /*snaplen=*/65535, /*promisc=*/0,
                                  /*to_ms=*/1000, errbuf);
        if (handle_ == nullptr) {
            return std::string("pcap_open_live failed: ") + errbuf;
        }

        // Everything CAP_NET_RAW/root was needed for is done: the
        // handle is open. Drop immediately, before the datalink check
        // or -w's file is even created - the latter is also why this
        // runs this early rather than at the very end of open(), since
        // it means a -w output file gets created as the real user, not
        // root, and doesn't need a manual chown to read back afterward.
        if (auto err = drop_privileges_if_root()) {
            return "failed to drop root privileges after opening the capture handle: " + *err;
        }

        datalink_ = pcap_datalink(handle_);
        if (!is_supported_datalink(datalink_)) {
            return std::string("unsupported datalink type on ") + device_ + ": " +
                   pcap_datalink_val_to_name(datalink_) + " (" +
                   pcap_datalink_val_to_description(datalink_) + ")";
        }

        if (options.filter_expr) {
            bpf_program program{};
            if (auto err = compile_filter(handle_, *options.filter_expr, &program)) {
                return "invalid filter '" + *options.filter_expr + "': " + *err;
            }
            if (pcap_setfilter(handle_, &program) == -1) {
                std::string err = std::string("pcap_setfilter failed: ") + pcap_geterr(handle_);
                pcap_freecode(&program);
                return err;
            }
            pcap_freecode(&program);  // bytecode is copied into the kernel by pcap_setfilter
        }

        if (options.pcapng_output_path) {
            if (auto err = open_pcapng_writer(*options.pcapng_output_path)) return err;
        }

        return std::nullopt;
    }

    // pcap_loop() blocks in a read/poll waiting for the next packet, so
    // a plain "stop requested" flag wouldn't unblock it promptly.
    // pcap_breakloop() is documented as signal-safe and is what
    // actually interrupts that wait. Replay mode has no handle to
    // breakloop, so it's interrupted via g_replay_stop_flag instead --
    // both are armed here so one signal handler covers either mode.
    void install_signal_handlers() {
        detail::g_capture_handle = handle_;
        detail::g_replay_stop_flag = &replay_stop_requested_;
        std::signal(SIGINT, detail::handle_stop_signal);
        std::signal(SIGTERM, detail::handle_stop_signal);
    }

    void request_stop() {
        if (handle_ != nullptr) pcap_breakloop(handle_);
        replay_stop_requested_.store(true);
    }

    // True once a stop has been explicitly requested - via
    // request_stop() or an external SIGINT/SIGTERM (the signal handler
    // sets the same flag). Lets a frontend tell "the producer stopped
    // because someone asked it to" apart from "the producer ran out of
    // data on its own" (replay reaching end-of-file), which call for
    // different UI behavior: the former should close the window, the
    // latter should leave it open so what's already loaded can still be
    // browsed.
    bool stop_requested() const { return replay_stop_requested_.load(); }

    // Must be called before close()/the destructor - pcap_stats()
    // needs a still-open handle. Safe to call after request_stop(),
    // since breakloop only stops pcap_loop(), it doesn't close handle_.
    // Always nullopt in replay mode (handle_ is never set there).
    std::optional<CaptureStats> stats() const {
        if (handle_ == nullptr) return std::nullopt;
        pcap_stat stat{};
        if (pcap_stats(handle_, &stat) == -1) return std::nullopt;
        return CaptureStats{stat.ps_recv, stat.ps_drop, stat.ps_ifdrop};
    }

    // Capture-thread side: copy each packet into the queue and return
    // immediately. No decoding, printing, or file I/O here - that's
    // every frontend's own consumer-side job.
    //
    // Live mode drops on backpressure (try_push, via capture_callback)
    // since a traffic spike can't be paused. Replay mode blocks instead
    // (push): a file has no real-time pressure forcing a drop, and
    // dropping from what's supposed to be a faithful replay of a fixed
    // historical record would defeat the point of replaying it.
    std::thread start_capture_thread(CaptureQueue& queue) {
        if (is_replay_) {
            return std::thread([this, &queue] {
                queue.push(to_captured_packet(std::move(*first_replay_packet_)));
                while (!replay_stop_requested_.load()) {
                    auto record = replay_reader_->next_packet();
                    if (!record) break;
                    if (!queue.push(to_captured_packet(std::move(*record)))) break;
                }
                queue.stop();
            });
        }
        return std::thread([this, &queue] {
            pcap_loop(handle_, /*count=*/-1, capture_callback,
                      reinterpret_cast<unsigned char*>(&queue));
            queue.stop();
        });
    }

    void close() {
        if (handle_ != nullptr) {
            pcap_close(handle_);
            handle_ = nullptr;
        }
        if (all_devices_ != nullptr) {
            pcap_freealldevs(all_devices_);
            all_devices_ = nullptr;
        }
        if (pcapng_file_ != nullptr) {
            std::fclose(pcapng_file_);
            pcapng_file_ = nullptr;
        }
        if (replay_file_ != nullptr) {
            std::fclose(replay_file_);
            replay_file_ = nullptr;
        }
    }

    pcap_t* handle() const { return handle_; }
    const std::string& device() const { return device_; }
    int datalink() const { return datalink_; }
    bool is_replay() const { return is_replay_; }
    pcapng::Writer* pcapng_writer() { return pcapng_writer_ ? &*pcapng_writer_ : nullptr; }

private:
    static void capture_callback(unsigned char* user, const pcap_pkthdr* header,
                                  const unsigned char* raw) {
        auto* queue = reinterpret_cast<CaptureQueue*>(user);
        CapturedPacket packet;
        packet.ts_sec = static_cast<std::uint32_t>(header->ts.tv_sec);
        packet.ts_usec = static_cast<std::uint32_t>(header->ts.tv_usec);
        packet.original_len = header->len;
        packet.data.assign(raw, raw + header->caplen);
        queue->try_push(std::move(packet));
    }

    static CapturedPacket to_captured_packet(pcapng::PacketRecord&& record) {
        CapturedPacket packet;
        packet.ts_sec = static_cast<std::uint32_t>(record.timestamp_us / 1'000'000ULL);
        packet.ts_usec = static_cast<std::uint32_t>(record.timestamp_us % 1'000'000ULL);
        packet.original_len = record.original_len;
        packet.data = std::move(record.data);
        return packet;
    }

    std::optional<std::string> open_pcapng_writer(const std::string& path) {
        pcapng_file_ = std::fopen(path.c_str(), "wb");
        if (pcapng_file_ == nullptr) {
            return "failed to open " + path + " for writing: " + std::strerror(errno);
        }
        pcapng_writer_.emplace(pcapng_file_);
        pcapng_writer_->write_section_header();
        pcapng_writer_->write_interface_description(65535,
                                                      static_cast<std::uint16_t>(datalink_));
        return std::nullopt;
    }

    std::optional<std::string> open_replay(const std::string& path,
                                            const std::optional<std::string>& pcapng_output_path) {
        replay_file_ = std::fopen(path.c_str(), "rb");
        if (replay_file_ == nullptr) {
            return "failed to open " + path + " for reading: " + std::strerror(errno);
        }

        replay_reader_.emplace(replay_file_);
        // Reading the first packet is also what makes the reader consume
        // the SHB/IDB blocks that precede it, which is what populates
        // link_type() below - there's no separate "just read the
        // header" step, so the packet itself is kept, not discarded.
        first_replay_packet_ = replay_reader_->next_packet();
        if (!first_replay_packet_) {
            return "no packets found in " + path + " (empty, or not a valid pcapng file)";
        }

        auto link_type = replay_reader_->link_type();
        if (!link_type || !is_supported_datalink(static_cast<int>(*link_type))) {
            return "unsupported or missing link type in " + path;
        }

        datalink_ = static_cast<int>(*link_type);
        device_ = path;
        is_replay_ = true;

        if (pcapng_output_path) {
            if (auto err = open_pcapng_writer(*pcapng_output_path)) return err;
        }

        return std::nullopt;
    }

    pcap_t* handle_ = nullptr;
    pcap_if_t* all_devices_ = nullptr;
    std::string device_;
    int datalink_ = 0;
    std::FILE* pcapng_file_ = nullptr;
    std::optional<pcapng::Writer> pcapng_writer_;

    bool is_replay_ = false;
    std::FILE* replay_file_ = nullptr;
    std::optional<pcapng::Reader> replay_reader_;
    std::optional<pcapng::PacketRecord> first_replay_packet_;
    std::atomic<bool> replay_stop_requested_{false};
};

}  // namespace wireframe