use async_trait::async_trait; use axum::{ extract::{State, Path, WebSocketUpgrade}, http::StatusCode, response::{Html, Json, Response}, routing::{get, post, put, delete}, Router, }; use std::sync::Arc; use tower_http::cors::CorsLayer; use sea_orm::DatabaseConnection; use crate::core::context::AppContext; use crate::core::module::Module; use crate::error::Result; use crate::node_manager::NodeManager; use crate::settings_manager::SettingsManager; use crate::yggdrasil::{Node, YggdrasilConfig}; #[derive(Clone)] struct AppState { node_manager: Arc, context: Arc, } pub struct WebModule { name: String, context: Option>, node_manager: Arc, settings_manager: Arc, } impl WebModule { pub fn new(db: DatabaseConnection, settings_manager: SettingsManager) -> Self { Self { name: "web".to_string(), context: None, node_manager: Arc::new(NodeManager::new(db)), settings_manager: Arc::new(settings_manager), } } } #[async_trait] impl Module for WebModule { fn name(&self) -> &str { &self.name } async fn init(&mut self, context: Arc) -> Result<()> { self.context = Some(context); tracing::info!("Web module initialized"); Ok(()) } async fn start(&self) -> Result<()> { let context = self.context.as_ref().unwrap(); let config = context.config_manager.get(); let port = config.server.port; tracing::info!("Starting web server on port {}", port); let app_state = AppState { node_manager: self.node_manager.clone(), context: context.clone(), }; let app = Router::new() .route("/", get(index_handler)) .route("/edit/:id", get(edit_page_handler)) .route("/api/nodes", get(get_nodes_handler)) .route("/api/nodes", post(add_node_handler)) .route("/api/nodes/:id", get(get_node_handler)) .route("/api/nodes/:id", put(update_node_handler)) .route("/api/nodes/:id", delete(delete_node_handler)) .route("/api/configs", get(get_configs_handler)) .route("/api/nodes/:id/config", get(get_node_config_handler)) .route("/api/settings/listen-template", get(get_listen_template_handler)) .route("/api/settings/listen-template", put(update_listen_template_handler)) .route("/ws/agent", get(ws_agent_handler)) .layer(CorsLayer::permissive()) .with_state(app_state); let bind_addr = format!("{}:{}", config.server.bind_address, port); let listener = tokio::net::TcpListener::bind(&bind_addr) .await .map_err(|e| crate::error::AppError::Io(e))?; tokio::spawn(async move { axum::serve(listener, app) .await .expect("Failed to run web server"); }); Ok(()) } async fn stop(&self) -> Result<()> { tracing::info!("Web module stopped"); Ok(()) } } async fn index_handler() -> Html<&'static str> { Html(include_str!("../../static/index.html")) } #[derive(serde::Serialize)] struct NodesResponse { nodes: Vec, } async fn get_nodes_handler( State(app_state): State, ) -> Json { let nodes = app_state.node_manager.get_all_nodes().await; Json(NodesResponse { nodes }) } #[derive(serde::Deserialize)] struct AddNodeRequest { name: String, listen: Vec, addresses: Vec, } #[derive(serde::Serialize)] struct AddNodeResponse { success: bool, message: String, } async fn add_node_handler( State(app_state): State, Json(payload): Json, ) -> Json { match app_state.node_manager.add_node(payload.name, payload.listen, payload.addresses).await { Ok(_) => { // Broadcast update to all connected agents crate::websocket_state::broadcast_configuration_update(&app_state.node_manager).await; Json(AddNodeResponse { success: true, message: "Node added successfully".to_string(), }) } Err(e) => Json(AddNodeResponse { success: false, message: format!("Failed to add node: {}", e), }), } } #[derive(serde::Serialize)] struct ConfigsResponse { configs: Vec, } #[derive(serde::Serialize)] struct NodeConfig { node_id: String, node_name: String, node_addresses: Vec, config: YggdrasilConfig, } async fn get_configs_handler( State(app_state): State, ) -> Json { let nodes = app_state.node_manager.get_all_nodes().await; let configs_map = app_state.node_manager.generate_configs().await; let mut configs = Vec::new(); for node in nodes { if let Some(config) = configs_map.get(&node.id) { configs.push(NodeConfig { node_id: node.id.clone(), node_name: node.name.clone(), node_addresses: node.addresses.clone(), config: config.clone(), }); } } Json(ConfigsResponse { configs }) } // Get single node handler async fn get_node_handler( State(app_state): State, Path(node_id): Path, ) -> std::result::Result, StatusCode> { match app_state.node_manager.get_node_by_id(&node_id).await { Some(node) => Ok(Json(node)), None => Err(StatusCode::NOT_FOUND), } } // Update node handler async fn update_node_handler( State(app_state): State, Path(node_id): Path, Json(payload): Json, ) -> std::result::Result, StatusCode> { match app_state.node_manager.update_node(&node_id, payload.name, payload.listen, payload.addresses).await { Ok(_) => { // Broadcast update to all connected agents crate::websocket_state::broadcast_configuration_update(&app_state.node_manager).await; Ok(Json(AddNodeResponse { success: true, message: "Node updated successfully".to_string(), })) } Err(e) => { if e.to_string().contains("Node not found") { Err(StatusCode::NOT_FOUND) } else { Ok(Json(AddNodeResponse { success: false, message: format!("Failed to update node: {}", e), })) } } } } // Delete node handler async fn delete_node_handler( State(app_state): State, Path(node_id): Path, ) -> std::result::Result, StatusCode> { match app_state.node_manager.remove_node(&node_id).await { Ok(_) => { // Broadcast update to all connected agents crate::websocket_state::broadcast_configuration_update(&app_state.node_manager).await; Ok(Json(AddNodeResponse { success: true, message: "Node deleted successfully".to_string(), })) } Err(e) => { if e.to_string().contains("Node not found") { Err(StatusCode::NOT_FOUND) } else { Ok(Json(AddNodeResponse { success: false, message: format!("Failed to delete node: {}", e), })) } } } } // Get node configuration for agent async fn get_node_config_handler( State(app_state): State, Path(node_id): Path, ) -> std::result::Result, StatusCode> { // Get the node let node = match app_state.node_manager.get_node_by_id(&node_id).await { Some(node) => node, None => return Err(StatusCode::NOT_FOUND), }; // Generate configurations for all nodes let configs_map = app_state.node_manager.generate_configs().await; // Get config for this specific node match configs_map.get(&node_id) { Some(config) => Ok(Json(NodeConfig { node_id: node.id.clone(), node_name: node.name.clone(), node_addresses: node.addresses.clone(), config: config.clone(), })), None => Err(StatusCode::INTERNAL_SERVER_ERROR), } } // WebSocket handler for agents async fn ws_agent_handler( ws: WebSocketUpgrade, State(app_state): State, ) -> Response { ws.on_upgrade(move |socket| crate::modules::websocket::handle_agent_socket(socket, app_state.node_manager, app_state.context)) } // Edit page handler async fn edit_page_handler(Path(node_id): Path) -> Html { let html = include_str!("../../static/edit.html"); let content = html.replace("{{NODE_ID}}", &node_id); Html(content) } // Listen template handlers #[derive(serde::Serialize, serde::Deserialize)] struct ListenTemplateResponse { template: Vec, } #[derive(serde::Deserialize)] struct UpdateListenTemplateRequest { template: Vec, } async fn get_listen_template_handler( State(app_state): State, ) -> Json { match app_state.context.settings_manager.get_listen_template().await { Ok(template) => Json(ListenTemplateResponse { template }), Err(e) => { tracing::error!("Failed to get listen template from database: {}", e); // Return fallback default Json(ListenTemplateResponse { template: vec!["tcp://0.0.0.0:9001".to_string()], }) } } } async fn update_listen_template_handler( State(app_state): State, Json(payload): Json, ) -> Json { tracing::info!("Listen template update request: {:?}", payload.template); // Save to database match app_state.context.settings_manager.set_listen_template(payload.template.clone()).await { Ok(_) => { // Update in-memory config app_state.context.config_manager.update_listen_template(payload.template); Json(serde_json::json!({ "success": true, "message": "Listen template updated successfully" })) } Err(e) => { tracing::error!("Failed to save listen template: {}", e); Json(serde_json::json!({ "success": false, "message": format!("Failed to save template: {}", e) })) } } }