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: