| 1 | from pathlib import Path |
| 2 | from typing import Any |
| 3 | from typing import Dict |
| 4 | |
| 5 | from fastapi import HTTPException |
| 6 | from loguru import logger |
| 7 | |
| 8 | from app.agents.routes.agents import get_wazuh_manager_version |
| 9 | from app.connectors.portainer.schema.stack import DeleteStackResponse |
| 10 | from app.connectors.portainer.schema.stack import StackResponse |
| 11 | from app.connectors.portainer.schema.stack import StacksResponse |
| 12 | from app.connectors.portainer.schema.stack import StackStatus |
| 13 | from app.connectors.portainer.utils.universal import get_endpoint_id |
| 14 | from app.connectors.portainer.utils.universal import get_swarm_id |
| 15 | from app.connectors.portainer.utils.universal import send_delete_request |
| 16 | from app.connectors.portainer.utils.universal import send_get_request |
| 17 | from app.connectors.portainer.utils.universal import send_post_request |
| 18 | from app.customer_provisioning.schema.provision import ProvisionNewCustomer |
| 19 | from app.customer_provisioning.services.portainer import list_node_ips |
| 20 | |
| 21 | |
| 22 | async def get_stacks() -> StackResponse: |
| 23 | """ |
| 24 | Get the list of stacks from Portainer. |
| 25 | |
| 26 | Returns: |
| 27 | StackResponse: The response from Portainer |
| 28 | """ |
| 29 | response = await send_get_request("/api/stacks") |
| 30 | logger.info(f"Stacks received: {response}") |
| 31 | return StacksResponse(data=response["data"], message="Stacks fetched successfully", success=True) |
| 32 | |
| 33 | |
| 34 | async def get_stack_details(stack_id: int) -> StackResponse: |
| 35 | """ |
| 36 | Get the details of a stack from Portainer. |
| 37 | |
| 38 | Args: |
| 39 | stack_id (int): The ID of the stack |
| 40 | |
| 41 | Returns: |
| 42 | StackResponse: The response from Portainer |
| 43 | """ |
| 44 | endpoint_id = await get_endpoint_id() |
| 45 | response = await send_get_request(f"/api/stacks/{stack_id}?endpointId={endpoint_id}") |
| 46 | logger.info(f"Stack details received: {response}") |
| 47 | return StackResponse(**response) |
| 48 | |
| 49 | |
| 50 | async def _load_stack_template(template_path: Path) -> str: |
| 51 | """ |
| 52 | Load the stack template from file. |
| 53 | |
| 54 | Args: |
| 55 | template_path (Path): Path to the template file |
| 56 | |
| 57 | Returns: |
| 58 | str: Contents of the template file |
| 59 | """ |
| 60 | with open(template_path, "r") as file: |
| 61 | return file.read() |
| 62 | |
| 63 | |
| 64 | # async def _prepare_template_variables(request: ProvisionNewCustomer) -> Dict[str, str]: |
| 65 | # """ |
| 66 | # Prepare variables for template replacement. |
| 67 | |
| 68 | # Args: |
| 69 | # request (ProvisionNewCustomer): The customer provisioning request |
| 70 | |
| 71 | # Returns: |
| 72 | # Dict[str, str]: Dictionary of template variables and their values |
| 73 | # """ |
| 74 | # formatted_customer_name = request.customer_name.replace(" ", "_") |
| 75 | # wazuh_manager_version = await get_wazuh_manager_version() |
| 76 | |
| 77 | # return { |
| 78 | # "{{ wazuh_worker_customer_code }}": formatted_customer_name, |
| 79 | # "{{ wazuh_manager_version }}": wazuh_manager_version, |
| 80 | # "REPLACE_LOG": request.wazuh_logs_port, |
| 81 | # "REPLACE_REGISTRATION": request.wazuh_registration_port, |
| 82 | # "REPLACE_API": request.wazuh_api_port, |
| 83 | # "customer_name": formatted_customer_name, |
| 84 | # } |
| 85 | |
| 86 | |
| 87 | async def _prepare_template_variables(request: ProvisionNewCustomer, node_count: int) -> Dict[str, str]: |
| 88 | """ |
| 89 | Prepare variables for template replacement. |
| 90 | |
| 91 | Args: |
| 92 | request (ProvisionNewCustomer): The customer provisioning request |
| 93 | node_count (int): Number of nodes in the swarm |
| 94 | |
| 95 | Returns: |
| 96 | Dict[str, str]: Dictionary of template variables and their values |
| 97 | """ |
| 98 | formatted_customer_name = request.customer_name.replace(" ", "_") |
| 99 | wazuh_manager_version = await get_wazuh_manager_version() |
| 100 | formatted_customer_code = request.customer_code.replace(" ", "_") |
| 101 | |
| 102 | return { |
| 103 | "{{ wazuh_worker_customer_code }}": formatted_customer_code, |
| 104 | "{{ wazuh_manager_version }}": wazuh_manager_version, |
| 105 | "REPLACE_LOG": request.wazuh_logs_port, |
| 106 | "REPLACE_REGISTRATION": request.wazuh_registration_port, |
| 107 | "REPLACE_API": request.wazuh_api_port, |
| 108 | "NUMBER_OF_NODES": str(node_count), |
| 109 | "customer_name": formatted_customer_name, |
| 110 | "customer_code": formatted_customer_code, |
| 111 | } |
| 112 | |
| 113 | |
| 114 | async def _create_stack_payload(template: str, variables: Dict[str, str], swarm_id: str) -> Dict[str, Any]: |
| 115 | """ |
| 116 | Create the payload for stack creation. |
| 117 | |
| 118 | Args: |
| 119 | template (str): The processed template |
| 120 | variables (Dict[str, str]): Template variables |
| 121 | swarm_id (str): The swarm ID |
| 122 | |
| 123 | Returns: |
| 124 | Dict[str, Any]: The payload for stack creation |
| 125 | """ |
| 126 | return { |
| 127 | "Name": f"wazuh-worker-{variables['customer_code']}", |
| 128 | "StackFileContent": template, |
| 129 | "SwarmID": swarm_id, |
| 130 | "Env": [], |
| 131 | } |
| 132 | |
| 133 | |
| 134 | # async def create_wazuh_customer_stack(request: ProvisionNewCustomer) -> StackResponse: |
| 135 | # """ |
| 136 | # Create a Wazuh stack for a customer. |
| 137 | |
| 138 | # Args: |
| 139 | # request (ProvisionNewCustomer): The customer provisioning request |
| 140 | |
| 141 | # Returns: |
| 142 | # StackResponse: The response from Portainer stack creation |
| 143 | # """ |
| 144 | # logger.info(f"Creating Wazuh stack for customer {request.customer_name}") |
| 145 | |
| 146 | # # Load template |
| 147 | # template_path = Path(__file__).parent.parent / "templates" / "wazuh_worker_stack.yml" |
| 148 | # template = await _load_stack_template(template_path) |
| 149 | |
| 150 | # # Prepare variables |
| 151 | # variables = await _prepare_template_variables(request) |
| 152 | # logger.info(f"Template variables prepared for customer: {variables['customer_name']}") |
| 153 | |
| 154 | # # Process template |
| 155 | # for placeholder, value in variables.items(): |
| 156 | # template = template.replace(placeholder, value) |
| 157 | # logger.info("Template processed with variables") |
| 158 | |
| 159 | # # Get required IDs |
| 160 | # endpoint_id = await get_endpoint_id() |
| 161 | # swarm_id = await get_swarm_id() |
| 162 | # logger.info(f"Retrieved endpoint ID: {endpoint_id} and swarm ID: {swarm_id}") |
| 163 | |
| 164 | # # Create and send request |
| 165 | # create_stack_url = f"/api/stacks?type=1&method=string&endpointId={endpoint_id}" |
| 166 | # payload = await _create_stack_payload(template, variables, swarm_id) |
| 167 | |
| 168 | # response = await send_post_request(endpoint=create_stack_url, data=payload) |
| 169 | # logger.info(f"Stack creation response received: {response}") |
| 170 | |
| 171 | # return StackResponse(**response) |
| 172 | |
| 173 | |
| 174 | async def create_wazuh_customer_stack(request: ProvisionNewCustomer) -> StackResponse: |
| 175 | """ |
| 176 | Create a Wazuh stack for a customer. |
| 177 | |
| 178 | Args: |
| 179 | request (ProvisionNewCustomer): The customer provisioning request |
| 180 | |
| 181 | Returns: |
| 182 | StackResponse: The response from Portainer stack creation |
| 183 | """ |
| 184 | logger.info(f"Creating Wazuh stack for customer {request.customer_name}") |
| 185 | |
| 186 | # Get the number of swarm nodes |
| 187 | swarm_node_ips = await list_node_ips() |
| 188 | node_count = len(swarm_node_ips) |
| 189 | logger.info(f"Found {node_count} swarm nodes: {swarm_node_ips}") |
| 190 | |
| 191 | # Load template |
| 192 | template_path = Path(__file__).parent.parent / "templates" / "wazuh_worker_stack.yml" |
| 193 | template = await _load_stack_template(template_path) |
| 194 | |
| 195 | # Prepare variables with node count |
| 196 | variables = await _prepare_template_variables(request, node_count) |
| 197 | logger.info(f"Template variables prepared for customer: {variables['customer_name']} with {node_count} nodes") |
| 198 | |
| 199 | # Process template |
| 200 | for placeholder, value in variables.items(): |
| 201 | template = template.replace(placeholder, value) |
| 202 | logger.info(f"Template processed with variables, replicas set to {node_count}") |
| 203 | |
| 204 | # Get required IDs |
| 205 | endpoint_id = await get_endpoint_id() |
| 206 | swarm_id = await get_swarm_id() |
| 207 | logger.info(f"Retrieved endpoint ID: {endpoint_id} and swarm ID: {swarm_id}") |
| 208 | |
| 209 | # Create and send request |
| 210 | create_stack_url = f"/api/stacks?type=1&method=string&endpointId={endpoint_id}" |
| 211 | payload = await _create_stack_payload(template, variables, swarm_id) |
| 212 | |
| 213 | response = await send_post_request(endpoint=create_stack_url, data=payload) |
| 214 | logger.info(f"Stack creation response received: {response}") |
| 215 | |
| 216 | return StackResponse(**response) |
| 217 | |
| 218 | |
| 219 | async def start_wazuh_customer_stack(stack_id: int) -> StackResponse: |
| 220 | """ |
| 221 | Start a Wazuh stack for a customer. |
| 222 | |
| 223 | Args: |
| 224 | stack_id (int): The ID of the stack to start |
| 225 | |
| 226 | Returns: |
| 227 | StackResponse: The response from Portainer stack start |
| 228 | |
| 229 | Raises: |
| 230 | HTTPException: If the stack is already active or in an unexpected state |
| 231 | """ |
| 232 | logger.info(f"Checking status of stack {stack_id} before starting") |
| 233 | |
| 234 | # Get current stack status |
| 235 | stack_details = await get_stack_details(stack_id) |
| 236 | |
| 237 | # Check if stack is already active |
| 238 | if stack_details.data.Status == StackStatus.ACTIVE: |
| 239 | logger.info(f"Stack {stack_id} is already active") |
| 240 | return stack_details |
| 241 | |
| 242 | # If stack is stopped, proceed with starting it |
| 243 | if stack_details.data.Status == StackStatus.DOWN: |
| 244 | logger.info(f"Starting stopped stack {stack_id}") |
| 245 | endpoint_id = await get_endpoint_id() |
| 246 | |
| 247 | # Create and send request |
| 248 | start_stack_url = f"/api/stacks/{stack_id}/start?endpointId={endpoint_id}" |
| 249 | response = await send_post_request(endpoint=start_stack_url) |
| 250 | logger.info(f"Stack start response received: {response}") |
| 251 | |
| 252 | return StackResponse(**response) |
| 253 | |
| 254 | # If stack is in any other state |
| 255 | logger.warning(f"Stack {stack_id} is in an unexpected state: {stack_details.data.Status}") |
| 256 | raise HTTPException(status_code=400, detail=f"Unexpected stack status: {stack_details.data.Status}") |
| 257 | |
| 258 | |
| 259 | async def stop_wazuh_customer_stack(stack_id: int) -> StackResponse: |
| 260 | """ |
| 261 | Stop a Wazuh stack for a customer. |
| 262 | |
| 263 | Args: |
| 264 | stack_id (int): The ID of the stack to stop |
| 265 | |
| 266 | Returns: |
| 267 | StackResponse: The response from Portainer stack stop |
| 268 | |
| 269 | Raises: |
| 270 | HTTPException: If the stack is already stopped |
| 271 | """ |
| 272 | logger.info(f"Checking status of stack {stack_id} before stopping") |
| 273 | |
| 274 | # Get current stack status |
| 275 | stack_details = await get_stack_details(stack_id) |
| 276 | |
| 277 | # Check if stack is already stopped |
| 278 | if stack_details.data.Status == StackStatus.DOWN: |
| 279 | logger.info(f"Stack {stack_id} is already stopped") |
| 280 | return stack_details |
| 281 | |
| 282 | # If stack is active, proceed with stopping it |
| 283 | if stack_details.data.Status == StackStatus.ACTIVE: |
| 284 | logger.info(f"Stopping active stack {stack_id}") |
| 285 | endpoint_id = await get_endpoint_id() |
| 286 | |
| 287 | # Create and send request |
| 288 | stop_stack_url = f"/api/stacks/{stack_id}/stop?endpointId={endpoint_id}" |
| 289 | response = await send_post_request(endpoint=stop_stack_url) |
| 290 | logger.info(f"Stack stop response received: {response}") |
| 291 | |
| 292 | return StackResponse(**response) |
| 293 | |
| 294 | # If stack is in any other state |
| 295 | logger.warning(f"Stack {stack_id} is in an unexpected state: {stack_details.data.Status}") |
| 296 | raise HTTPException(status_code=400, detail=f"Unexpected stack status: {stack_details.data.Status}") |
| 297 | |
| 298 | |
| 299 | async def delete_wazuh_customer_stack(stack_id: int) -> DeleteStackResponse: |
| 300 | """ |
| 301 | Delete a Wazuh stack for a customer. |
| 302 | |
| 303 | Args: |
| 304 | request (ProvisionNewCustomer): The customer provisioning request |
| 305 | |
| 306 | Returns: |
| 307 | DeleteStackResponse: The response from Portainer stack deletion |
| 308 | """ |
| 309 | logger.info(f"Deleting Wazuh stack for stack id {stack_id}") |
| 310 | |
| 311 | # Get required IDs |
| 312 | endpoint_id = await get_endpoint_id() |
| 313 | logger.info(f"Retrieved endpoint ID: {endpoint_id} and stack ID: {stack_id}") |
| 314 | |
| 315 | # Stop the stack first |
| 316 | await stop_wazuh_customer_stack(stack_id) |
| 317 | |
| 318 | # Create and send request |
| 319 | delete_stack_url = f"/api/stacks/{stack_id}" |
| 320 | response = await send_delete_request(endpoint=delete_stack_url, params={"endpointId": endpoint_id}) |
| 321 | logger.info(f"Stack deletion response received: {response}") |
| 322 | |
| 323 | return DeleteStackResponse(**response) |