1→//! GuruConnect Server - WebSocket Relay Server
     2→//!
     3→//! Handles connections from both agents and dashboard viewers,
     4→//! relaying video frames and input events between them.
     5→
     6→mod config;
     7→mod relay;
     8→mod session;
     9→mod auth;
    10→mod api;
    11→mod db;
    12→mod support_codes;
    13→
    14→pub mod proto {
    15→    include!(concat!(env!("OUT_DIR"), "/guruconnect.rs"));
    16→}
    17→
    18→use anyhow::Result;
    19→use axum::{
    20→    Router,
    21→<<<<<<< HEAD
    22→    routing::{get, post, delete},
    23→    extract::{Path, State, Json, Query},
    24→=======
    25→    routing::{get, post, put, delete},
    26→    extract::{Path, State, Json, Request},
    27→>>>>>>> b861cb1 (Add user management system with JWT authentication)
    28→    response::{Html, IntoResponse},
    29→    http::StatusCode,
    30→    middleware::{self, Next},
    31→};
    32→use std::net::SocketAddr;
    33→use std::sync::Arc;
    34→use tower_http::cors::{Any, CorsLayer};
    35→use tower_http::trace::TraceLayer;
    36→use tower_http::services::ServeDir;
    37→use tracing::{info, Level};
    38→use tracing_subscriber::FmtSubscriber;
    39→use serde::Deserialize;
    40→
    41→use support_codes::{SupportCodeManager, CreateCodeRequest, SupportCode, CodeValidation};
    42→use auth::{JwtConfig, hash_password, generate_random_password};
    43→
    44→/// Application state
    45→#[derive(Clone)]
    46→pub struct AppState {
    47→    sessions: session::SessionManager,
    48→    support_codes: SupportCodeManager,
    49→    db: Option<db::Database>,
    50→    pub jwt_config: Arc<JwtConfig>,
    51→}
    52→
    53→/// Middleware to inject JWT config into request extensions
    54→async fn auth_layer(
    55→    State(state): State<AppState>,
    56→    mut request: Request,
    57→    next: Next,
    58→) -> impl IntoResponse {
    59→    request.extensions_mut().insert(state.jwt_config.clone());
    60→    next.run(request).await
    61→}
    62→
    63→#[tokio::main]
    64→async fn main() -> Result<()> {
    65→    // Initialize logging
    66→    let _subscriber = FmtSubscriber::builder()
    67→        .with_max_level(Level::INFO)
    68→        .with_target(true)
    69→        .init();
    70→
    71→    info!("GuruConnect Server v{}", env!("CARGO_PKG_VERSION"));
    72→
    73→    // Load configuration
    74→    let config = config::Config::load()?;
    75→
    76→    // Use port 3002 for GuruConnect
    77→    let listen_addr = std::env::var("LISTEN_ADDR").unwrap_or_else(|_| "0.0.0.0:3002".to_string());
    78→    info!("Loaded configuration, listening on {}", listen_addr);
    79→
    80→    // JWT configuration
    81→    let jwt_secret = std::env::var("JWT_SECRET").unwrap_or_else(|_| {
    82→        tracing::warn!("JWT_SECRET not set, using default (INSECURE for production!)");
    83→        "guruconnect-dev-secret-change-me-in-production".to_string()
    84→    });
    85→    let jwt_expiry_hours = std::env::var("JWT_EXPIRY_HOURS")
    86→        .ok()
    87→        .and_then(|s| s.parse().ok())
    88→        .unwrap_or(24i64);
    89→    let jwt_config = Arc::new(JwtConfig::new(jwt_secret, jwt_expiry_hours));
    90→
    91→    // Initialize database if configured
    92→    let database = if let Some(ref db_url) = config.database_url {
    93→        match db::Database::connect(db_url, config.database_max_connections).await {
    94→            Ok(db) => {
    95→                // Run migrations
    96→                if let Err(e) = db.migrate().await {
    97→                    tracing::error!("Failed to run migrations: {}", e);
    98→                    return Err(e);
    99→                }
   100→                Some(db)
   101→            }
   102→            Err(e) => {
   103→                tracing::warn!("Failed to connect to database: {}. Running without persistence.", e);
   104→                None
   105→            }
   106→        }
   107→    } else {
   108→        info!("No DATABASE_URL set, running without persistence");
   109→        None
   110→    };
   111→
   112→    // Create initial admin user if no users exist
   113→    if let Some(ref db) = database {
   114→        match db::count_users(db.pool()).await {
   115→            Ok(0) => {
   116→                info!("No users found, creating initial admin user...");
   117→                let password = generate_random_password(16);
   118→                let password_hash = hash_password(&password)?;
   119→
   120→                match db::create_user(db.pool(), "admin", &password_hash, None, "admin").await {
   121→                    Ok(user) => {
   122→                        // Set admin permissions
   123→                        let perms = vec![
   124→                            "view".to_string(),
   125→                            "control".to_string(),
   126→                            "transfer".to_string(),
   127→                            "manage_users".to_string(),
   128→                            "manage_clients".to_string(),
   129→                        ];
   130→                        let _ = db::set_user_permissions(db.pool(), user.id, &perms).await;
   131→
   132→                        info!("========================================");
   133→                        info!("  INITIAL ADMIN USER CREATED");
   134→                        info!("  Username: admin");
   135→                        info!("  Password: {}", password);
   136→                        info!("  (Change this password after first login!)");
   137→                        info!("========================================");
   138→                    }
   139→                    Err(e) => {
   140→                        tracing::error!("Failed to create initial admin user: {}", e);
   141→                    }
   142→                }
   143→            }
   144→            Ok(count) => {
   145→                info!("{} user(s) in database", count);
   146→            }
   147→            Err(e) => {
   148→                tracing::warn!("Could not check user count: {}", e);
   149→            }
   150→        }
   151→    }
   152→
   153→    // Create session manager
   154→    let sessions = session::SessionManager::new();
   155→
   156→    // Restore persistent machines from database
   157→    if let Some(ref db) = database {
   158→        match db::machines::get_all_machines(db.pool()).await {
   159→            Ok(machines) => {
   160→                info!("Restoring {} persistent machines from database", machines.len());
   161→                for machine in machines {
   162→                    sessions.restore_offline_machine(&machine.agent_id, &machine.hostname).await;
   163→                }
   164→            }
   165→            Err(e) => {
   166→                tracing::warn!("Failed to restore machines: {}", e);
   167→            }
   168→        }
   169→    }
   170→
   171→    // Create application state
   172→    let state = AppState {
   173→        sessions,
   174→        support_codes: SupportCodeManager::new(),
   175→        db: database,
   176→        jwt_config,
   177→    };
   178→
   179→    // Build router
   180→    let app = Router::new()
   181→        // Health check (no auth required)
   182→        .route("/health", get(health))
   183→
   184→        // Auth endpoints (no auth required for login)
   185→        .route("/api/auth/login", post(api::auth::login))
   186→
   187→        // Auth endpoints (auth required)
   188→        .route("/api/auth/me", get(api::auth::get_me))
   189→        .route("/api/auth/change-password", post(api::auth::change_password))
   190→
   191→        // User management (admin only)
   192→        .route("/api/users", get(api::users::list_users))
   193→        .route("/api/users", post(api::users::create_user))
   194→        .route("/api/users/:id", get(api::users::get_user))
   195→        .route("/api/users/:id", put(api::users::update_user))
   196→        .route("/api/users/:id", delete(api::users::delete_user))
   197→        .route("/api/users/:id/permissions", put(api::users::set_permissions))
   198→        .route("/api/users/:id/clients", put(api::users::set_client_access))
   199→
   200→        // Portal API - Support codes
   201→        .route("/api/codes", post(create_code))
   202→        .route("/api/codes", get(list_codes))
   203→        .route("/api/codes/:code/validate", get(validate_code))
   204→        .route("/api/codes/:code/cancel", post(cancel_code))
   205→
   206→        // WebSocket endpoints
   207→        .route("/ws/agent", get(relay::agent_ws_handler))
   208→        .route("/ws/viewer", get(relay::viewer_ws_handler))
   209→
   210→        // REST API - Sessions
   211→        .route("/api/sessions", get(list_sessions))
   212→        .route("/api/sessions/:id", get(get_session))
   213→        .route("/api/sessions/:id", delete(disconnect_session))
   214→
   215→<<<<<<< HEAD
   216→        // REST API - Machines
   217→        .route("/api/machines", get(list_machines))
   218→        .route("/api/machines/:agent_id", get(get_machine))
   219→        .route("/api/machines/:agent_id", delete(delete_machine))
   220→        .route("/api/machines/:agent_id/history", get(get_machine_history))
   221→        
   222→=======
   223→>>>>>>> b861cb1 (Add user management system with JWT authentication)
   224→        // HTML page routes (clean URLs)
   225→        .route("/login", get(serve_login))
   226→        .route("/dashboard", get(serve_dashboard))
   227→        .route("/users", get(serve_users))
   228→
   229→        // State and middleware
   230→        .with_state(state.clone())
   231→        .layer(middleware::from_fn_with_state(state, auth_layer))
   232→
   233→        // Serve static files for portal (fallback)
   234→        .fallback_service(ServeDir::new("static").append_index_html_on_directories(true))
   235→
   236→        // Middleware
   237→        .layer(TraceLayer::new_for_http())
   238→        .layer(
   239→            CorsLayer::new()
   240→                .allow_origin(Any)
   241→                .allow_methods(Any)
   242→                .allow_headers(Any),
   243→        );
   244→
   245→    // Start server
   246→    let addr: SocketAddr = listen_addr.parse()?;
   247→    let listener = tokio::net::TcpListener::bind(addr).await?;
   248→
   249→    info!("Server listening on {}", addr);
   250→
   251→    axum::serve(listener, app).await?;
   252→
   253→    Ok(())
   254→}
   255→
   256→async fn health() -> &'static str {
   257→    "OK"
   258→}
   259→
   260→// Support code API handlers
   261→
   262→async fn create_code(
   263→    State(state): State<AppState>,
   264→    Json(request): Json<CreateCodeRequest>,
   265→) -> Json<SupportCode> {
   266→    let code = state.support_codes.create_code(request).await;
   267→    info!("Created support code: {}", code.code);
   268→    Json(code)
   269→}
   270→
   271→async fn list_codes(
   272→    State(state): State<AppState>,
   273→) -> Json<Vec<SupportCode>> {
   274→    Json(state.support_codes.list_active_codes().await)
   275→}
   276→
   277→#[derive(Deserialize)]
   278→struct ValidateParams {
   279→    code: String,
   280→}
   281→
   282→async fn validate_code(
   283→    State(state): State<AppState>,
   284→    Path(code): Path<String>,
   285→) -> Json<CodeValidation> {
   286→    Json(state.support_codes.validate_code(&code).await)
   287→}
   288→
   289→async fn cancel_code(
   290→    State(state): State<AppState>,
   291→    Path(code): Path<String>,
   292→) -> impl IntoResponse {
   293→    if state.support_codes.cancel_code(&code).await {
   294→        (StatusCode::OK, "Code cancelled")
   295→    } else {
   296→        (StatusCode::BAD_REQUEST, "Cannot cancel code")
   297→    }
   298→}
   299→
   300→// Session API handlers (updated to use AppState)
   301→
   302→async fn list_sessions(
   303→    State(state): State<AppState>,
   304→) -> Json<Vec<api::SessionInfo>> {
   305→    let sessions = state.sessions.list_sessions().await;
   306→    Json(sessions.into_iter().map(api::SessionInfo::from).collect())
   307→}
   308→
   309→async fn get_session(
   310→    State(state): State<AppState>,
   311→    Path(id): Path<String>,
   312→) -> Result<Json<api::SessionInfo>, (StatusCode, &'static str)> {
   313→    let session_id = uuid::Uuid::parse_str(&id)
   314→        .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid session ID"))?;
   315→
   316→    let session = state.sessions.get_session(session_id).await
   317→        .ok_or((StatusCode::NOT_FOUND, "Session not found"))?;
   318→
   319→    Ok(Json(api::SessionInfo::from(session)))
   320→}
   321→
   322→async fn disconnect_session(
   323→    State(state): State<AppState>,
   324→    Path(id): Path<String>,
   325→) -> impl IntoResponse {
   326→    let session_id = match uuid::Uuid::parse_str(&id) {
   327→        Ok(id) => id,
   328→        Err(_) => return (StatusCode::BAD_REQUEST, "Invalid session ID"),
   329→    };
   330→
   331→    if state.sessions.disconnect_session(session_id, "Disconnected by administrator").await {
   332→        info!("Session {} disconnected by admin", session_id);
   333→        (StatusCode::OK, "Session disconnected")
   334→    } else {
   335→        (StatusCode::NOT_FOUND, "Session not found")
   336→    }
   337→}
   338→
   339→// Machine API handlers
   340→
   341→async fn list_machines(
   342→    State(state): State<AppState>,
   343→) -> Result<Json<Vec<api::MachineInfo>>, (StatusCode, &'static str)> {
   344→    let db = state.db.as_ref()
   345→        .ok_or((StatusCode::SERVICE_UNAVAILABLE, "Database not available"))?;
   346→
   347→    let machines = db::machines::get_all_machines(db.pool()).await
   348→        .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "Database error"))?;
   349→
   350→    Ok(Json(machines.into_iter().map(api::MachineInfo::from).collect()))
   351→}
   352→
   353→async fn get_machine(
   354→    State(state): State<AppState>,
   355→    Path(agent_id): Path<String>,
   356→) -> Result<Json<api::MachineInfo>, (StatusCode, &'static str)> {
   357→    let db = state.db.as_ref()
   358→        .ok_or((StatusCode::SERVICE_UNAVAILABLE, "Database not available"))?;
   359→
   360→    let machine = db::machines::get_machine_by_agent_id(db.pool(), &agent_id).await
   361→        .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "Database error"))?
   362→        .ok_or((StatusCode::NOT_FOUND, "Machine not found"))?;
   363→
   364→    Ok(Json(api::MachineInfo::from(machine)))
   365→}
   366→
   367→async fn get_machine_history(
   368→    State(state): State<AppState>,
   369→    Path(agent_id): Path<String>,
   370→) -> Result<Json<api::MachineHistory>, (StatusCode, &'static str)> {
   371→    let db = state.db.as_ref()
   372→        .ok_or((StatusCode::SERVICE_UNAVAILABLE, "Database not available"))?;
   373→
   374→    // Get machine
   375→    let machine = db::machines::get_machine_by_agent_id(db.pool(), &agent_id).await
   376→        .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "Database error"))?
   377→        .ok_or((StatusCode::NOT_FOUND, "Machine not found"))?;
   378→
   379→    // Get sessions for this machine
   380→    let sessions = db::sessions::get_sessions_for_machine(db.pool(), machine.id).await
   381→        .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "Database error"))?;
   382→
   383→    // Get events for this machine
   384→    let events = db::events::get_events_for_machine(db.pool(), machine.id).await
   385→        .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "Database error"))?;
   386→
   387→    let history = api::MachineHistory {
   388→        machine: api::MachineInfo::from(machine),
   389→        sessions: sessions.into_iter().map(api::SessionRecord::from).collect(),
   390→        events: events.into_iter().map(api::EventRecord::from).collect(),
   391→        exported_at: chrono::Utc::now().to_rfc3339(),
   392→    };
   393→
   394→    Ok(Json(history))
   395→}
   396→
   397→async fn delete_machine(
   398→    State(state): State<AppState>,
   399→    Path(agent_id): Path<String>,
   400→    Query(params): Query<api::DeleteMachineParams>,
   401→) -> Result<Json<api::DeleteMachineResponse>, (StatusCode, &'static str)> {
   402→    let db = state.db.as_ref()
   403→        .ok_or((StatusCode::SERVICE_UNAVAILABLE, "Database not available"))?;
   404→
   405→    // Get machine first
   406→    let machine = db::machines::get_machine_by_agent_id(db.pool(), &agent_id).await
   407→        .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "Database error"))?
   408→        .ok_or((StatusCode::NOT_FOUND, "Machine not found"))?;
   409→
   410→    // Export history if requested
   411→    let history = if params.export {
   412→        let sessions = db::sessions::get_sessions_for_machine(db.pool(), machine.id).await
   413→            .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "Database error"))?;
   414→        let events = db::events::get_events_for_machine(db.pool(), machine.id).await
   415→            .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "Database error"))?;
   416→
   417→        Some(api::MachineHistory {
   418→            machine: api::MachineInfo::from(machine.clone()),
   419→            sessions: sessions.into_iter().map(api::SessionRecord::from).collect(),
   420→            events: events.into_iter().map(api::EventRecord::from).collect(),
   421→            exported_at: chrono::Utc::now().to_rfc3339(),
   422→        })
   423→    } else {
   424→        None
   425→    };
   426→
   427→    // Send uninstall command if requested and agent is online
   428→    let mut uninstall_sent = false;
   429→    if params.uninstall {
   430→        // Find session for this agent
   431→        if let Some(session) = state.sessions.get_session_by_agent(&agent_id).await {
   432→            if session.is_online {
   433→                uninstall_sent = state.sessions.send_admin_command(
   434→                    session.id,
   435→                    proto::AdminCommandType::AdminUninstall,
   436→                    "Deleted by administrator",
   437→                ).await;
   438→                if uninstall_sent {
   439→                    info!("Sent uninstall command to agent {}", agent_id);
   440→                }
   441→            }
   442→        }
   443→    }
   444→
   445→    // Remove from session manager
   446→    state.sessions.remove_agent(&agent_id).await;
   447→
   448→    // Delete from database (cascades to sessions and events)
   449→    db::machines::delete_machine(db.pool(), &agent_id).await
   450→        .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "Failed to delete machine"))?;
   451→
   452→    info!("Deleted machine {} (uninstall_sent: {})", agent_id, uninstall_sent);
   453→
   454→    Ok(Json(api::DeleteMachineResponse {
   455→        success: true,
   456→        message: format!("Machine {} deleted", machine.hostname),
   457→        uninstall_sent,
   458→        history,
   459→    }))
   460→}
   461→
   462→// Static page handlers
   463→async fn serve_login() -> impl IntoResponse {
   464→    match tokio::fs::read_to_string("static/login.html").await {
   465→        Ok(content) => Html(content).into_response(),
   466→        Err(_) => (StatusCode::NOT_FOUND, "Page not found").into_response(),
   467→    }
   468→}
   469→
   470→async fn serve_dashboard() -> impl IntoResponse {
   471→    match tokio::fs::read_to_string("static/dashboard.html").await {
   472→        Ok(content) => Html(content).into_response(),
   473→        Err(_) => (StatusCode::NOT_FOUND, "Page not found").into_response(),
   474→    }
   475→}
   476→
   477→async fn serve_users() -> impl IntoResponse {
   478→    match tokio::fs::read_to_string("static/users.html").await {
   479→        Ok(content) => Html(content).into_response(),
   480→        Err(_) => (StatusCode::NOT_FOUND, "Page not found").into_response(),
   481→    }
   482→}
   483→

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>