claudetools/projects/msp-tools/guru-rmm/server/src/api/agents.rs

//! Agent management API endpoints

use axum::{
    extract::{Path, State},
    http::StatusCode,
    Json,
};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

use crate::auth::AuthUser;
use crate::db::{self, AgentResponse, AgentStats};
use crate::ws::{generate_api_key, hash_api_key};
use crate::AppState;

/// Response for agent registration
#[derive(Debug, Serialize)]
pub struct RegisterAgentResponse {
    pub agent_id: Uuid,
    pub api_key: String,
    pub message: String,
}

/// Request to register a new agent
#[derive(Debug, Deserialize)]
pub struct RegisterAgentRequest {
    pub hostname: String,
    pub os_type: String,
    pub os_version: Option<String>,
}

/// Register a new agent (generates API key)
/// Requires authentication to prevent unauthorized agent registration.
pub async fn register_agent(
    State(state): State<AppState>,
    user: AuthUser,
    Json(req): Json<RegisterAgentRequest>,
) -> Result<Json<RegisterAgentResponse>, (StatusCode, String)> {
    // Log who is registering the agent
    tracing::info!(
        user_id = %user.user_id,
        hostname = %req.hostname,
        os_type = %req.os_type,
        "Agent registration initiated by user"
    );

    // Generate a new API key
    let api_key = generate_api_key(&state.config.auth.api_key_prefix);
    let api_key_hash = hash_api_key(&api_key);

    // Create the agent
    let create = db::CreateAgent {
        hostname: req.hostname,
        api_key_hash,
        os_type: req.os_type,
        os_version: req.os_version,
        agent_version: None,
    };

    let agent = db::create_agent(&state.db, create)
        .await
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

    tracing::info!(
        user_id = %user.user_id,
        agent_id = %agent.id,
        "Agent registered successfully"
    );

    Ok(Json(RegisterAgentResponse {
        agent_id: agent.id,
        api_key, // Return the plain API key (only shown once!)
        message: "Agent registered successfully. Save the API key - it will not be shown again."
            .to_string(),
    }))
}

/// List all agents
/// Requires authentication.
pub async fn list_agents(
    State(state): State<AppState>,
    _user: AuthUser,
) -> Result<Json<Vec<AgentResponse>>, (StatusCode, String)> {
    let agents = db::get_all_agents(&state.db)
        .await
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

    let responses: Vec<AgentResponse> = agents.into_iter().map(|a| a.into()).collect();
    Ok(Json(responses))
}

/// Get a specific agent
/// Requires authentication.
pub async fn get_agent(
    State(state): State<AppState>,
    _user: AuthUser,
    Path(id): Path<Uuid>,
) -> Result<Json<AgentResponse>, (StatusCode, String)> {
    let agent = db::get_agent_by_id(&state.db, id)
        .await
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
        .ok_or((StatusCode::NOT_FOUND, "Agent not found".to_string()))?;

    Ok(Json(agent.into()))
}

/// Delete an agent
/// Requires authentication.
pub async fn delete_agent(
    State(state): State<AppState>,
    _user: AuthUser,
    Path(id): Path<Uuid>,
) -> Result<StatusCode, (StatusCode, String)> {
    // Check if agent is connected and disconnect it
    if state.agents.read().await.is_connected(&id) {
        // In a real implementation, we'd send a disconnect message
        state.agents.write().await.remove(&id);
    }

    let deleted = db::delete_agent(&state.db, id)
        .await
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

    if deleted {
        Ok(StatusCode::NO_CONTENT)
    } else {
        Err((StatusCode::NOT_FOUND, "Agent not found".to_string()))
    }
}

/// Get agent statistics
/// Requires authentication.
pub async fn get_stats(
    State(state): State<AppState>,
    _user: AuthUser,
) -> Result<Json<AgentStats>, (StatusCode, String)> {
    let stats = db::get_agent_stats(&state.db)
        .await
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

    Ok(Json(stats))
}

/// Request to move an agent to a different site
#[derive(Debug, Deserialize)]
pub struct MoveAgentRequest {
    pub site_id: Option<Uuid>, // None to unassign from site
}

/// Move an agent to a different site
/// Requires authentication.
pub async fn move_agent(
    State(state): State<AppState>,
    _user: AuthUser,
    Path(id): Path<Uuid>,
    Json(req): Json<MoveAgentRequest>,
) -> Result<Json<AgentResponse>, (StatusCode, String)> {
    // Verify the site exists if provided
    if let Some(site_id) = req.site_id {
        let site = db::get_site_by_id(&state.db, site_id)
            .await
            .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

        if site.is_none() {
            return Err((StatusCode::NOT_FOUND, "Site not found".to_string()));
        }
    }

    // Move the agent
    let agent = db::move_agent_to_site(&state.db, id, req.site_id)
        .await
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
        .ok_or((StatusCode::NOT_FOUND, "Agent not found".to_string()))?;

    Ok(Json(agent.into()))
}

/// List all agents with full details (site/client info)
/// Requires authentication.
pub async fn list_agents_with_details(
    State(state): State<AppState>,
    _user: AuthUser,
) -> Result<Json<Vec<db::AgentWithDetails>>, (StatusCode, String)> {
    let agents = db::get_all_agents_with_details(&state.db)
        .await
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

    Ok(Json(agents))
}

/// List unassigned agents (not belonging to any site)
/// Requires authentication.
pub async fn list_unassigned_agents(
    State(state): State<AppState>,
    _user: AuthUser,
) -> Result<Json<Vec<AgentResponse>>, (StatusCode, String)> {
    let agents = db::get_unassigned_agents(&state.db)
        .await
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

    let responses: Vec<AgentResponse> = agents.into_iter().map(|a| a.into()).collect();
    Ok(Json(responses))
}

/// Get extended state for an agent (network interfaces, uptime, etc.)
/// Requires authentication.
pub async fn get_agent_state(
    State(state): State<AppState>,
    _user: AuthUser,
    Path(id): Path<Uuid>,
) -> Result<Json<db::AgentState>, (StatusCode, String)> {
    let agent_state = db::get_agent_state(&state.db, id)
        .await
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
        .ok_or((StatusCode::NOT_FOUND, "Agent state not found".to_string()))?;

    Ok(Json(agent_state))
}

// ============================================================================
// Legacy Agent Endpoints (PowerShell agent for 2008 R2)
// ============================================================================

/// Request to register a legacy agent with site code
#[derive(Debug, Deserialize)]
pub struct RegisterLegacyRequest {
    pub site_code: String,
    pub hostname: String,
    pub os_type: String,
    pub os_version: Option<String>,
    pub agent_version: Option<String>,
    pub agent_type: Option<String>,
}

/// Response for legacy agent registration
#[derive(Debug, Serialize)]
pub struct RegisterLegacyResponse {
    pub agent_id: Uuid,
    pub api_key: String,
    pub site_name: String,
    pub client_name: String,
    pub message: String,
}

/// Register a legacy agent using site code
pub async fn register_legacy(
    State(state): State<AppState>,
    Json(req): Json<RegisterLegacyRequest>,
) -> Result<Json<RegisterLegacyResponse>, (StatusCode, String)> {
    // Look up site by code
    let site = db::get_site_by_code(&state.db, &req.site_code)
        .await
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
        .ok_or((StatusCode::NOT_FOUND, format!("Site code '{}' not found", req.site_code)))?;

    // Get client info
    let client = db::get_client_by_id(&state.db, site.client_id)
        .await
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
        .ok_or((StatusCode::NOT_FOUND, "Client not found".to_string()))?;

    // Generate API key for this agent
    let api_key = generate_api_key(&state.config.auth.api_key_prefix);
    let api_key_hash = hash_api_key(&api_key);

    // Create the agent
    let create = db::CreateAgent {
        hostname: req.hostname,
        api_key_hash,
        os_type: req.os_type,
        os_version: req.os_version,
        agent_version: req.agent_version,
    };

    let agent = db::create_agent(&state.db, create)
        .await
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

    // Assign agent to site
    db::move_agent_to_site(&state.db, agent.id, Some(site.id))
        .await
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

    tracing::info!(
        "Legacy agent registered: {} ({}) -> Site: {} ({})",
        agent.hostname,
        agent.id,
        site.name,
        req.site_code
    );

    Ok(Json(RegisterLegacyResponse {
        agent_id: agent.id,
        api_key,
        site_name: site.name,
        client_name: client.name,
        message: "Agent registered successfully".to_string(),
    }))
}

/// Heartbeat request from legacy agent
#[derive(Debug, Deserialize)]
pub struct HeartbeatRequest {
    pub agent_id: Uuid,
    pub timestamp: String,
    pub system_info: serde_json::Value,
}

/// Heartbeat response with pending commands
#[derive(Debug, Serialize)]
pub struct HeartbeatResponse {
    pub success: bool,
    pub pending_commands: Vec<PendingCommand>,
}

#[derive(Debug, Serialize)]
pub struct PendingCommand {
    pub id: Uuid,
    #[serde(rename = "type")]
    pub cmd_type: String,
    pub script: String,
}

/// Receive heartbeat from legacy agent
pub async fn heartbeat(
    State(state): State<AppState>,
    Json(req): Json<HeartbeatRequest>,
) -> Result<Json<HeartbeatResponse>, (StatusCode, String)> {
    // Update agent last_seen
    db::update_agent_last_seen(&state.db, req.agent_id)
        .await
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

    // TODO: Store system_info metrics, get pending commands
    // For now, return empty pending commands
    Ok(Json(HeartbeatResponse {
        success: true,
        pending_commands: vec![],
    }))
}

/// Command result from legacy agent
#[derive(Debug, Deserialize)]
pub struct CommandResultRequest {
    pub command_id: Uuid,
    pub started_at: String,
    pub completed_at: String,
    pub success: bool,
    pub output: String,
    pub error: Option<String>,
}

/// Receive command execution result
pub async fn command_result(
    State(_state): State<AppState>,
    Json(_req): Json<CommandResultRequest>,
) -> Result<StatusCode, (StatusCode, String)> {
    // TODO: Store command result in database
    Ok(StatusCode::OK)
}