Added agent

This commit is contained in:
Ultradesu
2025-08-26 13:08:43 +03:00
parent 98ba8846c3
commit a2eabe2c69
13 changed files with 1089 additions and 105 deletions

View File

@@ -14,20 +14,29 @@ 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<NodeManager>,
context: Arc<AppContext>,
}
pub struct WebModule {
name: String,
context: Option<Arc<AppContext>>,
node_manager: Arc<NodeManager>,
settings_manager: Arc<SettingsManager>,
}
impl WebModule {
pub fn new(db: DatabaseConnection) -> Self {
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),
}
}
}
@@ -51,10 +60,14 @@ impl Module for WebModule {
tracing::info!("Starting web server on port {}", port);
let node_manager = self.node_manager.clone();
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))
@@ -62,9 +75,11 @@ impl Module for WebModule {
.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(node_manager);
.with_state(app_state);
let bind_addr = format!("{}:{}", config.server.bind_address, port);
let listener = tokio::net::TcpListener::bind(&bind_addr)
@@ -96,9 +111,9 @@ struct NodesResponse {
}
async fn get_nodes_handler(
State(node_manager): State<Arc<NodeManager>>,
State(app_state): State<AppState>,
) -> Json<NodesResponse> {
let nodes = node_manager.get_all_nodes().await;
let nodes = app_state.node_manager.get_all_nodes().await;
Json(NodesResponse { nodes })
}
@@ -116,13 +131,13 @@ struct AddNodeResponse {
}
async fn add_node_handler(
State(node_manager): State<Arc<NodeManager>>,
State(app_state): State<AppState>,
Json(payload): Json<AddNodeRequest>,
) -> Json<AddNodeResponse> {
match node_manager.add_node(payload.name, payload.listen, payload.addresses).await {
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(&node_manager).await;
crate::websocket_state::broadcast_configuration_update(&app_state.node_manager).await;
Json(AddNodeResponse {
success: true,
@@ -150,10 +165,10 @@ struct NodeConfig {
}
async fn get_configs_handler(
State(node_manager): State<Arc<NodeManager>>,
State(app_state): State<AppState>,
) -> Json<ConfigsResponse> {
let nodes = node_manager.get_all_nodes().await;
let configs_map = node_manager.generate_configs().await;
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 {
@@ -172,10 +187,10 @@ async fn get_configs_handler(
// Get single node handler
async fn get_node_handler(
State(node_manager): State<Arc<NodeManager>>,
State(app_state): State<AppState>,
Path(node_id): Path<String>,
) -> std::result::Result<Json<Node>, StatusCode> {
match node_manager.get_node_by_id(&node_id).await {
match app_state.node_manager.get_node_by_id(&node_id).await {
Some(node) => Ok(Json(node)),
None => Err(StatusCode::NOT_FOUND),
}
@@ -183,14 +198,14 @@ async fn get_node_handler(
// Update node handler
async fn update_node_handler(
State(node_manager): State<Arc<NodeManager>>,
State(app_state): State<AppState>,
Path(node_id): Path<String>,
Json(payload): Json<AddNodeRequest>,
) -> std::result::Result<Json<AddNodeResponse>, StatusCode> {
match node_manager.update_node(&node_id, payload.name, payload.listen, payload.addresses).await {
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(&node_manager).await;
crate::websocket_state::broadcast_configuration_update(&app_state.node_manager).await;
Ok(Json(AddNodeResponse {
success: true,
@@ -212,13 +227,13 @@ async fn update_node_handler(
// Delete node handler
async fn delete_node_handler(
State(node_manager): State<Arc<NodeManager>>,
State(app_state): State<AppState>,
Path(node_id): Path<String>,
) -> std::result::Result<Json<AddNodeResponse>, StatusCode> {
match node_manager.remove_node(&node_id).await {
match app_state.node_manager.remove_node(&node_id).await {
Ok(_) => {
// Broadcast update to all connected agents
crate::websocket_state::broadcast_configuration_update(&node_manager).await;
crate::websocket_state::broadcast_configuration_update(&app_state.node_manager).await;
Ok(Json(AddNodeResponse {
success: true,
@@ -240,17 +255,17 @@ async fn delete_node_handler(
// Get node configuration for agent
async fn get_node_config_handler(
State(node_manager): State<Arc<NodeManager>>,
State(app_state): State<AppState>,
Path(node_id): Path<String>,
) -> std::result::Result<Json<NodeConfig>, StatusCode> {
// Get the node
let node = match node_manager.get_node_by_id(&node_id).await {
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 = node_manager.generate_configs().await;
let configs_map = app_state.node_manager.generate_configs().await;
// Get config for this specific node
match configs_map.get(&node_id) {
@@ -267,7 +282,67 @@ async fn get_node_config_handler(
// WebSocket handler for agents
async fn ws_agent_handler(
ws: WebSocketUpgrade,
State(node_manager): State<Arc<NodeManager>>,
State(app_state): State<AppState>,
) -> Response {
ws.on_upgrade(move |socket| crate::modules::websocket::handle_agent_socket(socket, node_manager))
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<String>) -> Html<String> {
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<String>,
}
#[derive(serde::Deserialize)]
struct UpdateListenTemplateRequest {
template: Vec<String>,
}
async fn get_listen_template_handler(
State(app_state): State<AppState>,
) -> Json<ListenTemplateResponse> {
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<AppState>,
Json(payload): Json<UpdateListenTemplateRequest>,
) -> Json<serde_json::Value> {
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)
}))
}
}
}