markbase/markbase-core/src/main.rs

use axum::{extract::Request, response::IntoResponse, routing::any, Extension, Router};
use clap::{Parser, Subcommand};
use std::path::Path;

#[derive(Parser)]
#[command(name = "markbase", about = "Momentry Display Engine")]
struct Cli {
    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
enum Commands {
    /// Start display server
    Display {
        #[arg(short, long, default_value = "11438")]
        port: u16,
        /// Optional initial markdown file
        file: Option<String>,
    },
    /// Render markdown to HTML (stdout)
    Render {
        file: String,
        #[arg(short, long)]
        output: Option<String>,
    },
    /// Configuration management
    Config {
        #[command(subcommand)]
        action: ConfigCommands,
    },
    /// Scan and import files from directory
    Scan {
        /// User ID
        #[arg(short, long)]
        user: String,
        /// Directory to scan
        #[arg(short, long)]
        dir: String,
        /// Batch size for database insertion
        #[arg(short, long, default_value = "100")]
        batch: usize,
        /// Skip SHA256 hash calculation (faster import)
        #[arg(short, long, default_value = "true")]
        skip_hash: bool,
        /// Number of threads for hash calculation (if skip_hash=false)
        #[arg(short, long, default_value = "4")]
        threads: usize,
    },
    /// Compute SHA256 hashes for imported files
    Hash {
        /// User ID
        #[arg(short, long)]
        user: String,
        /// Number of threads for parallel hash calculation
        #[arg(short, long, default_value = "4")]
        threads: usize,
    },
    /// Start WebDAV server
    WebDAV {
        #[command(subcommand)]
        action: WebDAVCommands,
    },
    /// Manage iSCSI target (gotgt)
    Iscsi {
        #[command(subcommand)]
        action: IscsiCommands,
    },
    /// Start SFTP server
    Sftp {
        /// Port to listen on
        #[arg(short, long, default_value = "2023")]
        port: u16,
        /// User ID for database
        #[arg(short, long)]
        user: String,
    },
    /// Test bcrypt password hash
    BcryptTest {
        /// Password to hash
        #[arg(short, long, default_value = "demo123")]
        password: String,
        /// Hash to verify (optional)
        #[arg(long)]
        verify_hash: Option<String>,
    },
    /// Start SSH server (hand-written implementation)
    SshServer {
        /// Port to listen on (default 2024)
        #[arg(short, long, default_value = "2024")]
        port: u16,
    },
}

#[derive(Subcommand)]
pub enum IscsiCommands {
    /// Start iSCSI target daemon
    Start {
        #[arg(short, long)]
        user: String,
        #[arg(short, long, default_value = "3260")]
        port: u16,
        #[arg(short, long, default_value = "5GB")]
        lun_size: String,
        #[arg(short, long)]
        force: bool,
        /// Block device path (e.g., /dev/disk5). Overrides file-backed LUN.
        #[arg(long)]
        device: Option<String>,
    },
    /// Stop iSCSI target daemon
    Stop,
    /// Show iSCSI target status
    Status,
}

#[derive(Subcommand)]
enum WebDAVCommands {
    /// Start WebDAV server for user
    Start {
        /// Port to listen on
        #[arg(short, long, default_value = "8002")]
        port: u16,
        /// User ID for database
        #[arg(short, long)]
        user: String,
    },
}

#[derive(Subcommand)]
enum ConfigCommands {
    /// Initialize default configuration file
    Init {
        #[arg(short, long)]
        force: bool,
    },
    /// Show current configuration
    Show {
        #[arg(short, long)]
        section: Option<String>,
    },
    /// Edit configuration
    Edit {
        #[arg(short, long)]
        key: String,
        #[arg(short, long)]
        value: String,
    },
    /// Validate configuration
    Validate,
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let cli = Cli::parse();

    match cli.command {
        Commands::Display { port, file } => {
            markbase_core::server::run(port, file).await?;
        }
        Commands::Render { file, output } => {
            let md = std::fs::read_to_string(&file)?;
            let html = markbase_core::render::md_to_html(&md);
            if let Some(path) = &output {
                std::fs::write(path, html)?;
            } else {
                println!("{html}");
            }
        }
        Commands::Config { action } => {
            handle_config_command(action)?;
        }
        Commands::Scan {
            user,
            dir,
            batch,
            skip_hash,
            threads,
        } => {
            use markbase_core::scan::ScanOptions;
            let options = ScanOptions { skip_hash, threads };
            markbase_core::scan::scan_directory(&user, &dir, batch, options)?;
        }
        Commands::Hash { user, threads } => {
            markbase_core::scan::compute_hashes(&user, threads)?;
        }
        Commands::WebDAV { action } => {
            handle_webdav_command(action).await?;
        }
        Commands::Iscsi { action } => {
            handle_iscsi_command(action)?;
        }
Commands::Sftp { port, user } => {
            println!("SFTP server command is currently disabled (old implementation)");
            println!("Use 'ssh-server' command for the new SSH+SFTP implementation");
            // handle_sftp_command(port, user).await?;
        }
        Commands::BcryptTest { password, verify_hash } => {
            handle_bcrypt_test(password, verify_hash)?;
        }
Commands::SshServer { port } => {
            println!("=== MarkBase SSH Server (Hand-written Implementation) ===");
            println!("Port: {}", port);  // port已经是u16，不是Option<u16>
            println!("Implementation: SSH-2.0-MarkBaseSSH_1.0");
            println!("Features: SSH + SFTP + SCP + rsync");
            println!("Security: ⭐⭐⭐⭐⭐ (RustCrypto authoritative libraries)");
            println!();
            
            markbase_core::ssh_server::server::run_ssh_server(Some(port))?;
        }
        Commands::BcryptTest {
            password,
            verify_hash,
        } => {
            handle_bcrypt_test(password, verify_hash)?;
        }
    }
    Ok(())
}

fn handle_config_command(action: ConfigCommands) -> anyhow::Result<()> {
    match action {
        ConfigCommands::Init { force } => {
            let config_path = Path::new("config/markbase.toml");

            if config_path.exists() && !force {
                println!("Configuration file already exists at config/markbase.toml");
                println!("Use --force to overwrite");
                return Ok(());
            }

            let config = markbase_core::config::MarkBaseConfig::default_config();
            config.save(config_path)?;

            println!("✓ Configuration file created: config/markbase.toml");
            println!("Default values:");
            println!("  Server port: {}", config.server.port);
            println!("  PostgreSQL host: {}", config.postgresql.host);
            println!("  Test users: {}", config.test.users.join(", "));
        }
        ConfigCommands::Show { section } => {
            let config_path = Path::new("config/markbase.toml");

            if !config_path.exists() {
                println!("Configuration file not found. Run 'markbase config init' first.");
                return Ok(());
            }

            let config = markbase_core::config::MarkBaseConfig::load(config_path)?;

            if let Some(s) = section {
                show_section(&config, &s);
            } else {
                println!("{}", toml::to_string_pretty(&config)?);
            }
        }
        ConfigCommands::Edit { key, value } => {
            let config_path = Path::new("config/markbase.toml");

            if !config_path.exists() {
                println!("Configuration file not found. Run 'markbase config init' first.");
                return Ok(());
            }

            let mut config = markbase_core::config::MarkBaseConfig::load(config_path)?;

            match config.get(&key) {
                Some(old_value) => {
                    config.set(&key, &value)?;
                    config.validate()?;
                    config.save(config_path)?;
                    println!("✓ Updated {}: {} → {}", key, old_value, value);
                }
                None => {
                    println!("Invalid config key: {}", key);
                    println!(
                        "Valid keys: server.*, postgresql.*, authentication.*, test.*, logging.*"
                    );
                }
            }
        }
        ConfigCommands::Validate => {
            let config_path = Path::new("config/markbase.toml");

            if !config_path.exists() {
                println!("Configuration file not found. Run 'markbase config init' first.");
                return Ok(());
            }

            let config = markbase_core::config::MarkBaseConfig::load(config_path)?;

            match config.validate() {
                Ok(_) => {
                    println!("✓ Configuration is valid");
                }
                Err(e) => {
                    println!("✗ Configuration validation failed: {}", e);
                }
            }
        }
    }

    Ok(())
}

fn show_section(config: &markbase_core::config::MarkBaseConfig, section: &str) {
    match section {
        "server" => println!("{}", toml::to_string_pretty(&config.server).unwrap()),
        "postgresql" => println!("{}", toml::to_string_pretty(&config.postgresql).unwrap()),
        "authentication" => println!("{}", toml::to_string_pretty(&config.authentication).unwrap()),
        "test" => println!("{}", toml::to_string_pretty(&config.test).unwrap()),
        "logging" => println!("{}", toml::to_string_pretty(&config.logging).unwrap()),
        _ => println!("Invalid section: {}. Valid sections: server, postgresql, authentication, test, logging", section),
    }
}

fn handle_iscsi_command(action: IscsiCommands) -> anyhow::Result<()> {
    let binary = find_binary("markbase-iscsi");
    let mut cmd = std::process::Command::new(&binary);
    cmd.arg("iscsi");
    match action {
        IscsiCommands::Start {
            user,
            port,
            lun_size,
            force,
            device,
        } => {
            cmd.arg("start")
                .args(["--user", &user])
                .args(["--port", &port.to_string()])
                .args(["--lun-size", &lun_size]);
            if force {
                cmd.arg("--force");
            }
            if let Some(d) = device {
                cmd.args(["--device", &d]);
            }
        }
        IscsiCommands::Stop => {
            cmd.arg("stop");
        }
        IscsiCommands::Status => {
            cmd.arg("status");
        }
    }
    let status = cmd.status()?;
    std::process::exit(status.code().unwrap_or(1));
}

async fn handle_webdav_command(action: WebDAVCommands) -> anyhow::Result<()> {
    match action {
        WebDAVCommands::Start { port, user } => {
            let db_path = std::path::PathBuf::from(filetree::FileTree::user_db_path(&user));

            if !db_path.exists() {
                return Err(anyhow::anyhow!(
                    "User database not found: {}",
                    db_path.display()
                ));
            }

            println!("=== MarkBase WebDAV Server ===");
            println!("User: {}", user);
            println!("Port: {}", port);
            println!("Database: {}", db_path.display());
            println!("");

            run_webdav_server(port, user, db_path).await?;
        }
    }
    Ok(())
}

async fn run_webdav_server(
    port: u16,
    user: String,
    db_path: std::path::PathBuf,
) -> anyhow::Result<()> {
    use tokio::net::TcpListener;

    let webdav = markbase_webdav::webdav::MarkBaseWebDAV::new(user, db_path);
    let dav_handler = webdav.create_handler();

    let app = Router::new()
        .route("/webdav", any(handle_dav))
        .route("/webdav/", any(handle_dav))
        .route("/webdav/*path", any(handle_dav))
        .layer(Extension(dav_handler));

    let addr = format!("127.0.0.1:{}", port);
    let listener = TcpListener::bind(&addr).await?;

    println!("WebDAV server listening on http://{}", addr);
    println!("Mount point: /webdav");
    println!("");
    println!("Press Ctrl+C to stop");

    axum::serve(listener, app).await?;

    Ok(())
}

async fn handle_dav(
    Extension(dav): Extension<dav_server::DavHandler>,
    req: Request,
) -> impl IntoResponse {
    dav.handle(req).await
}

fn find_binary(name: &str) -> std::path::PathBuf {
    let exe = std::env::current_exe().unwrap();
    let dir = exe.parent().unwrap();
    dir.join(name)
}

// async fn handle_sftp_command(port: u16, user: String) -> anyhow::Result<()> {
//     println!("=== MarkBase SFTP Server ===");
//     println!("User: {}", user);
//     println!("Port: {}", port);
//     println!("Auth DB: data/auth.sqlite");
//     println!("FileTree DB: data/users/{}.sqlite", user);
//     println!("");
//
//     let config = markbase_core::sftp::SftpConfig::load_default()?;
//
//     if port != config.sftp.port {
//         println!(
//             "Warning: CLI port {} overrides config port {}",
//             port, config.sftp.port
//         );
//     }
//
//     markbase_core::sftp::run_server(config, &user).await?;
//
//     Ok(())
// }

fn handle_bcrypt_test(password: String, verify_hash: Option<String>) -> anyhow::Result<()> {
    use bcrypt::{hash, verify, DEFAULT_COST};

    println!("=== bcrypt Hash Test ===");
    println!("Password: {}", password);
    println!("");

    // Generate new hash
    let new_hash = hash(&password, DEFAULT_COST)?;
    println!("Generated hash:");
    println!("{}", new_hash);
    println!("");

    // Verify current hash if provided
    if let Some(hash_to_verify) = verify_hash {
        println!("Verifying hash: {}", hash_to_verify);
        let valid = verify(&password, &hash_to_verify)?;
        println!("Valid: {}", valid);
        println!("");
    }

    // Verify database hash
    let db_hash = "$2b$10$ha5wU.mOi8fHLJCfun860u2cfVopa04jwe/q82IKOwqp5uG70qsH6";
    println!("Database hash: {}", db_hash);
    let valid = verify(&password, db_hash)?;
    println!("Database hash valid for '{}': {}", password, valid);
    println!("");

    if !valid {
        println!("❌ Database hash is incorrect!");
        println!("Update SQL:");
        println!("UPDATE sftpgo_users SET password_hash = '{}' WHERE username IN ('testuser', 'demo', 'warren', 'momentry');", new_hash);
        println!("");
        println!("Execute:");
        println!("sqlite3 data/auth.sqlite \"UPDATE sftpgo_users SET password_hash = '{}' WHERE username IN ('testuser', 'demo', 'warren', 'momentry');\"", new_hash);
    } else {
        println!("✅ Database hash is correct!");
    }

    Ok(())
}