logic to run remote commands via velociraptor (#42)

backend/app/routes/velociraptor.py

@@ -118,3 +118,39 @@ def collect_artifact():
         artifact=artifact_name,
     )
     return artifact_results
+
+
+@bp.route("/velociraptor/remotecommand", methods=["POST"])
+def run_remote_command():
+    """
+    Endpoint to run a remote command.
+    It collects the command and client name from the request body and returns the results.
+
+    Returns:
+        json: A JSON response containing the result of the PowerShell command execution.
+    """
+    req_data = request.get_json()
+    command = req_data["command"]
+    client_name = req_data["client_name"]
+    artifact_name = req_data["artifact_name"]
+    service = UniversalService()
+    client_id = service.get_client_id(client_name=client_name)["results"][0]["client_id"]
+    if client_id is None:
+        return (
+            jsonify(
+                {
+                    "message": f"{client_name} has not been seen in the last 30 seconds and may not be online with the "
+                    "Velociraptor server.",
+                    "success": False,
+                },
+            ),
+            500,
+        )
+
+    artifact_service = ArtifactsService()
+    artifact_results = artifact_service.run_remote_command(
+        client_id=client_id,
+        artifact=artifact_name,
+        command=command,
+    )
+    return artifact_results

backend/app/services/Velociraptor/artifacts.py

@@ -24,18 +24,22 @@ def _create_query(self, query: str) -> str:
         """
         return query
 
-    def _get_artifact_key(self, client_id: str, artifact: str) -> str:
+    def _get_artifact_key(self, client_id: str, artifact: str, command: str = None) -> str:
         """
         Construct the artifact key.
 
         Args:
             client_id (str): The ID of the client.
             artifact (str): The name of the artifact.
+            command (str): The command that was run, if applicable.
 
         Returns:
             str: The constructed artifact key.
         """
-        return f"collect_client(client_id='{client_id}', artifacts=['{artifact}'])"
+        if command:
+            return f"collect_client(client_id='{client_id}', urgent=true, artifacts=['{artifact}'], env=dict(Command='{command}'))"
+        else:
+            return f"collect_client(client_id='{client_id}', artifacts=['{artifact}'])"
 
     def collect_artifacts(self) -> dict:
         """
@@ -163,3 +167,44 @@ def run_artifact_collection(self, client_id: str, artifact: str) -> dict:
                 "message": "Failed to run artifact collection",
                 "success": False,
             }
+
+    def run_remote_command(self, client_id: str, artifact: str, command: str) -> dict:
+        """
+        Run a remote command on a specific client.
+        Accepted artifact names are `Windows.System.PowerShell`, `Windows.System.CmdShell`.
+
+        Args:
+            client_id (str): The ID of the client.
+            artifact (str): The name of the artifact.
+            command (str): The command to be executed.
+
+        Returns:
+            dict: A dictionary with the success status, a message, and potentially the results.
+        """
+        try:
+            query = self._create_query(
+                f"SELECT collect_client(client_id='{client_id}', urgent=true, artifacts=['{artifact}'], env=dict(Command='{command}')) "
+                "FROM scope()",
+            )
+            flow = self.universal_service.execute_query(query)
+            logger.info(f"Successfully ran artifact collection on {flow}")
+
+            artifact_key = self._get_artifact_key(client_id, artifact, command)
+            flow_id = flow["results"][0][artifact_key]["flow_id"]
+            logger.info(f"Successfully ran artifact collection on {flow_id}")
+
+            completed = self.universal_service.watch_flow_completion(flow_id)
+            logger.info(f"Successfully watched flow completion on {completed}")
+
+            results = self.universal_service.read_collection_results(
+                client_id,
+                flow_id,
+                artifact,
+            )
+            return results
+        except Exception as err:
+            logger.error(f"Failed to run artifact collection: {err}")
+            return {
+                "message": "Failed to run artifact collection",
+                "success": False,
+            }

backend/app/static/swagger.json

@@ -1734,9 +1734,9 @@
                 "description": "Endpoint to get all artifacts for a hostname.",
                 "parameters": [
                     {
-                        "name": "hostname",
+                        "name": "client_name",
                         "in": "path",
-                        "description": "The hostname",
+                        "description": "The client name",
                         "required": true,
                         "schema": {
                             "type": "string"
@@ -1834,6 +1834,69 @@
                 "tags": ["Velociraptor"]
             }
         },
+        "/velociraptor/remotecommand": {
+            "post": {
+                "summary": "Run a remote command",
+                "description": "Endpoint to run a remote command.",
+                "requestBody": {
+                    "description": "Remote command details",
+                    "content": {
+                        "application/json": {
+                            "schema": {
+                                "type": "object",
+                                "properties": {
+                                    "client_name": {
+                                        "type": "string",
+                                        "value": "WIN-39O01J5F7G5",
+                                        "description": "The hostname of the client to run the command on."
+                                    },
+                                    "artifact_name": {
+                                        "type": "string",
+                                        "value": "Windows.System.PowerShell",
+                                        "description": "The name of the artifact to run."
+                                    },
+                                    "command": {
+                                        "type": "string",
+                                        "value": "ping 8.8.8.8",
+                                        "description": "The powershell command to run."
+                                    }
+                                }
+                            }
+                        }
+                    }
+                },
+                "responses": {
+                    "200": {
+                        "description": "Successful operation",
+                        "content": {
+                            "application/json": {
+                                "schema": {
+                                    "type": "object",
+                                    "properties": {
+                                        "output": {
+                                            "type": "string",
+                                            "description": "The output of the command."
+                                        }
+                                    }
+                                }
+                            }
+                        }
+                    },
+                    "default": {
+                        "description": "Unexpected error",
+                        "content": {
+                            "application/json": {
+                                "schema": {
+                                    "$ref": "#/components/schemas/Error"
+                                }
+                            }
+                        }
+                    }
+                },
+                "operationId": "runPowershellCommand",
+                "tags": ["Velociraptor"]
+            }
+        },
         "/dfir_iris/cases": {
             "get": {
                 "summary": "Get all cases",