hai/src/main.rs

/*
 *  SPDX-FileCopyrightText: 2025 April Faye John <april.john@denic.de> & Contributors
 *  SPDX-License-Identifier: AGPL-3.0-or-later
 */

#![cfg_attr(not(debug_assertions), windows_subsystem = "windows")] // hide console window on Windows in release

mod bininfo;
mod gui;
mod pman;
mod quic;

use std::fs::File;
use std::net::SocketAddr;
use std::path::{Path, PathBuf};
use std::str::FromStr;
use bunt::println;
use clap::{Parser, Subcommand};
use config::{Config, FileFormat, Source};
use log::{debug, error, info};
use pman::{init_process_manager, ProcessCommand, ProcessManager};
use shadow_rs::shadow;
use std::sync::{Arc, Mutex, OnceLock};
use anyhow::{bail, Context};
use quinn::crypto::rustls::QuicServerConfig;
use rustls::pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer};
use tokio::{signal, sync::mpsc::Sender};
use trust_dns_resolver::config::{ResolverConfig, ResolverOpts};
use trust_dns_resolver::{Resolver, TokioAsyncResolver};
use crate::quic::ALPN_QUIC_HTTP;
use crate::quic::server::handle_connection;

shadow!(build);

static PMAN_SENDER: OnceLock<Sender<ProcessCommand>> = OnceLock::new();

/// Simple program to greet a person
#[derive(Parser, Debug)]
#[command(version, about, long_about = None)]
struct Args {
    /// Config File to load
    #[arg(short, long)]
    config: Option<String>,

    /// Number of times to greet
    //#[arg(short, long, default_value_t = 1)]
    //count: u8,

    #[command(subcommand)]
    command: Commands,
}

#[derive(Debug, Subcommand)]
enum Commands {
    #[command(about = "Start client without GUI")]
    CliClient,
    #[command(about = "List compile time backed info to audit binary")]
    Shadow,
    #[command(about = "Start client as GUI")]
    GuiClient,
    Devtest,
    Server {
        /// file to log TLS keys to for debugging
        #[clap(long = "keylog")]
        keylog: bool,
        /// TLS private key in PEM format
        #[clap(short = 'k', long = "key", requires = "cert")]
        key: Option<PathBuf>,
        /// TLS certificate in PEM format
        #[clap(short = 'c', long = "cert", requires = "key")]
        cert: Option<PathBuf>,
        /// Enable stateless retries
        #[clap(long = "stateless-retry")]
        stateless_retry: bool,
        /// Address to listen on
        #[clap(long = "listen", default_value = "[::1]:4433")]
        listen: SocketAddr,
        /// Client address to block
        #[clap(long = "block")]
        block: Option<SocketAddr>,
        /// Maximum number of concurrent connections to allow
        #[clap(long = "connection-limit")]
        connection_limit: Option<usize>,
    }
}

fn config() -> &'static Config {
    static CONFIG: OnceLock<Config> = OnceLock::new();
    CONFIG.get_or_init(|| {
        Config::builder()
            .add_source(config::Environment::with_prefix("HAI").separator("_"))
            .add_source(get_config_file_source())
            .build()
            .unwrap()
    })
}

static CONFIG_FILE: OnceLock<Option<String>> = OnceLock::new();

fn get_config_file_source() -> impl Source {
    let file = CONFIG_FILE.get();
    let default_dir = dirs::config_dir().unwrap();
    let file_buf = default_dir.join("hai").join("config.toml");
    println!("{}", file_buf.to_string_lossy());
    let file_content = std::fs::read_to_string(file_buf).unwrap();
    config::File::from_str(&*file_content, FileFormat::Toml)
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {

    env_logger::init();
    let args = Args::parse();
    CONFIG_FILE.get_or_init(|| {
        args.config
    });

    rustls::crypto::aws_lc_rs::default_provider().install_default().expect("Couldnt install rustls crypto into process");

    let _ = init_process_manager();

    match args.command {
        Commands::CliClient => {
            let host = config().get_string("remote_host")?;
            let port = config().get::<u16>("remote_port")?;
            let resolver = TokioAsyncResolver::tokio(
                ResolverConfig::quad9_tls(),
                ResolverOpts::default(),
            );
            let response = resolver.lookup_ip(&host).await?;
            let ip_addr = response.iter().next().expect("no addresses returned!");
            let ip_sock = SocketAddr::new(ip_addr, port);
            let client_endpoint = quic::client::make_client_endpoint(ip_sock.clone(), None).unwrap();
            let client = client_endpoint.connect(ip_sock, &host)?.await?;
            println!("[client] connected: addr={}", client.remote_address());
        }

        Commands::Shadow => {
            bininfo::print_info();
            return Ok(());
        }

        Commands::GuiClient => {
            let res = gui::gui_main();
            if let Err(e) = res {
                println!("{}", e);
            }
            return Ok(());
        }

        Commands::Devtest => {
            PMAN_SENDER.get().unwrap().send(ProcessCommand::SpawnShellCmd {
                cmd: "sleep 100 && echo hello meow".to_string(),
            }).await.expect("TODO: panic message");
            tokio::select! {
                _ = signal::ctrl_c() => {
                    info!("Ctrl-c received");
                    let sender = PMAN_SENDER.get().unwrap();
                    let _ = sender.send(ProcessCommand::Shutdown).await;
                }
            }
        }

        Commands::Server {
            keylog, key, cert, stateless_retry, listen, block, connection_limit
        } => {
            let (certs, key) = if let (Some(key_path), Some(cert_path)) = (&key, &cert) {
                let key = std::fs::read(key_path).context("failed to read private key")?;
                let key = if key_path.extension().is_some_and(|x| x == "der") {
                    PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(key))
                } else {
                    rustls_pemfile::private_key(&mut &*key)
                        .context("malformed PKCS #1 private key")?
                        .ok_or_else(|| anyhow::Error::msg("no private keys found"))?
                };
                let cert_chain = std::fs::read(cert_path).context("failed to read certificate chain")?;
                let cert_chain = if cert_path.extension().is_some_and(|x| x == "der") {
                    vec![CertificateDer::from(cert_chain)]
                } else {
                    rustls_pemfile::certs(&mut &*cert_chain)
                        .collect::<Result<_, _>>()
                        .context("invalid PEM-encoded certificate")?
                };

                (cert_chain, key)
            } else {
                let default_dir = dirs::data_dir().unwrap();
                let path = default_dir.join("hai-server");
                let cert_path = path.join("cert.der");
                let key_path = path.join("key.der");
                let (cert, key) = match std::fs::read(&cert_path).and_then(|x| Ok((x, std::fs::read(&key_path)?))) {
                    Ok((cert, key)) => (
                        CertificateDer::from(cert),
                        PrivateKeyDer::try_from(key).map_err(anyhow::Error::msg)?,
                    ),
                    Err(ref e) if e.kind() == std::io::ErrorKind::NotFound => {
                        info!("generating self-signed certificate");
                        let cert = rcgen::generate_simple_self_signed(vec!["localhost".into()]).unwrap();
                        let key = PrivatePkcs8KeyDer::from(cert.key_pair.serialize_der());
                        let cert = cert.cert.into();
                        std::fs::create_dir_all(path).context("failed to create certificate directory")?;
                        std::fs::write(&cert_path, &cert).context("failed to write certificate")?;
                        std::fs::write(&key_path, key.secret_pkcs8_der())
                            .context("failed to write private key")?;
                        (cert, key.into())
                    }
                    Err(e) => {
                        bail!("failed to read certificate: {}", e);
                    }
                };

                (vec![cert], key)
            };
            let mut server_crypto = rustls::ServerConfig::builder()
                //TODO change to auth
                .with_no_client_auth()
                .with_single_cert(certs, key)?;
            server_crypto.alpn_protocols = ALPN_QUIC_HTTP.iter().map(|&x| x.into()).collect();
            if keylog {
                server_crypto.key_log = Arc::new(rustls::KeyLogFile::new());
            }

            let mut server_config =
                quinn::ServerConfig::with_crypto(Arc::new(QuicServerConfig::try_from(server_crypto)?));
            let transport_config = Arc::get_mut(&mut server_config.transport).unwrap();
            transport_config.max_concurrent_uni_streams(0_u8.into());

            let endpoint = quinn::Endpoint::server(server_config, listen)?;
            eprintln!("listening on {}", endpoint.local_addr()?);

            while let Some(conn) = endpoint.accept().await {
                if connection_limit
                    .is_some_and(|n| endpoint.open_connections() >= n)
                {
                    info!("refusing due to open connection limit");
                    conn.refuse();
                } else if Some(conn.remote_address()) == block {
                    info!("refusing blocked client IP address");
                    conn.refuse();
                } else if stateless_retry && !conn.remote_address_validated() {
                    info!("requiring connection to validate its address");
                    conn.retry().unwrap();
                } else {
                    info!("accepting connection");
                    let fut = handle_connection(conn);
                    tokio::spawn(async move {
                        if let Err(e) = fut.await {
                            error!("connection failed: {reason}", reason = e.to_string())
                        }
                    });
                }
            }
        }
    };

    //handling anything here for gui wont work
    println!("exit");
    Ok(())
}