@cryptotaxi247 / CoPilot / commits / 525e6e75

Copilot action (#499)

* Add InvokeCopilotActionBody model and refactor action invocation logic - Introduced InvokeCopilotActionBody for invoking Copilot actions with parameters. - Refactored invoke_action endpoint to utilize new model and improve modularity. - Enhanced error handling and parameter validation in action invocation process. * precommit-fixes * Add required parameters RepoURL and ScriptName to Copilot Actions; update repo_url type to str and enhance validation logic * Add optional arg_position field to ScriptParameter model * Enhance build_velociraptor_parameters to include script_params for arg_position mapping * precommit fixes

taylor_socfortress committed Sep 5, 2025 at 10:58 UTC 525e6e75cbd684e92e7f4832417c57ff8ec339fe
4 files changed +266 -69
backend/app/connectors/velociraptor/schema/artifacts.py
+12
@@ -151,6 +151,18 @@ class CollectArtifactBody(BaseBody):
151 }
152
153
154 +class InvokeCopilotActionBody(BaseModel):
155 + """Request body for invoking a Copilot action."""
156 +
157 + copilot_action_name: str = Field(..., description="Name of the action to invoke")
158 + agent_name: str = Field(..., description="Name of the agent to invoke the action on")
159 + artifact_name: Optional[str] = Field(None, description="Name of the artifact to use")
160 + parameters: Optional[Dict[str, Union[str, List[ParameterKeyValue]]]] = Field(
161 + None,
162 + description="Optional parameters for the action",
163 + )
164 +
165 +
166 class CollectFileBody(BaseBody):
167 artifact_name: str = Field(
168 "Generic.Collectors.File",
backend/app/integrations/copilot_action/routes/copilot_action.py
+236 -59
@@ -2,12 +2,21 @@ import os
2 from typing import Optional
3
4 from fastapi import APIRouter
5 +from fastapi import Depends
6 from fastapi import HTTPException
7 from fastapi import Query
8 from fastapi import Security
9 from loguru import logger
10 +from sqlalchemy.ext.asyncio import AsyncSession
11 +from sqlalchemy.future import select
12
13 from app.auth.routes.auth import AuthHandler
14 +from app.connectors.velociraptor.schema.artifacts import CollectArtifactBody
15 +from app.connectors.velociraptor.schema.artifacts import CollectArtifactResponse
16 +from app.connectors.velociraptor.schema.artifacts import InvokeCopilotActionBody
17 +from app.connectors.velociraptor.services.artifacts import run_artifact_collection
18 +from app.db.db_session import get_db
19 +from app.db.universal_models import Agents
20 from app.integrations.copilot_action.schema.copilot_action import ActionDetailResponse
21 from app.integrations.copilot_action.schema.copilot_action import (
22 InventoryMetricsResponse,
@@ -19,6 +28,152 @@ from app.integrations.copilot_action.services.copilot_action import CopilotActio
28 copilot_action_router = APIRouter()
29 auth_handler = AuthHandler()
30
31 +# Helper functions for better modularity
32 +
33 +
34 +def get_license_key() -> str:
35 + """Get and validate the COPILOT_API_KEY environment variable."""
36 + license_key = os.getenv("COPILOT_API_KEY")
37 + if not license_key:
38 + logger.error("COPILOT_API_KEY environment variable not set")
39 + raise HTTPException(status_code=500, detail="COPILOT_API_KEY environment variable not configured")
40 + return license_key
41 +
42 +
43 +async def get_agent_by_hostname(session: AsyncSession, hostname: str) -> Agents:
44 + """Retrieve agent from database by hostname."""
45 + agent_details = await session.execute(
46 + select(Agents).filter(Agents.hostname == hostname),
47 + )
48 + agent = agent_details.scalars().first()
49 +
50 + if not agent:
51 + raise HTTPException(
52 + status_code=404,
53 + detail=f"Agent with hostname {hostname} not found",
54 + )
55 +
56 + logger.info(f"Found agent: {agent.hostname} with OS: {agent.os}")
57 + return agent
58 +
59 +
60 +def determine_artifact_name(agent_os: str) -> str:
61 + """Determine the appropriate Velociraptor artifact based on OS."""
62 + if "Windows" in agent_os:
63 + return "Windows.Execute.RemotePowerShellScript"
64 + elif "Linux" in agent_os:
65 + return "Linux.Execute.RemoteBashScript"
66 + else:
67 + raise HTTPException(
68 + status_code=400,
69 + detail=f"Unsupported OS: {agent_os}",
70 + )
71 +
72 +
73 +def build_velociraptor_parameters(copilot_params: dict, script_params: list) -> dict:
74 + """
75 + Convert Copilot Action parameters to Velociraptor format.
76 +
77 + Args:
78 + copilot_params: Dictionary of parameters provided by the user
79 + script_params: List of ScriptParameter objects from the action details
80 +
81 + Returns:
82 + Dictionary with env array for Velociraptor
83 + """
84 + if not copilot_params:
85 + return {}
86 +
87 + env_array = []
88 +
89 + # Always add RepoURL and ScriptName first
90 + if "RepoURL" in copilot_params:
91 + env_array.append({"key": "RepoURL", "value": str(copilot_params["RepoURL"])})
92 + if "ScriptName" in copilot_params:
93 + env_array.append({"key": "ScriptName", "value": str(copilot_params["ScriptName"])})
94 +
95 + # Create a mapping of parameter names to their arg_position
96 + param_position_map = {}
97 + for param in script_params:
98 + if param.arg_position is not None:
99 + param_position_map[param.name] = param.arg_position
100 +
101 + # Add parameters with arg_position as Arg{position}
102 + for param_name, param_value in copilot_params.items():
103 + if param_name in param_position_map:
104 + arg_key = f"Arg{param_position_map[param_name]}"
105 + env_array.append({"key": arg_key, "value": str(param_value)})
106 +
107 + # Add other parameters (those without arg_position and not RepoURL/ScriptName)
108 + for param_name, param_value in copilot_params.items():
109 + if param_name not in param_position_map and param_name not in ["RepoURL", "ScriptName"]:
110 + env_array.append({"key": param_name, "value": str(param_value)})
111 +
112 + return {"env": env_array}
113 +
114 +
115 +async def validate_parameters(provided_params: dict, script_params: list) -> None:
116 + """
117 + Validate provided parameters against the script's expected parameters.
118 +
119 + Args:
120 + provided_params: Dictionary of parameters provided by the user
121 + script_params: List of ScriptParameter objects defining expected parameters
122 +
123 + Raises:
124 + HTTPException: If invalid or missing parameters are provided
125 + """
126 + if not provided_params:
127 + provided_params = {}
128 +
129 + logger.info(f"Validating {len(provided_params)} provided parameters against {len(script_params)} script parameters")
130 +
131 + # Extract parameter info from script
132 + valid_param_names = {param.name for param in script_params}
133 + required_params = {param.name for param in script_params if param.required}
134 + provided_param_keys = set(provided_params.keys())
135 +
136 + # Add RepoURL and ScriptName as required parameters for Copilot Actions
137 + required_params.add("RepoURL")
138 + required_params.add("ScriptName")
139 + valid_param_names.add("RepoURL")
140 + valid_param_names.add("ScriptName")
141 +
142 + logger.info(f"Valid parameters: {valid_param_names}")
143 + logger.info(f"Required parameters: {required_params}")
144 + logger.info(f"Provided parameters: {provided_param_keys}")
145 +
146 + # Validate: no invalid parameters
147 + invalid_params = provided_param_keys - valid_param_names
148 + if invalid_params:
149 + logger.error(f"Invalid parameters provided: {invalid_params}")
150 + raise HTTPException(
151 + status_code=400,
152 + detail=f"Invalid parameters provided: {list(invalid_params)}. Valid parameters are: {list(valid_param_names)}",
153 + )
154 +
155 + # Validate: all required parameters present
156 + missing_params = required_params - provided_param_keys
157 + if missing_params:
158 + logger.error(f"Missing required parameters: {missing_params}")
159 + raise HTTPException(status_code=400, detail=f"Missing required parameters: {list(missing_params)}")
160 +
161 + logger.info("Parameter validation successful")
162 +
163 +
164 +async def build_artifact_collection_body(agent: Agents, artifact_name: str, velociraptor_params: dict) -> CollectArtifactBody:
165 + """Build the artifact collection request body for Velociraptor."""
166 + return CollectArtifactBody(
167 + hostname=agent.hostname,
168 + velociraptor_id=agent.velociraptor_id,
169 + velociraptor_org=agent.velociraptor_org,
170 + artifact_name=artifact_name,
171 + parameters=velociraptor_params,
172 + )
173 +
174 +
175 +# Route handlers (keeping existing routes but updating invoke_action)
176 +
177
178 @copilot_action_router.get(
179 "/inventory",
@@ -36,34 +191,11 @@ async def get_inventory(
191 refresh: bool = Query(False, description="Force refresh cache"),
192 include: Optional[str] = Query(None, description="Comma-separated extra fields to include"),
193 ) -> InventoryResponse:
39 - """
40 - Retrieve inventory of available active response scripts.
41 -
42 - This endpoint fetches the catalog of active response scripts from the
43 - Copilot Action service, with optional filtering and pagination.
44 -
45 - Args:
46 - technology: Filter by technology type (e.g., Windows, Linux, Wazuh)
47 - category: Filter by category if present
48 - tag: Filter by tag contained in the tags list
49 - q: Free-text search in name/description
50 - limit: Maximum number of results (1-1000)
51 - offset: Offset for pagination
52 - refresh: Force refresh the remote cache
53 - include: Extra fields to include (e.g., 'category,tags')
54 -
55 - Returns:
56 - InventoryResponse: List of active response scripts with metadata
57 - """
194 + """Retrieve inventory of available active response scripts."""
195 logger.info(f"Fetching active response inventory with filters: tech={technology}, category={category}, tag={tag}, q={q}")
196
60 - # Get license key from environment variable
61 - license_key = os.getenv("COPILOT_API_KEY")
62 - if not license_key:
63 - logger.error("COPILOT_API_KEY environment variable not set")
64 - raise HTTPException(status_code=500, detail="COPILOT_API_KEY environment variable not configured")
197 + license_key = get_license_key()
198
66 - # Fetch inventory from service
199 try:
200 response = await CopilotActionService.get_inventory(
201 license_key=license_key,
@@ -92,24 +224,11 @@ async def get_inventory(
224 dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],
225 )
226 async def get_action_by_name(copilot_action_name: str) -> ActionDetailResponse:
95 - """
96 - Get detailed information for a specific active response script.
97 -
98 - Args:
99 - copilot_action_name: Name of the action to retrieve
100 -
101 - Returns:
102 - ActionDetailResponse: Detailed information about the action
103 - """
227 + """Get detailed information for a specific active response script."""
228 logger.info(f"Fetching action details for: {copilot_action_name}")
229
106 - # Get license key from environment variable
107 - license_key = os.getenv("COPILOT_API_KEY")
108 - if not license_key:
109 - logger.error("COPILOT_API_KEY environment variable not set")
110 - raise HTTPException(status_code=500, detail="COPILOT_API_KEY environment variable not configured")
230 + license_key = get_license_key()
231
112 - # Fetch action details from service
232 try:
233 response = await CopilotActionService.get_action_by_name(license_key=license_key, copilot_action_name=copilot_action_name)
234
@@ -120,6 +239,7 @@ async def get_action_by_name(copilot_action_name: str) -> ActionDetailResponse:
239 raise HTTPException(status_code=500, detail=response.message)
240
241 logger.info(f"Successfully fetched action details for: {copilot_action_name}")
242 + logger.info(f"Raw action feteched: {response}")
243 return response
244
245 except HTTPException:
@@ -136,24 +256,13 @@ async def get_action_by_name(copilot_action_name: str) -> ActionDetailResponse:
256 dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],
257 )
258 async def get_metrics() -> InventoryMetricsResponse:
139 - """
140 - Get metrics and status information for the inventory service.
141 -
142 - Returns:
143 - InventoryMetricsResponse: Service metrics and status
144 - """
259 + """Get metrics and status information for the inventory service."""
260 logger.info("Fetching inventory metrics")
261
147 - # Get license key from environment variable
148 - license_key = os.getenv("COPILOT_API_KEY")
149 - if not license_key:
150 - logger.error("COPILOT_API_KEY environment variable not set")
151 - raise HTTPException(status_code=500, detail="COPILOT_API_KEY environment variable not configured")
262 + license_key = get_license_key()
263
153 - # Fetch metrics from service
264 try:
265 response = await CopilotActionService.get_metrics(license_key=license_key)
156 -
266 logger.info("Successfully fetched inventory metrics")
267 return response
268
@@ -168,12 +277,7 @@ async def get_metrics() -> InventoryMetricsResponse:
277 dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],
278 )
279 async def get_technologies() -> dict:
171 - """
172 - Get list of available technology types for filtering.
173 -
174 - Returns:
175 - Dictionary containing available technology types
176 - """
280 + """Get list of available technology types for filtering."""
281 technologies = [tech.value for tech in Technology]
282
283 return {
@@ -182,3 +286,76 @@ async def get_technologies() -> dict:
286 "message": "Successfully retrieved available technologies",
287 "success": True,
288 }
289 +
290 +
291 +@copilot_action_router.post(
292 + "/invoke",
293 + response_model=CollectArtifactResponse,
294 + description="Invoke a Copilot Action on a target agent",
295 + dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],
296 +)
297 +async def invoke_action(body: InvokeCopilotActionBody, session: AsyncSession = Depends(get_db)) -> CollectArtifactResponse:
298 + """
299 + Invoke a Copilot Action on a target agent.
300 +
301 + This endpoint orchestrates the process of:
302 + 1. Finding the target agent
303 + 2. Determining the appropriate Velociraptor artifact
304 + 3. Fetching action details and validating parameters
305 + 4. Building and executing the artifact collection request
306 +
307 + Args:
308 + body: Request body containing action name, agent name, and parameters
309 + session: Database session
310 +
311 + Returns:
312 + CollectArtifactResponse: Response from the artifact collection
313 + """
314 + logger.info(f"Invoking Copilot action '{body.copilot_action_name}' on agent '{body.agent_name}'")
315 +
316 + try:
317 + # Step 1: Get the target agent
318 + agent = await get_agent_by_hostname(session, body.agent_name)
319 +
320 + # Step 2: Determine the appropriate artifact based on OS
321 + artifact_name = determine_artifact_name(agent.os)
322 + logger.info(f"Using artifact: {artifact_name} for OS: {agent.os}")
323 +
324 + # Step 3: Fetch action details
325 + copilot_action_details = await get_action_by_name(body.copilot_action_name)
326 + logger.info(f"Found action details for: {body.copilot_action_name}")
327 +
328 + # Add the `repo_url` and the `script_name` to the parameters
329 + if copilot_action_details.copilot_action.repo_url:
330 + if not body.parameters:
331 + body.parameters = {}
332 + body.parameters["RepoURL"] = copilot_action_details.copilot_action.repo_url
333 + if copilot_action_details.copilot_action.script_name:
334 + if not body.parameters:
335 + body.parameters = {}
336 + body.parameters["ScriptName"] = copilot_action_details.copilot_action.script_name
337 +
338 + logger.info(f"Parameters after adding repo and script: {body.parameters}")
339 +
340 + # Step 4: Validate parameters
341 + await validate_parameters(body.parameters or {}, copilot_action_details.copilot_action.script_parameters)
342 +
343 + # Step 5: Build Velociraptor parameters (now includes script_params for arg_position mapping)
344 + velociraptor_params = build_velociraptor_parameters(body.parameters or {}, copilot_action_details.copilot_action.script_parameters)
345 +
346 + # Step 6: Build artifact collection request
347 + artifact_body = await build_artifact_collection_body(agent, artifact_name, velociraptor_params)
348 + logger.info(f"Built artifact collection request for {agent.hostname}")
349 +
350 + # Step 7: Execute the collection
351 + response = await run_artifact_collection(artifact_body)
352 + logger.info(f"Successfully invoked Copilot action on {agent.hostname}")
353 +
354 + return response
355 +
356 + except HTTPException:
357 + # Re-raise HTTP exceptions (validation errors, not found, etc.)
358 + raise
359 + except Exception as e:
360 + logger.error(f"Unexpected error invoking Copilot action: {str(e)}")
361 + raise HTTPException(status_code=500, detail=f"Error invoking Copilot action: {str(e)}")
backend/app/integrations/copilot_action/schema/copilot_action.py
+12 -6
@@ -8,7 +8,6 @@ from typing import Union
8
9 from pydantic import BaseModel
10 from pydantic import Field
11 -from pydantic import HttpUrl
11 from pydantic import validator
12
13
@@ -33,6 +32,7 @@ class ScriptParameter(BaseModel):
32 description: Optional[str] = None
33 default: Optional[Union[str, int, float, bool, list, dict]] = None
34 enum: Optional[List[str]] = None
35 + arg_position: Optional[str] = None
36
37 @validator("type")
38 def validate_type(cls, v):
@@ -50,7 +50,7 @@ class ActiveResponseItem(BaseModel):
50 technology: Technology
51 icon: Optional[str] = None
52 script_parameters: List[ScriptParameter] = Field(default_factory=list)
53 - repo_url: HttpUrl
53 + repo_url: str # Changed from HttpUrl to str
54 script_name: Optional[str] = None
55 version: Optional[str] = None
56 last_updated: Optional[datetime] = None
@@ -65,10 +65,16 @@ class ActiveResponseItem(BaseModel):
65
66 @validator("repo_url")
67 def ensure_repo_url_ends_with_main(cls, v):
68 + # Keep it as string, just ensure it ends with /main/
69 repo_str = str(v)
69 - if not repo_str.endswith("/main"):
70 - return HttpUrl(f"{repo_str}/main")
71 - return v
70 + if not repo_str.endswith("/main/") and not repo_str.endswith("/main"):
71 + if repo_str.endswith("/"):
72 + return f"{repo_str}main/"
73 + else:
74 + return f"{repo_str}/main/"
75 + elif repo_str.endswith("/main"):
76 + return f"{repo_str}/"
77 + return repo_str
78
79
80 class InventoryQueryRequest(BaseModel):
@@ -95,7 +101,7 @@ class InventoryResponse(BaseModel):
101 class ActionDetailResponse(BaseModel):
102 """Response model for single action details"""
103
98 - active_response: ActiveResponseItem
104 + copilot_action: ActiveResponseItem
105 message: str
106 success: bool
107
backend/app/integrations/copilot_action/services/copilot_action.py
+6 -4
@@ -1,6 +1,7 @@
1 from typing import Optional
2
3 import httpx
4 +from fastapi import HTTPException
5 from loguru import logger
6
7 from app.integrations.copilot_action.schema.copilot_action import ActionDetailResponse
@@ -119,9 +120,10 @@ class CopilotActionService:
120
121 try:
122 data = response.json()
123 + logger.debug(f"Action detail response data: {data}")
124 except ValueError:
125 logger.error(f"Non-JSON response from action API: {response.text[:200]}")
124 - return ActionDetailResponse(active_response=None, message="Invalid response format from action API", success=False)
126 + return ActionDetailResponse(copilot_action=None, message="Invalid response format from action API", success=False)
127
128 logger.info(f"Successfully fetched action details for: {copilot_action_name}")
129 return ActionDetailResponse(**data)
@@ -129,12 +131,12 @@ class CopilotActionService:
131 except httpx.HTTPStatusError as e:
132 if e.response.status_code == 404:
133 logger.warning(f"Action not found: {copilot_action_name}")
132 - return ActionDetailResponse(active_response=None, message=f"Action '{copilot_action_name}' not found", success=False)
134 + raise HTTPException(status_code=404, detail=f"Action '{copilot_action_name}' not found")
135 logger.error(f"HTTP error fetching action details: {str(e)}")
134 - return ActionDetailResponse(active_response=None, message=f"HTTP error fetching action details: {str(e)}", success=False)
136 + return ActionDetailResponse(copilot_action=None, message=f"HTTP error fetching action details: {str(e)}", success=False)
137 except Exception as e:
138 logger.error(f"Unexpected error fetching action details: {str(e)}")
137 - return ActionDetailResponse(active_response=None, message=f"Unexpected error: {str(e)}", success=False)
139 + return ActionDetailResponse(copilot_action=None, message=f"Unexpected error: {str(e)}", success=False)
140
141 @classmethod
142 async def get_metrics(cls, license_key: str) -> InventoryMetricsResponse: