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 {
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);
let mut warned_loglevel_deprecated = false;
let mut warned_daemon_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_daemon_deprecated {
warned_daemon_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 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 {
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; }
"-t" | "--fake-ttl" => { fake = true; fake_ttl = take_value(&mut args, argv)?; }
"-a" | "--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!(" -t, --fake-ttl <u8> Override ttl of fake clienthello (default: {DEFAULT_FAKE_TTL})");
println!(" -a, --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(())
}