path: root/src/opt.rs

// SPDX-FileCopyrightText: 2026 Dilluti0n <[email protected]>
// SPDX-License-Identifier: GPL-3.0-or-later

use anyhow::{Result, anyhow, Context};
use std::sync::OnceLock;

use crate::log;

use log::LogLevel;

#[derive(Copy, Clone)]
pub struct Segment(pub u32, pub u32);

impl std::fmt::Display for Segment {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        let end = if self.1 == u32::MAX { "end".to_string() } else { self.1.to_string() };
        write!(f, "[{},{})", self.0, end)
    }
}

impl std::fmt::Debug for Segment {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "{self}")
    }
}

pub struct SegmentOrder {
    raw: String,
    segments: Vec<Segment>
}

impl SegmentOrder {
    /// Parse 5,1,0,3 to (5, u32::MAX), (1, 3), (0, 1), (3, 5).
    pub fn new(s: &str) -> Result<Self> {
        let mut points: Vec<u32> = s
            .split(',')
            .map(|x| x.trim().parse::<u32>())
            .collect::<std::result::Result<_, _>>()
            .with_context(|| format!("--segment-order: invalid value '{s}'"))?;

        if points.is_empty() {
            return Err(anyhow!("--segment-order: empty"));
        }

        let order = points.clone();
        points.sort_unstable();
        points.dedup();

        if !points.contains(&0) {
            return Err(anyhow!("--segment-order: must contain 0"));
        }

        let sorted_ranges: Vec<Segment> = points.windows(2)
            .map(|w| Segment(w[0], w[1]))
            .chain(std::iter::once(Segment(*points.last().unwrap(), u32::MAX)))
            .collect();

        let segments = order.iter()
            .map(|&p| {
                sorted_ranges.iter()
                    .find(|&&Segment(start, _)| start == p)
                    .copied()
                    .ok_or_else(|| anyhow!("--segment-order: internal error"))
            })
            .collect::<Result<Vec<_>>>()?;

        Ok(Self {
            raw: s.to_string(),
            segments,
        })
    }

    pub fn segments(&self) -> &[Segment] {
        &self.segments
    }
}

impl std::fmt::Display for SegmentOrder {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "{} (", self.raw)?;
        for (i, seg) in self.segments.iter().enumerate() {
            if i > 0 { write!(f, ", ")?; }
            write!(f, "{seg}")?;
        }
        write!(f, ")")
    }
}

static OPT_DAEMON: OnceLock<bool> = OnceLock::new();
static OPT_LOG_LEVEL: OnceLock<LogLevel> = OnceLock::new();
static OPT_NO_SPLASH: OnceLock<bool> = OnceLock::new();
static OPT_FAKE: OnceLock<bool> = OnceLock::new();
static OPT_FAKE_TTL: OnceLock<u8> = OnceLock::new();
static OPT_FAKE_AUTOTTL: OnceLock<bool> = OnceLock::new();
static OPT_FAKE_BADSUM: OnceLock<bool> = OnceLock::new();
static OPT_DELAY_MS: OnceLock<u64> = OnceLock::new();
#[cfg(target_os = "linux")] static OPT_QUEUE_NUM: OnceLock<u16> = OnceLock::new();
#[cfg(target_os = "linux")] static OPT_NFT_COMMAND: OnceLock<String> = OnceLock::new();
static OPT_SEGMENT_ORDER: OnceLock<SegmentOrder> = OnceLock::new();

const DEFAULT_DAEMON: bool = false;
#[cfg(debug_assertions)]      const DEFAULT_LOG_LEVEL: LogLevel = LogLevel::Debug;
#[cfg(not(debug_assertions))] const DEFAULT_LOG_LEVEL: LogLevel = LogLevel::Warning;
const DEFAULT_NO_SPLASH: bool = false;
const DEFAULT_FAKE: bool = false;
const DEFAULT_FAKE_TTL: u8 = 8;
const DEFAULT_FAKE_AUTOTTL: bool = false;
const DEFAULT_FAKE_BADSUM: bool = false;
const DEFAULT_DELAY_MS: u64 = 0;
#[cfg(target_os = "linux")] const DEFAULT_QUEUE_NUM: u16 = 1;
#[cfg(target_os = "linux")] const DEFAULT_NFT_COMMAND: &str = "nft";
const DEFAULT_SEGMENT_ORDER: &str = "0,1";

pub struct Opt {
    daemon: bool,
    log_level: LogLevel,
    no_splash: bool,
    fake: bool,
    fake_ttl: u8,
    fake_autottl: bool,
    fake_badsum: bool,
    delay_ms: u64,
    #[cfg(target_os = "linux")] queue_num: u16,
    #[cfg(target_os = "linux")] nft_command: String,
    segment_order: SegmentOrder,
}

impl Opt {
    pub fn from_args() -> Result<Self> {
        let mut daemon = DEFAULT_DAEMON;
        let mut log_level     = DEFAULT_LOG_LEVEL;
        let mut delay_ms      = DEFAULT_DELAY_MS;
        let mut no_splash     = DEFAULT_NO_SPLASH;
        let mut fake          = DEFAULT_FAKE;
        let mut fake_ttl      = DEFAULT_FAKE_TTL;
        let mut fake_autottl  = DEFAULT_FAKE_AUTOTTL;
        let mut fake_badsum   = DEFAULT_FAKE_BADSUM;
        let mut segment_order = SegmentOrder::new(DEFAULT_SEGMENT_ORDER)?;

        #[cfg(target_os = "linux")]
        let mut queue_num: u16 = DEFAULT_QUEUE_NUM;
        #[cfg(target_os = "linux")]
        let mut nft_command = String::from(DEFAULT_NFT_COMMAND);

        let mut args = std::env::args().skip(1); // program name

        let mut warned_loglevel_deprecated = false;
        let mut warned_D_deprecated = false;

        while let Some(arg) = args.next() {
            let argv = arg.as_str();

            match argv {
                "-h" | "--help" => { usage(); std::process::exit(0); }
                "-d" | "-D" | "--daemon" => {
                    if argv == "--loglevel" && !warned_D_deprecated {
                        // FIXME(on release): remove this on v1.0.0
                        warned_D_deprecated = true;
                        eprintln!("Note: `{arg}' has been deprecated since v0.6.0 and planned to be removed on v1.0.0. Use `-d' instead.");
                    }
                    no_splash = true;
                    // if it is unchanged explicitly by argument, set it to info
                    if log_level == DEFAULT_LOG_LEVEL {
                        log_level = LogLevel::Info;
                    }
                    daemon = true;
                }
                "--delay-ms" => { delay_ms = take_value(&mut args, argv)?; }
                "--log-level" | "--loglevel" => {
                    if argv == "--loglevel" && !warned_loglevel_deprecated {
                        // FIXME(on release): remove this on v1.0.0
                        warned_loglevel_deprecated = true;
                        eprintln!("Note: `{arg}' has been deprecated since v0.1.1 and planned to be removed on v1.0.0. Use `--log-level' instead.");
                    }
                    log_level = take_value(&mut args, argv)?;
                }
                "--no-splash" => { no_splash = true; }

                "-o" | "--segment-order" => {
                    let s: String = take_value(&mut args, argv)?;
                    segment_order = SegmentOrder::new(&s)?;
                }

                "--fake" => { fake = true; }
                "--fake-ttl" => { fake = true; fake_ttl = take_value(&mut args, argv)?; }
                "--fake-autottl" => { fake = true; fake_autottl = true }
                "--fake-badsum" => { fake = true; fake_badsum = true }

                #[cfg(target_os = "linux")]
                "--queue-num" => { queue_num = take_value(&mut args, argv)?; }

                #[cfg(target_os = "linux")]
                "--nft-command" => { nft_command = take_value(&mut args, argv)?; }

                _ => { return Err(anyhow!("argument: unknown: {}", arg)); }
            }
        }

        Ok(Opt {
            daemon: daemon,
            log_level: log_level,
            no_splash: no_splash,
            segment_order: segment_order,
            fake: fake,
            fake_ttl: fake_ttl,
            fake_autottl: fake_autottl,
            fake_badsum: fake_badsum,
            delay_ms: delay_ms,
            #[cfg(target_os = "linux")] queue_num: queue_num,
            #[cfg(target_os = "linux")] nft_command: nft_command,
        })
    }

    pub fn set_opt(self) -> Result<InitializedOpts> {
        set_opt("OPT_DAEMON", &OPT_DAEMON, self.daemon)?;
        set_opt("OPT_LOG_LEVEL", &OPT_LOG_LEVEL, self.log_level)?;
        set_opt("OPT_NO_SPLASH", &OPT_NO_SPLASH, self.no_splash)?;

        set_opt("OPT_SEGMENT_ORDER", &OPT_SEGMENT_ORDER, self.segment_order)?;

        set_opt("OPT_DELAY_MS", &OPT_DELAY_MS, self.delay_ms)?;
        set_opt("OPT_FAKE", &OPT_FAKE, self.fake)?;
        set_opt("OPT_FAKE_TTL", &OPT_FAKE_TTL, self.fake_ttl)?;
        set_opt("OPT_FAKE_AUTOTTL", &OPT_FAKE_AUTOTTL, self.fake_autottl)?;
        set_opt("OPT_FAKE_BADSUM", &OPT_FAKE_BADSUM, self.fake_badsum)?;

        #[cfg(target_os = "linux")] set_opt("OPT_QUEUE_NUM", &OPT_QUEUE_NUM, self.queue_num)?;
        #[cfg(target_os = "linux")] set_opt("OPT_NFT_COMMAND", &OPT_NFT_COMMAND, self.nft_command)?;

        Ok(InitializedOpts)
    }
}

pub struct InitializedOpts;

impl InitializedOpts {
    pub fn log(&self) {
        crate::info!("OPT_DAEMON: {}", daemon());
        crate::info!("OPT_NO_SPLASH: {}", no_splash());
        crate::info!("OPT_LOG_LEVEL: {}", log_level());
        crate::info!("OPT_DELAY_MS: {}", delay_ms());
        crate::info!("OPT_FAKE: {}", fake());
        crate::info!("OPT_FAKE_TTL: {}", fake_ttl());
        crate::info!("OPT_FAKE_AUTOTTL: {}", fake_autottl());
        crate::info!("OPT_FAKE_BADSUM: {}", fake_badsum());
        #[cfg(target_os = "linux")]
        crate::info!("OPT_QUEUE_NUM: {}", queue_num());
        #[cfg(target_os = "linux")]
        crate::info!("OPT_NFT_COMMAND: {}", nft_command());
        crate::info!("OPT_SEGMENT_ORDER: {}", segment_order());
    }
}

pub fn daemon() -> bool {
    *OPT_DAEMON.get().unwrap_or(&DEFAULT_DAEMON)
}

pub fn no_splash() -> bool {
    *OPT_NO_SPLASH.get().unwrap_or(&DEFAULT_NO_SPLASH)
}

pub fn segment_order() -> &'static SegmentOrder {
    OPT_SEGMENT_ORDER.get().unwrap()
}

pub fn log_level() -> LogLevel {
    *OPT_LOG_LEVEL.get().unwrap_or(&DEFAULT_LOG_LEVEL)
}

pub fn fake() -> bool {
    *OPT_FAKE.get().unwrap_or(&DEFAULT_FAKE)
}

pub fn fake_ttl() -> u8 {
    *OPT_FAKE_TTL.get().unwrap_or(&DEFAULT_FAKE_TTL)
}

pub fn fake_autottl() -> bool {
    *OPT_FAKE_AUTOTTL.get().unwrap_or(&DEFAULT_FAKE_AUTOTTL)
}

pub fn fake_badsum() -> bool {
    *OPT_FAKE_BADSUM.get().unwrap_or(&DEFAULT_FAKE_BADSUM)
}

pub fn delay_ms() -> u64 {
    *OPT_DELAY_MS.get().unwrap_or(&DEFAULT_DELAY_MS)
}

#[cfg(target_os = "linux")]
pub fn queue_num() -> u16 {
    *OPT_QUEUE_NUM.get().unwrap_or(&DEFAULT_QUEUE_NUM)
}

#[cfg(target_os = "linux")]
pub fn nft_command() -> &'static str {
    OPT_NFT_COMMAND.get().map(String::as_str).unwrap_or(DEFAULT_NFT_COMMAND)
}

fn take_value<T, I>(args: &mut I, arg_name: &str) -> Result<T>
where
    T: std::str::FromStr,
    T::Err: std::error::Error + Send + Sync + 'static,
    I: Iterator<Item = String>,
{
    let raw = args
        .next()
        .ok_or_else(|| anyhow!("argument: missing value after {}", arg_name))?;
    raw.parse::<T>()
        .with_context(|| format!("argument: {}: invalid value '{}'", arg_name, raw))
}

fn usage() {
    println!("Usage: dpibreak [OPTIONS]\n");
    println!("Options:");
    println!("  -d, --daemon                            Run as daemon; kill `pidof dpibreak` to stop.");
    println!("  --delay-ms    <u64>                       (default: {DEFAULT_DELAY_MS})");
    #[cfg(target_os = "linux")]
    println!("  --queue-num   <u16>                       (default: {DEFAULT_QUEUE_NUM})");
    #[cfg(target_os = "linux")]
    println!("  --nft-command <string>                    (default: {DEFAULT_NFT_COMMAND})");

    println!("  --log-level   <debug|info|warning|error>  (default: {DEFAULT_LOG_LEVEL})");
    println!("  --no-splash                             Do not print splash messages\n");

    println!("  --fake                                  Enable fake clienthello injection");
    println!("  --fake-ttl    <u8>                      Override ttl of fake clienthello (default: {DEFAULT_FAKE_TTL})");
    println!("  --fake-autottl                          Override ttl of fake clienthello automatically");
    println!("  --fake-badsum                           Modifies the TCP checksum of the fake packet to an invalid value.");
    println!("");
    println!("  -o, --segment-order <u32,u32,...>       Byte offsets defining segment boundaries and transmission order.");
    println!("                                          Must include 0. (default: {DEFAULT_SEGMENT_ORDER})");
    println!("");

    println!("  -h, --help                              Show this help");
}

fn set_opt<T: std::fmt::Display>(
    name: &str,
    cell: &OnceLock<T>,
    value: T,
) -> Result<()> {
    cell.set(value).map_err(|_| anyhow!("{name} already initialized"))?;
    Ok(())
}