use std::collections::HashSet; use std::future::Future; use std::io::Write; use std::net::{IpAddr, SocketAddr}; use std::path::PathBuf; use axum_server::tls_openssl::OpenSSLAcceptor; use clap::{Args, Parser, Subcommand, ValueEnum}; use futures::{stream::FuturesUnordered, StreamExt}; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; use one_kvm::auth::{SessionStore, TwoFactorService, UserStore}; use one_kvm::config; use one_kvm::db::open_database_pool; use one_kvm::platform::PlatformCapabilities; use one_kvm::runtime::{RuntimeBuilder, WebConfigOverrides}; use one_kvm::state::ShutdownAction; use one_kvm::utils::bind_tcp_listener; use one_kvm::utils::tls::server_config_from_pem_file; #[derive(Debug, Clone, Copy, Default, ValueEnum)] enum LogLevel { Error, Warn, #[default] Info, Debug, Trace, } #[derive(Parser, Debug)] #[command(name = "one-kvm")] #[command(version, about = "A open and lightweight IP-KVM solution", long_about = None)] struct CliArgs { /// User management commands #[command(subcommand)] command: Option, /// Listen address (overrides database config) #[arg(short = 'a', long, value_name = "ADDRESS")] address: Option, /// HTTP port (overrides database config) #[arg( short = 'p', long = "port", visible_alias = "http-port", value_name = "PORT" )] http_port: Option, /// HTTPS port (overrides database config) #[arg(long, value_name = "PORT")] https_port: Option, /// Enable HTTPS (overrides database config) #[arg(long)] enable_https: bool, /// Path to SSL certificate file (generates self-signed if not provided) #[arg(long, value_name = "FILE", requires = "ssl_key")] ssl_cert: Option, /// Path to SSL private key file #[arg(long, value_name = "FILE", requires = "ssl_cert")] ssl_key: Option, /// Data directory path #[arg(short = 'd', long, value_name = "DIR")] data_dir: Option, /// Log level (error, warn, info, debug, trace) #[arg(short = 'l', long, value_name = "LEVEL", default_value = "info")] log_level: LogLevel, } #[derive(Subcommand, Debug)] enum CliCommand { /// Manage local users User(UserCommand), } #[derive(Args, Debug)] struct UserCommand { #[command(subcommand)] action: UserAction, } #[derive(Subcommand, Debug)] enum UserAction { /// Set password for the single local user (interactive terminal prompt) SetPassword, /// Disable TOTP for the single local user DisableTotp, } #[tokio::main] async fn main() -> anyhow::Result<()> { let args = CliArgs::parse(); init_logging(args.log_level); tracing::info!("Starting One-KVM v{}", env!("CARGO_PKG_VERSION")); let platform = PlatformCapabilities::current(); tracing::info!( "Platform mode: {:?} ({})", platform.mode, platform.mode_label ); let data_dir = args.data_dir.clone().unwrap_or_else(get_data_dir); tracing::info!("Data directory: {}", data_dir.display()); if let Some(command) = args.command { run_cli_command(command, data_dir).await?; return Ok(()); } let overrides = WebConfigOverrides { address: args.address, http_port: args.http_port, https_port: args.https_port, enable_https: args.enable_https, ssl_cert: args.ssl_cert, ssl_key: args.ssl_key, }; let mut runtime = RuntimeBuilder::new(data_dir.clone()) .with_web_overrides(overrides) .build() .await?; let config = runtime.config(); let state = runtime.state().clone(); let bind_ips = resolve_bind_addresses(&config.web)?; let scheme = if config.web.https_enabled { "https" } else { "http" }; let bind_port = if config.web.https_enabled { config.web.https_port } else { config.web.http_port }; for ip in &bind_ips { let addr = SocketAddr::new(*ip, bind_port); tracing::info!("Server will listen on: {}://{}", scheme, addr); } let app = runtime.router(); let listeners = bind_tcp_listeners(&bind_ips, bind_port)?; let shutdown_signal = { let shutdown_tx = state.shutdown_tx.clone(); let mut shutdown_rx = shutdown_tx.subscribe(); async move { tokio::select! { result = shutdown_signal() => { if let Err(e) = result { tracing::error!("Failed while waiting for shutdown signal: {}", e); } tracing::info!("SIGINT or SIGTERM received"); ShutdownAction::Exit } request = shutdown_rx.recv() => { match request { Ok(action) => { tracing::info!("Shutdown request received: {:?}", action); action } Err(e) => { tracing::warn!("Shutdown request channel closed: {}", e); ShutdownAction::Exit } } } } } }; let shutdown_action = if config.web.https_enabled { let tls_config = if let (Some(cert_path), Some(key_path)) = (&config.web.ssl_cert_path, &config.web.ssl_key_path) { server_config_from_pem_file(cert_path, key_path).await? } else { let cert_dir = data_dir.join("certs"); let cert_path = cert_dir.join("server.crt"); let key_path = cert_dir.join("server.key"); if !cert_path.exists() || !key_path.exists() { tracing::info!("Generating new self-signed TLS certificate"); let cert = generate_self_signed_cert()?; tokio::fs::create_dir_all(&cert_dir).await?; tokio::fs::write(&cert_path, cert.cert.pem()).await?; tokio::fs::write(&key_path, cert.signing_key.serialize_pem()).await?; } else { tracing::info!("Using existing TLS certificate from {}", cert_dir.display()); } server_config_from_pem_file(&cert_path, &key_path).await? }; let servers = FuturesUnordered::new(); for listener in listeners { let local_addr = listener.local_addr()?; tracing::info!("Starting HTTPS server on {}", local_addr); let server = axum_server::from_tcp(listener)? .acceptor(OpenSSLAcceptor::new(tls_config.clone())) .serve(app.clone().into_make_service()); servers.push(server); } run_servers_until_shutdown(servers, shutdown_signal, "HTTPS").await } else { let servers = FuturesUnordered::new(); for listener in listeners { let local_addr = listener.local_addr()?; tracing::info!("Starting HTTP server on {}", local_addr); let listener = tokio::net::TcpListener::from_std(listener)?; let server = axum::serve(listener, app.clone()); servers.push(async move { server.await }); } run_servers_until_shutdown(servers, shutdown_signal, "HTTP").await }; runtime.shutdown().await; tracing::info!("Server shutdown complete"); if let ShutdownAction::Restart { exe_path } = shutdown_action { restart_current_process(exe_path)?; } Ok(()) } fn init_logging(level: LogLevel) { let app_level = match level { LogLevel::Error => "error", LogLevel::Warn => "warn", LogLevel::Info => "info", LogLevel::Debug => "debug", LogLevel::Trace => "trace", }; let env_filter = tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| app_level.into()); if let Err(err) = tracing_subscriber::registry() .with(env_filter) .with(tracing_subscriber::fmt::layer()) .try_init() { eprintln!("failed to initialize tracing: {}", err); } } fn get_data_dir() -> PathBuf { if let Ok(path) = std::env::var("ONE_KVM_DATA_DIR") { return PathBuf::from(path); } #[cfg(windows)] { if let Ok(exe_path) = std::env::current_exe() { if let Some(exe_dir) = exe_path.parent() { return exe_dir.join("one-kvm"); } } return std::env::current_dir() .map(|dir| dir.join("one-kvm")) .unwrap_or_else(|_| PathBuf::from("one-kvm")); } #[cfg(not(windows))] PathBuf::from("/etc/one-kvm") } #[cfg(unix)] async fn shutdown_signal() -> anyhow::Result<()> { use tokio::signal::unix::{signal, SignalKind}; let mut terminate = signal(SignalKind::terminate())?; tokio::select! { result = tokio::signal::ctrl_c() => result?, _ = terminate.recv() => {}, } Ok(()) } #[cfg(not(unix))] async fn shutdown_signal() -> anyhow::Result<()> { tokio::signal::ctrl_c().await?; Ok(()) } async fn run_servers_until_shutdown( mut servers: FuturesUnordered, shutdown_signal: impl Future, protocol: &'static str, ) -> ShutdownAction where F: Future> + Send, E: std::fmt::Display, { tokio::select! { action = shutdown_signal => { action } result = servers.next() => { if let Some(Err(e)) = result { tracing::error!("{} server error: {}", protocol, e); } ShutdownAction::Exit } } } fn restart_current_process(exe_path: Option) -> anyhow::Result<()> { let exe = exe_path.unwrap_or(std::env::current_exe()?); let args: Vec = std::env::args().skip(1).collect(); tracing::info!("Restarting: {:?} {:?}", exe, args); #[cfg(unix)] { use std::os::unix::process::CommandExt; let err = std::process::Command::new(&exe).args(&args).exec(); Err(anyhow::anyhow!("Failed to restart: {}", err)) } #[cfg(not(unix))] { std::process::Command::new(&exe).args(&args).spawn()?; std::process::exit(0); } } async fn run_cli_command(command: CliCommand, data_dir: PathBuf) -> anyhow::Result<()> { let db = open_database_pool(&data_dir).await?; let users = UserStore::new(db.clone_pool()); let two_factor = TwoFactorService::new(db.clone_pool()); let sessions = SessionStore::new(0); match command { CliCommand::User(user) => { run_user_action(user.action, &users, &sessions, &two_factor).await } } } async fn run_user_action( action: UserAction, users: &UserStore, sessions: &SessionStore, two_factor: &TwoFactorService, ) -> anyhow::Result<()> { match action { UserAction::SetPassword => set_user_password(users, sessions).await, UserAction::DisableTotp => disable_user_totp(users, two_factor).await, } } async fn disable_user_totp(users: &UserStore, two_factor: &TwoFactorService) -> anyhow::Result<()> { let user = users.single_user().await?.ok_or_else(|| { anyhow::anyhow!("No local user exists yet; complete setup in the web UI first.") })?; if two_factor.disable_without_code(&user.id).await? { println!("TOTP disabled for user '{}'.", user.username); } else { println!("TOTP is already disabled for user '{}'.", user.username); } Ok(()) } async fn set_user_password(users: &UserStore, sessions: &SessionStore) -> anyhow::Result<()> { let user = users.single_user().await?.ok_or_else(|| { anyhow::anyhow!("No local user exists yet; complete setup in the web UI first.") })?; let new_password = read_new_password_interactive()?; if new_password.len() < 4 { anyhow::bail!("Password must be at least 4 characters"); } users.update_password(&user.id, &new_password).await?; let revoked = sessions.delete_all().await?; tracing::info!( "Password updated for user '{}' and {} sessions revoked", user.username, revoked ); println!( "Password updated for user '{}' (revoked {} sessions).", user.username, revoked ); Ok(()) } fn read_new_password_interactive() -> anyhow::Result { let once = |label: &str| -> anyhow::Result { print!("{}", label); std::io::stdout().flush()?; let mut line = String::new(); std::io::stdin().read_line(&mut line)?; let s = line.trim_end_matches(['\r', '\n']).to_string(); if s.is_empty() { anyhow::bail!("Password cannot be empty"); } Ok(s) }; let a = once("New password: ")?; let b = once("Confirm password: ")?; if a != b { anyhow::bail!("Passwords do not match"); } Ok(a) } fn resolve_bind_addresses(web: &config::WebConfig) -> anyhow::Result> { let raw_addrs = if !web.bind_addresses.is_empty() { web.bind_addresses.as_slice() } else { std::slice::from_ref(&web.bind_address) }; let mut seen = HashSet::new(); let mut addrs = Vec::new(); for addr in raw_addrs { let ip: IpAddr = addr .parse() .map_err(|_| anyhow::anyhow!("Invalid bind address: {}", addr))?; if seen.insert(ip) { addrs.push(ip); } } Ok(addrs) } fn bind_tcp_listeners(addrs: &[IpAddr], port: u16) -> anyhow::Result> { let mut listeners = Vec::new(); for ip in addrs { let addr = SocketAddr::new(*ip, port); match bind_tcp_listener(addr) { Ok(listener) => listeners.push(listener), Err(err) => { tracing::warn!("Failed to bind {}: {}", addr, err); } } } if listeners.is_empty() { anyhow::bail!("Failed to bind any addresses on port {}", port); } Ok(listeners) } fn generate_self_signed_cert() -> anyhow::Result> { use rcgen::generate_simple_self_signed; let subject_alt_names = vec![ "localhost".to_string(), "127.0.0.1".to_string(), "::1".to_string(), ]; let certified_key = generate_simple_self_signed(subject_alt_names)?; Ok(certified_key) }