One-KVM/build.rs

use std::fs;
use std::path::Path;

fn main() {
    // Set BUILD_DATE environment variable for compile-time access
    // Use system time to avoid adding chrono as a build dependency
    let now = std::time::SystemTime::now();
    let duration = now.duration_since(std::time::UNIX_EPOCH).unwrap();
    let secs = duration.as_secs();

    // Convert Unix timestamp to date (simplified calculation)
    // Days since epoch
    let days = secs / 86400;
    // Calculate year, month, day from days since 1970-01-01
    let (year, month, day) = days_to_ymd(days as i64);
    let build_date = format!("{:04}-{:02}-{:02}", year, month, day);

    println!("cargo:rustc-env=BUILD_DATE={}", build_date);

    // Compile protobuf files for RustDesk protocol
    compile_protos();

    // Generate secrets module from secrets.toml
    generate_secrets();

    // Rerun if the script itself changes
    println!("cargo:rerun-if-changed=build.rs");
    println!("cargo:rerun-if-changed=protos/rendezvous.proto");
    println!("cargo:rerun-if-changed=protos/message.proto");
    println!("cargo:rerun-if-changed=secrets.toml");
}

/// Compile protobuf files using protobuf-codegen (same as RustDesk server)
fn compile_protos() {
    let out_dir = std::path::PathBuf::from(std::env::var("OUT_DIR").unwrap());
    let protos_dir = out_dir.join("protos");
    std::fs::create_dir_all(&protos_dir).unwrap();

    protobuf_codegen::Codegen::new()
        .pure()
        .out_dir(&protos_dir)
        .inputs(["protos/rendezvous.proto", "protos/message.proto"])
        .include("protos")
        .customize(protobuf_codegen::Customize::default().tokio_bytes(true))
        .run()
        .expect("Failed to compile protobuf files");

    // Generate mod.rs for the protos module
    let mod_content = r#"pub mod rendezvous;
pub mod message;
"#;
    std::fs::write(protos_dir.join("mod.rs"), mod_content).unwrap();
}

/// Generate secrets module from secrets.toml
///
/// This reads the secrets.toml file and generates a Rust module with
/// compile-time constants for sensitive configuration values.
fn generate_secrets() {
    let out_dir = std::env::var("OUT_DIR").unwrap();
    let dest_path = Path::new(&out_dir).join("secrets_generated.rs");

    // Default values if secrets.toml doesn't exist
    let mut rustdesk_public_server = String::new();
    let mut rustdesk_public_key = String::new();
    let mut rustdesk_relay_key = String::new();
    let mut ice_stun_server = String::new();
    let mut ice_turn_urls = String::new();
    let mut ice_turn_username = String::new();
    let mut ice_turn_password = String::new();

    // Try to read secrets.toml
    if let Ok(content) = fs::read_to_string("secrets.toml") {
        if let Ok(value) = content.parse::<toml::Value>() {
            // RustDesk section
            if let Some(rustdesk) = value.get("rustdesk") {
                if let Some(v) = rustdesk.get("public_server").and_then(|v| v.as_str()) {
                    rustdesk_public_server = v.to_string();
                }
                if let Some(v) = rustdesk.get("public_key").and_then(|v| v.as_str()) {
                    rustdesk_public_key = v.to_string();
                }
                if let Some(v) = rustdesk.get("relay_key").and_then(|v| v.as_str()) {
                    rustdesk_relay_key = v.to_string();
                }
            }

            // ICE section (for WebRTC)
            if let Some(ice) = value.get("ice") {
                if let Some(v) = ice.get("stun_server").and_then(|v| v.as_str()) {
                    ice_stun_server = v.to_string();
                }
                if let Some(v) = ice.get("turn_urls").and_then(|v| v.as_str()) {
                    ice_turn_urls = v.to_string();
                }
                if let Some(v) = ice.get("turn_username").and_then(|v| v.as_str()) {
                    ice_turn_username = v.to_string();
                }
                if let Some(v) = ice.get("turn_password").and_then(|v| v.as_str()) {
                    ice_turn_password = v.to_string();
                }
            }
        } else {
            println!("cargo:warning=Failed to parse secrets.toml");
        }
    } else {
        println!("cargo:warning=secrets.toml not found, using empty defaults");
    }

    // Generate the secrets module
    let code = format!(
        r#"// Auto-generated secrets module
// DO NOT EDIT - This file is generated by build.rs from secrets.toml

/// RustDesk public server configuration
pub mod rustdesk {{
    /// Public RustDesk ID server address (used when user leaves field empty)
    pub const PUBLIC_SERVER: &str = "{}";

    /// Public key for the RustDesk server (for client connection)
    pub const PUBLIC_KEY: &str = "{}";

    /// Relay server authentication key (licence_key for relay server)
    pub const RELAY_KEY: &str = "{}";

    /// Check if public server is configured
    pub const fn has_public_server() -> bool {{
        !PUBLIC_SERVER.is_empty()
    }}
}}

/// ICE server configuration (for WebRTC NAT traversal)
pub mod ice {{
    /// Public STUN server URL
    pub const STUN_SERVER: &str = "{}";

    /// Public TURN server URLs (comma-separated)
    pub const TURN_URLS: &str = "{}";

    /// TURN authentication username
    pub const TURN_USERNAME: &str = "{}";

    /// TURN authentication password
    pub const TURN_PASSWORD: &str = "{}";

    /// Check if public ICE servers are configured
    pub const fn is_configured() -> bool {{
        !STUN_SERVER.is_empty() || !TURN_URLS.is_empty()
    }}

    /// Check if TURN servers are configured (requires credentials)
    pub const fn has_turn() -> bool {{
        !TURN_URLS.is_empty() && !TURN_USERNAME.is_empty() && !TURN_PASSWORD.is_empty()
    }}
}}
"#,
        escape_string(&rustdesk_public_server),
        escape_string(&rustdesk_public_key),
        escape_string(&rustdesk_relay_key),
        escape_string(&ice_stun_server),
        escape_string(&ice_turn_urls),
        escape_string(&ice_turn_username),
        escape_string(&ice_turn_password),
    );

    fs::write(&dest_path, code).expect("Failed to write secrets_generated.rs");
}

/// Escape special characters in a string for use in Rust string literals
fn escape_string(s: &str) -> String {
    s.replace('\\', "\\\\").replace('"', "\\\"")
}

/// Convert days since Unix epoch to year-month-day
fn days_to_ymd(days: i64) -> (i32, u32, u32) {
    // Algorithm from http://howardhinnant.github.io/date_algorithms.html
    let z = days + 719468;
    let era = if z >= 0 { z } else { z - 146096 } / 146097;
    let doe = (z - era * 146097) as u32;
    let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146096) / 365;
    let y = yoe as i64 + era * 400;
    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
    let mp = (5 * doy + 2) / 153;
    let d = doy - (153 * mp + 2) / 5 + 1;
    let m = if mp < 10 { mp + 3 } else { mp - 9 };
    let year = if m <= 2 { y + 1 } else { y };
    (year as i32, m, d)
}