//! 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, } /// Register a new agent (generates API key) /// Requires authentication to prevent unauthorized agent registration. pub async fn register_agent( State(state): State, user: AuthUser, Json(req): Json, ) -> Result, (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, _user: AuthUser, ) -> Result>, (StatusCode, String)> { let agents = db::get_all_agents(&state.db) .await .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; let responses: Vec = 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, _user: AuthUser, Path(id): Path, ) -> Result, (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, _user: AuthUser, Path(id): Path, ) -> Result { // 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, _user: AuthUser, ) -> Result, (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, // None to unassign from site } /// Move an agent to a different site /// Requires authentication. pub async fn move_agent( State(state): State, _user: AuthUser, Path(id): Path, Json(req): Json, ) -> Result, (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, _user: AuthUser, ) -> Result>, (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, _user: AuthUser, ) -> Result>, (StatusCode, String)> { let agents = db::get_unassigned_agents(&state.db) .await .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; let responses: Vec = 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, _user: AuthUser, Path(id): Path, ) -> Result, (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, pub agent_version: Option, pub agent_type: Option, } /// 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, Json(req): Json, ) -> Result, (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, } #[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, Json(req): Json, ) -> Result, (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, } /// Receive command execution result pub async fn command_result( State(_state): State, Json(_req): Json, ) -> Result { // TODO: Store command result in database Ok(StatusCode::OK) }