velociraptor: get_client_flow_results - optional artifact, no IndexError on multi-source artifacts (fixes #458)

velociraptor/1.0.0/api.yaml

@@ -109,6 +109,13 @@ actions:
         required: true
         schema:
           type: string
+      - name: artifact
+        description: Artifact, or Artifact/Source, to read (needed for artifacts with several sources, e.g. Generic.Client.Info/BasicInformation). Leave empty for the previous behaviour.
+        multiline: false
+        example: Generic.Client.Info/BasicInformation
+        required: false
+        schema:
+          type: string
     returns:
       schema:
         type: string

velociraptor/1.0.0/src/app.py

@@ -1,4 +1,5 @@
 import json
+import re
 import grpc
 import ipaddress
 import time
@@ -123,17 +124,27 @@ def get_client_flows(self, api_config, client_id):
         results = self.request(api_config, query)
         return results
 
-    def get_client_flow_results(self, api_config, client_id, flow_id):
+    def get_client_flow_results(self, api_config, client_id, flow_id, artifact=""):
+        if artifact:
+            # Only an artifact name, optionally /Source - it becomes part of the VQL below.
+            if not re.fullmatch(r"[A-Za-z0-9_.]+(/[A-Za-z0-9_.]+)?", artifact):
+                raise ValueError("Velociraptor: invalid artifact name %r" % artifact)
         state = self.get_client_flow_status(api_config, client_id, flow_id)
         while (state == "RUNNING"):
             state = self.get_client_flow_status(api_config, client_id, flow_id)
             if state == "FINISHED":
                 break
             else:
                 time.sleep(5)
-        query = "SELECT * FROM flow_results(flow_id='" + flow_id  + "', client_id='" + client_id  + "')"
+        query = "SELECT * FROM flow_results(flow_id='" + flow_id  + "', client_id='" + client_id  + "'"
+        if artifact:
+            # Artifacts with several sources return no rows unless the source is named.
+            query += ", artifact='" + artifact + "'"
+        query += ")"
         results = self.request(api_config, query)
-        return results[0]
+        if artifact:
+            return results
+        return results[0] if results else []
 
     def get_client_flow_status(self, api_config, client_id, flow_id):
         query = "SELECT * FROM flows(flow_id='" + flow_id  + "', client_id='" + client_id  + "')"