wraith/src-tauri/src/lib.rs

pub mod db;
pub mod vault;
pub mod settings;
pub mod connections;
pub mod credentials;
pub mod ssh;
pub mod sftp;
pub mod rdp;
pub mod theme;
pub mod workspace;
pub mod pty;
pub mod mcp;
pub mod scanner;
pub mod commands;

use std::path::PathBuf;
use std::sync::Mutex;

use db::Database;
use vault::VaultService;
use credentials::CredentialService;
use settings::SettingsService;
use connections::ConnectionService;
use sftp::SftpService;
use ssh::session::SshService;
use rdp::RdpService;
use theme::ThemeService;
use workspace::WorkspaceService;
use pty::PtyService;
use mcp::ScrollbackRegistry;
use mcp::error_watcher::ErrorWatcher;

pub struct AppState {
    pub db: Database,
    pub vault: Mutex<Option<VaultService>>,
    pub settings: SettingsService,
    pub connections: ConnectionService,
    pub credentials: Mutex<Option<CredentialService>>,
    pub ssh: SshService,
    pub sftp: SftpService,
    pub rdp: RdpService,
    pub theme: ThemeService,
    pub workspace: WorkspaceService,
    pub pty: PtyService,
    pub scrollback: ScrollbackRegistry,
    pub error_watcher: std::sync::Arc<ErrorWatcher>,
}

impl AppState {
    pub fn new(data_dir: PathBuf) -> Result<Self, Box<dyn std::error::Error>> {
        std::fs::create_dir_all(&data_dir)?;
        let database = Database::open(&data_dir.join("wraith.db"))?;
        database.migrate()?;
        Ok(Self {
            db: database.clone(),
            vault: Mutex::new(None),
            settings: SettingsService::new(database.clone()),
            connections: ConnectionService::new(database.clone()),
            credentials: Mutex::new(None),
            ssh: SshService::new(database.clone()),
            sftp: SftpService::new(),
            rdp: RdpService::new(),
            theme: ThemeService::new(database.clone()),
            workspace: WorkspaceService::new(SettingsService::new(database.clone())),
            pty: PtyService::new(),
            scrollback: ScrollbackRegistry::new(),
            error_watcher: std::sync::Arc::new(ErrorWatcher::new()),
        })
    }

    pub fn is_first_run(&self) -> bool {
        self.settings.get("vault_salt").unwrap_or_default().is_empty()
    }

    pub fn is_unlocked(&self) -> bool {
        self.vault.lock().unwrap().is_some()
    }
}

pub fn data_directory() -> PathBuf {
    if let Ok(appdata) = std::env::var("APPDATA") { return PathBuf::from(appdata).join("Wraith"); }
    if let Ok(home) = std::env::var("HOME") {
        if cfg!(target_os = "macos") { return PathBuf::from(home).join("Library").join("Application Support").join("Wraith"); }
        if let Ok(xdg) = std::env::var("XDG_DATA_HOME") { return PathBuf::from(xdg).join("wraith"); }
        return PathBuf::from(home).join(".local").join("share").join("wraith");
    }
    PathBuf::from(".")
}

#[cfg_attr(mobile, tauri::mobile_entry_point)]
pub fn run() {
    let app_state = AppState::new(data_directory()).expect("Failed to init AppState");
    app_state.theme.seed_builtins();

    tauri::Builder::default()
        .plugin(tauri_plugin_shell::init())
        .manage(app_state)
        .setup(|app| {
            #[cfg(debug_assertions)]
            {
                use tauri::Manager;
                if let Some(window) = app.get_webview_window("main") {
                    window.open_devtools();
                }
            }

            // Start MCP HTTP server for bridge binary communication
            {
                use tauri::Manager;
                let state = app.state::<AppState>();
                // Start error watcher background task
                let watcher_clone = state.error_watcher.clone();
                let scrollback_for_watcher = state.scrollback.clone();
                let app_for_watcher = app.handle().clone();
                mcp::error_watcher::start_error_watcher(watcher_clone, scrollback_for_watcher, app_for_watcher);

                let ssh_clone = state.ssh.clone();
                let rdp_clone = state.rdp.clone();
                let sftp_clone = state.sftp.clone();
                let scrollback_clone = state.scrollback.clone();
                tauri::async_runtime::spawn(async move {
                    match mcp::server::start_mcp_server(ssh_clone, rdp_clone, sftp_clone, scrollback_clone).await {
                        Ok(port) => log::info!("MCP server started on localhost:{}", port),
                        Err(e) => log::error!("Failed to start MCP server: {}", e),
                    }
                });
            }

            Ok(())
        })
        .invoke_handler(tauri::generate_handler![
            commands::vault::is_first_run, commands::vault::create_vault, commands::vault::unlock, commands::vault::is_unlocked,
            commands::settings::get_setting, commands::settings::set_setting,
            commands::connections::list_connections, commands::connections::create_connection, commands::connections::get_connection, commands::connections::update_connection, commands::connections::delete_connection,
            commands::connections::list_groups, commands::connections::create_group, commands::connections::delete_group, commands::connections::rename_group, commands::connections::search_connections,
            commands::credentials::list_credentials, commands::credentials::create_password, commands::credentials::create_ssh_key, commands::credentials::delete_credential, commands::credentials::decrypt_password, commands::credentials::decrypt_ssh_key,
            commands::ssh_commands::connect_ssh, commands::ssh_commands::connect_ssh_with_key, commands::ssh_commands::ssh_write, commands::ssh_commands::ssh_resize, commands::ssh_commands::disconnect_ssh, commands::ssh_commands::disconnect_session, commands::ssh_commands::list_ssh_sessions,
            commands::sftp_commands::sftp_list, commands::sftp_commands::sftp_read_file, commands::sftp_commands::sftp_write_file, commands::sftp_commands::sftp_mkdir, commands::sftp_commands::sftp_delete, commands::sftp_commands::sftp_rename,
            commands::rdp_commands::connect_rdp, commands::rdp_commands::rdp_get_frame, commands::rdp_commands::rdp_send_mouse, commands::rdp_commands::rdp_send_key, commands::rdp_commands::rdp_send_clipboard, commands::rdp_commands::disconnect_rdp, commands::rdp_commands::list_rdp_sessions,
            commands::theme_commands::list_themes, commands::theme_commands::get_theme,
            commands::pty_commands::list_available_shells, commands::pty_commands::spawn_local_shell, commands::pty_commands::pty_write, commands::pty_commands::pty_resize, commands::pty_commands::disconnect_pty,
            commands::mcp_commands::mcp_list_sessions, commands::mcp_commands::mcp_terminal_read, commands::mcp_commands::mcp_terminal_execute, commands::mcp_commands::mcp_get_session_context,
            commands::scanner_commands::scan_network, commands::scanner_commands::scan_ports, commands::scanner_commands::quick_scan,
            commands::tools_commands::tool_ping, commands::tools_commands::tool_traceroute, commands::tools_commands::tool_wake_on_lan, commands::tools_commands::tool_generate_ssh_key, commands::tools_commands::tool_generate_password,
            commands::tools_commands_r2::tool_dns_lookup, commands::tools_commands_r2::tool_whois, commands::tools_commands_r2::tool_bandwidth_iperf, commands::tools_commands_r2::tool_bandwidth_speedtest, commands::tools_commands_r2::tool_subnet_calc,
        ])
        .run(tauri::generate_context!())
        .expect("error while running tauri application");
}