added docker management

WatcherAgent/src/docker/serverclientcomm.rs

@@ -1,15 +1,13 @@
-
 //! Server-client communication utilities for WatcherAgent
 //!
 //! Handles server commands, Docker image updates, and container management using the Bollard library.
 //!
-use crate::models::{DockerContainer, ServerMessage};
-use crate::docker::container::{get_available_container};
+use crate::models::ServerMessage;
 
-use std::error::Error;
-use bollard::Docker;
+use super::container::{restart_container, update_docker_image};
 use bollard::query_parameters::{CreateImageOptions, RestartContainerOptions};
-use futures_util::StreamExt;
+use bollard::Docker;
+use std::error::Error;
 
 /// Handles a message from the backend server and dispatches the appropriate action.
 ///
@@ -19,7 +17,10 @@ use futures_util::StreamExt;
 ///
 /// # Returns
 /// * `Result<(), Box<dyn Error + Send + Sync>>` - Ok if handled successfully, error otherwise.
-pub async fn handle_server_message(docker: &Docker, msg: ServerMessage) -> Result<(), Box<dyn Error + Send + Sync>> {
+pub async fn handle_server_message(
+    docker: &Docker,
+    msg: ServerMessage,
+) -> Result<(), Box<dyn Error + Send + Sync>> {
     let msg = msg.clone();
     println!("Handling server message: {:?}", msg);
 
@@ -36,10 +37,14 @@ pub async fn handle_server_message(docker: &Docker, msg: ServerMessage) -> Resul
             }
         }
         "restart_container" => {
-            println!("Received restart container command");
-            // Call your restart_container function here
-            restart_container(docker).await?;
-            Ok(())
+            if let Some(image_name) = msg.data.get("image").and_then(|v| v.as_str()) {
+                println!("Received restart command for image: {}", image_name);
+                // Call your update_docker_image function here
+                update_docker_image(docker, image_name).await?;
+                Ok(())
+            } else {
+                Err("Missing image name in update message".into())
+            }
         }
         "stop_agent" => {
             println!("Received stop agent command");
@@ -52,87 +57,3 @@ pub async fn handle_server_message(docker: &Docker, msg: ServerMessage) -> Resul
         }
     }
 }
-
-/// Pulls a new Docker image and restarts the current container.
-///
-/// # Arguments
-/// * `docker` - Reference to a Bollard Docker client.
-/// * `image` - The name of the Docker image to pull.
-///
-/// # Returns
-/// * `Result<(), Box<dyn Error + Send + Sync>>` - Ok if updated successfully, error otherwise.
-pub async fn update_docker_image(docker: &Docker, image: &str) -> Result<(), Box<dyn Error + Send + Sync>> {
-    println!("Updating to {}", image);
-
-    // 1. Pull new image
-    let mut stream = docker.create_image(
-        Some(CreateImageOptions {
-            from_image: Some(image.to_string()),
-            ..Default::default()
-        }),
-        None,
-        None,
-    );
-
-    // Use the stream with proper trait bounds
-    while let Some(result) = StreamExt::next(&mut stream).await {
-        match result {
-            Ok(progress) => {
-                if let Some(status) = progress.status {
-                    println!("Pull status: {}", status);
-                }
-            }
-            Err(e) => {
-                eprintln!("Error pulling image: {}", e);
-                break;
-            }
-        }
-    }
-
-    // 2. Restart the current container
-    let _ = restart_container(docker).await;
-
-    Ok(())
-}
-
-/// Finds the Docker container running the agent by image name.
-///
-/// # Arguments
-/// * `docker` - Reference to a Bollard Docker client.
-///
-/// # Returns
-/// * `Result<Option<DockerContainer>, Box<dyn Error + Send + Sync>>` - The agent's container info if found.
-pub async fn get_client_container(docker: &Docker) -> Result<Option<DockerContainer>, Box<dyn Error + Send + Sync>> {
-    let containers = get_available_container(docker).await;
-    let client_image = "watcher-agent";
-    
-    // Find container with the specific image
-    if let Some(container) = containers.iter().find(|c| c.image == client_image) {
-        Ok(Some(container.clone()))
-    } else {
-        Ok(None)
-    }
-}
-
-/// Restarts the agent's own Docker container.
-///
-/// # Arguments
-/// * `docker` - Reference to a Bollard Docker client.
-///
-/// # Returns
-/// * `Result<(), Box<dyn Error + Send + Sync>>` - Ok if restarted successfully, error otherwise.
-pub async fn restart_container(docker: &Docker) -> Result<(), Box<dyn Error + Send + Sync>> {
-    if let Ok(Some(container)) = get_client_container(docker).await {
-        let container_id = container.clone().ID;
-        println!("Restarting container {}", container_id);
-        if let Err(e) = docker.restart_container(&container_id.to_string(), Some(RestartContainerOptions { signal: None, t: Some(0) }))
-            .await
-        {
-            eprintln!("Failed to restart container: {}", e);
-        }
-    } else {
-        eprintln!("No container ID found (HOSTNAME not set?)");
-    }
-
-    Ok(())
-}