use anyhow::{Context, Result};
use clap::{Parser, Subcommand};
use std::net::{IpAddr, Ipv4Addr};
use tracing::{error, info, instrument, warn};
mod client_server;
mod config;
mod connection_pool;
mod forward;
mod proxy_server;
mod scanner;
mod system;
mod web;
mod wol;
#[cfg(test)]
mod test_support;
#[derive(Parser, Debug)]
#[command(author, version, about)]
pub struct Cli {
#[command(subcommand)]
command: Commands,
}
#[derive(Subcommand, Debug)]
pub enum Commands {
Send(SendArgs),
ProxyServer(ServeArgs),
ClientServer(ClientServerArgs),
}
#[derive(Parser, Debug)]
#[command()]
pub struct ServeArgs {
#[arg(
short,
long,
default_value_t = 3000,
help_heading = "Proxy Server Options"
)]
port: u16,
}
#[derive(Parser, Debug)]
#[command()]
pub struct ClientServerArgs {
#[arg(
short,
long,
default_value_t = 3001,
help_heading = "Client Server Options"
)]
port: u16,
}
#[derive(Parser, Debug)]
#[command()]
pub struct SendArgs {
mac: String,
#[arg(short, long)]
broadcast: Option<Ipv4Addr>,
#[arg(short, long, default_value_t = 9)]
port: u16,
#[arg(short = 'n', long, default_value_t = 3)]
count: u32,
#[arg(long, value_name = "IP")]
check_ip: Option<IpAddr>,
#[arg(long, default_value_t = 22)]
check_tcp_port: u16,
#[arg(long, default_value_t = 90)]
wait_secs: u64,
#[arg(long, default_value_t = 1000)]
interval_ms: u64,
#[arg(long, default_value_t = 700)]
connect_timeout_ms: u64,
}
#[tokio::main]
#[instrument(name = "wakezilla_main", skip_all)]
async fn main() -> Result<()> {
let env_filter =
tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| "info".into());
tracing_subscriber::fmt()
.with_writer(std::io::stdout)
.with_env_filter(env_filter)
.init();
let config = config::Config::from_env().unwrap_or_else(|e| {
warn!(
"Failed to load configuration from environment: {} - using defaults",
e
);
Default::default()
});
info!(
"Using configuration: server_proxy_port={}, server_client_port={}, wol_default_port={}, machines_db_path={}",
config.server.proxy_port, config.server.client_port, config.wol.default_port, config.storage.machines_db_path
);
let cli = Cli::parse();
match cli.command {
Commands::Send(args) => {
handle_send_command(args, &config)?;
}
Commands::ProxyServer(_args) => {
if let Err(e) = proxy_server::start(config.server.proxy_port).await {
error!("Proxy server error: {}", e);
std::process::exit(1);
}
}
Commands::ClientServer(_args) => {
if let Err(e) = client_server::start(config.server.client_port).await {
error!("Client server error: {}", e);
std::process::exit(1);
}
}
}
Ok(())
}
#[instrument(name = "handle_send_command", skip(args, config))]
fn handle_send_command(args: SendArgs, config: &config::Config) -> Result<()> {
info!("Processing WOL send command");
let mac = wol::parse_mac(&args.mac).context("Failed to parse MAC address")?;
let bcast = args
.broadcast
.unwrap_or(config.get_default_broadcast_addr());
match tokio::runtime::Handle::try_current() {
Ok(handle) => {
let result = handle.block_on(async {
wol::send_packets(&mac, bcast, args.port, args.count, config)
.await
.context("Failed to send WOL packets")?;
info!(
"Sent WOL magic packet to {:02x}:{:02x}:{:02x}:{:02x}:{:02x}:{:02x} via {}:{}",
mac[0], mac[1], mac[2], mac[3], mac[4], mac[5], bcast, args.port
);
if let Some(ip) = args.check_ip {
info!("Performing post-WOL reachability check for {}", ip);
if !wol::check_host(
ip,
args.check_tcp_port,
args.wait_secs,
args.interval_ms,
args.connect_timeout_ms,
config,
) {
anyhow::bail!(
"Host {}:{} did not become reachable within {} seconds",
ip,
args.check_tcp_port,
args.wait_secs
);
}
info!("Host {}:{} is now reachable", ip, args.check_tcp_port);
}
Ok(())
});
match result {
Ok(_) => Ok(()),
Err(e) => Err(e),
}
}
Err(_) => Err(anyhow::anyhow!("No runtime context available")),
}
}