Copilot action agent list (#500)
* Refactor invoke action to support multiple agents and update request body to accept a list of agent names * precommit-fixes
taylor_socfortress committed
Sep 5, 2025 at 12:00 UTC
bd9fc583b1a101a928746fa3a0459a8d7fafbfda
2 files changed
+110
-53
backend/app/connectors/velociraptor/schema/artifacts.py
+1
-1
@@ -155,7 +155,7 @@ 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")
158
+ agent_names: List[str] = Field(..., description="List of agent names to invoke the action on") # Changed from agent_name to agent_names
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,
backend/app/integrations/copilot_action/routes/copilot_action.py
+109
-52
@@ -1,4 +1,5 @@
1
import os
2
+from typing import List
3
from typing import Optional
4
5
from fastapi import APIRouter
@@ -40,21 +41,22 @@ def get_license_key() -> str:
41
return license_key
42
43
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()
44
+async def get_agents_by_hostnames(session: AsyncSession, hostnames: List[str]) -> List[Agents]:
45
+ """Retrieve multiple agents from database by hostnames."""
46
+ agent_details = await session.execute(select(Agents).filter(Agents.hostname.in_(hostnames)))
47
+ agents = agent_details.scalars().all()
48
+
49
+ found_hostnames = {agent.hostname for agent in agents}
50
+ missing_hostnames = set(hostnames) - found_hostnames
51
50
- if not agent:
52
+ if missing_hostnames:
53
raise HTTPException(
54
status_code=404,
53
- detail=f"Agent with hostname {hostname} not found",
55
+ detail=f"Agents with hostnames {list(missing_hostnames)} not found",
56
)
57
56
- logger.info(f"Found agent: {agent.hostname} with OS: {agent.os}")
57
- return agent
58
+ logger.info(f"Found {len(agents)} agents")
59
+ return agents
60
61
62
def determine_artifact_name(agent_os: str) -> str:
@@ -288,70 +290,125 @@ async def get_technologies() -> dict:
290
}
291
292
293
+async def invoke_action_on_agent(
294
+ agent: Agents,
295
+ copilot_action_name: str,
296
+ copilot_action_details: ActionDetailResponse,
297
+ parameters: dict,
298
+) -> CollectArtifactResponse:
299
+ """
300
+ Invoke a Copilot Action on a single agent.
301
+
302
+ Args:
303
+ agent: The agent to invoke the action on
304
+ copilot_action_name: Name of the action
305
+ copilot_action_details: Action details from the service
306
+ parameters: Parameters for the action
307
+
308
+ Returns:
309
+ CollectArtifactResponse: Response from the artifact collection
310
+ """
311
+ try:
312
+ # Determine the appropriate artifact based on OS
313
+ artifact_name = determine_artifact_name(agent.os)
314
+ logger.info(f"Using artifact: {artifact_name} for OS: {agent.os} on agent {agent.hostname}")
315
+
316
+ # Build Velociraptor parameters
317
+ velociraptor_params = build_velociraptor_parameters(parameters, copilot_action_details.copilot_action.script_parameters)
318
+
319
+ # Build artifact collection request
320
+ artifact_body = await build_artifact_collection_body(agent, artifact_name, velociraptor_params)
321
+ logger.info(f"Built artifact collection request for {agent.hostname}")
322
+
323
+ # Execute the collection
324
+ response = await run_artifact_collection(artifact_body)
325
+ logger.info(f"Successfully invoked Copilot action on {agent.hostname}")
326
+
327
+ return response
328
+
329
+ except Exception as e:
330
+ logger.error(f"Error invoking action on agent {agent.hostname}: {str(e)}")
331
+ # You might want to return a failed response instead of raising
332
+ raise HTTPException(status_code=500, detail=f"Error invoking action on agent {agent.hostname}: {str(e)}")
333
+
334
+
335
@copilot_action_router.post(
336
"/invoke",
293
- response_model=CollectArtifactResponse,
294
- description="Invoke a Copilot Action on a target agent",
337
+ response_model=List[CollectArtifactResponse], # Now returns a list of responses
338
+ description="Invoke a Copilot Action on multiple target agents",
339
dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],
340
)
297
-async def invoke_action(body: InvokeCopilotActionBody, session: AsyncSession = Depends(get_db)) -> CollectArtifactResponse:
341
+async def invoke_action(body: InvokeCopilotActionBody, session: AsyncSession = Depends(get_db)) -> List[CollectArtifactResponse]:
342
"""
299
- Invoke a Copilot Action on a target agent.
343
+ Invoke a Copilot Action on multiple target agents.
344
345
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
346
+ 1. Finding all target agents
347
+ 2. Fetching action details and validating parameters once
348
+ 3. Building and executing the artifact collection request for each agent
349
350
Args:
308
- body: Request body containing action name, agent name, and parameters
351
+ body: Request body containing action name, agent names, and parameters
352
session: Database session
353
354
Returns:
312
- CollectArtifactResponse: Response from the artifact collection
355
+ List[CollectArtifactResponse]: List of responses from the artifact collections
356
"""
314
- logger.info(f"Invoking Copilot action '{body.copilot_action_name}' on agent '{body.agent_name}'")
357
+ logger.info(f"Invoking Copilot action '{body.copilot_action_name}' on {len(body.agent_names)} agents")
358
359
try:
317
- # Step 1: Get the target agent
318
- agent = await get_agent_by_hostname(session, body.agent_name)
360
+ # Step 1: Get all target agents
361
+ agents = await get_agents_by_hostnames(session, body.agent_names)
362
+ logger.info(f"Found agents: {[agent.hostname for agent in agents]}")
363
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
364
+ # Step 2: Fetch action details (do this once for all agents)
365
copilot_action_details = await get_action_by_name(body.copilot_action_name)
366
logger.info(f"Found action details for: {body.copilot_action_name}")
367
368
+ # Step 3: Prepare parameters (do this once for all agents)
369
+ final_parameters = body.parameters or {}
370
+
371
# Add the `repo_url` and the `script_name` to the parameters
372
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
373
+ final_parameters["RepoURL"] = copilot_action_details.copilot_action.repo_url
374
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
375
+ final_parameters["ScriptName"] = copilot_action_details.copilot_action.script_name
376
+
377
+ logger.info(f"Parameters after adding repo and script: {final_parameters}")
378
+
379
+ # Step 4: Validate parameters (do this once for all agents)
380
+ await validate_parameters(final_parameters, copilot_action_details.copilot_action.script_parameters)
381
+
382
+ # Step 5: Execute action on each agent
383
+ responses = []
384
+ successful_agents = []
385
+ failed_agents = []
386
+
387
+ for agent in agents:
388
+ try:
389
+ response = await invoke_action_on_agent(agent, body.copilot_action_name, copilot_action_details, final_parameters)
390
+ responses.append(response)
391
+ successful_agents.append(agent.hostname)
392
+
393
+ except Exception as e:
394
+ logger.error(f"Failed to invoke action on agent {agent.hostname}: {str(e)}")
395
+ failed_agents.append(agent.hostname)
396
+ # Add a failed response to maintain order
397
+ failed_response = CollectArtifactResponse(
398
+ message=f"Failed to invoke action on {agent.hostname}: {str(e)}",
399
+ success=False,
400
+ results=[],
401
+ )
402
+ responses.append(failed_response)
403
+
404
+ # Log summary
405
+ logger.info(f"Action invocation complete. Successful: {len(successful_agents)}, Failed: {len(failed_agents)}")
406
+ if successful_agents:
407
+ logger.info(f"Successful agents: {successful_agents}")
408
+ if failed_agents:
409
+ logger.warning(f"Failed agents: {failed_agents}")
410
+
411
+ return responses
412
413
except HTTPException:
414
# Re-raise HTTP exceptions (validation errors, not found, etc.)