Update Velociraptor Responder

responders/Velociraptor/velociraptor_flow.json

@@ -1,10 +1,10 @@
 {
-  "name": "Velociraptor_Flow",
-  "version": "0.1",
-  "author": "Wes Lambert",
+  "name": "Velociraptor",
+  "version": "0.2",
+  "author": "Wes Lambert, @therealwlambert",
   "url": "https://github.com/TheHive-Project/Cortex-Analyzers",
   "license": "AGPL-V3",
-  "description": "Run Velociraptor flow",
+  "description": "Run Velociraptor artifact collection",
   "dataTypeList": ["thehive:case_artifact"],
   "command": "Velociraptor/velociraptor_flow.py",
   "baseConfig": "Velociraptor",
@@ -24,7 +24,14 @@
       "multi": false,
       "required": true,
       "defaultValue": ""
+    },
+    {
+      "name": "query_max_duration",
+      "description": "Max query duration (in seconds)",
+      "type": "number",
+      "multi": false,
+      "required": false,
+      "defaultValue": "600"
     }
   ]
 }
-

responders/Velociraptor/velociraptor_flow.py

@@ -15,9 +15,12 @@ def __init__(self):
     self.configpath = self.get_param('config.velociraptor_client_config', None, "File path missing!")
     self.config = yaml.load(open(self.configpath).read(), Loader=yaml.FullLoader)
     self.artifact = self.get_param('config.velociraptor_artifact', None, 'Artifact missing!')
+    self.artifact_args = self.get_param('config.velociraptor_artifact_args', None)
     self.observable_type = self.get_param('data.dataType', None, "Data type is empty")
     self.observable = self.get_param('data.data', None, 'Data missing!')
-
+    self.max_wait = self.get_param('config.query_max_duration', 600)
+    self.query = self.get_param('config.velociraptor_query', None)
+
   def run(self):
     Responder.run(self)
     creds = grpc.ssl_channel_credentials(
@@ -40,9 +43,10 @@ def run(self):
             self.report({'message': "Not a valid data type!" })
             return
 
-        # Send initial request
+        # Query to get client ID
         client_request = api_pb2.VQLCollectorArgs(
-            max_wait=1,
+            # Setting static max_wait here, because it should not take as long as other queries
+            max_wait=60,
             Query=[api_pb2.VQLRequest(
                 Name="TheHive-ClientQuery",
                 VQL=client_query,
@@ -56,65 +60,27 @@ def run(self):
           except:
               self.report({'message': 'Could not find a suitable client.'})
               pass
+
+        # Free-form query
+        freeform_query = self.query  
+
+        # Artifact query
+        artifact_query = "LET collection <= collect_client(client_id='"+ client_id +"',artifacts=['" + self.artifact + "'], spec=dict()) LET collection_completed <= SELECT * FROM watch_monitoring(artifact='System.Flow.Completion') WHERE FlowId = collection.flow_id  LIMIT 1 SELECT * FROM source(client_id=collection.request.client_id, flow_id=collection.flow_id, artifact=collection_completed.Flow.artifacts_with_results[0])"
 
-        # Define initial query
-        init_query = "SELECT collect_client(client_id='"+ client_id +"',artifacts=['" + self.artifact + "']) FROM scope()"
+
 
-        # Send initial request
         request = api_pb2.VQLCollectorArgs(
-            max_wait=1,
+            max_wait=self.max_wait,
             Query=[api_pb2.VQLRequest(
                 Name="TheHive-Query",
-                VQL=init_query,
+                VQL=artifact_query,
             )])
 
         for response in stub.Query(request):
           try:
-            init_results = json.loads(response.Response)
-            flow=list(init_results[0].values())[0]
-            self.report({'message': init_results })
-
-            # Define second query
-            flow_query = "SELECT * from flows(client_id='" + str(flow['request']['client_id']) + "', flow_id='" + str(flow['flow_id']) + "')"
-
-            state=0
-
-            # Check to see if the flow has completed        
-            while (state == 0):
-
-              followup_request = api_pb2.VQLCollectorArgs(
-                max_wait=1,
-                Query=[api_pb2.VQLRequest(
-                     Name="TheHive-QueryForFlow",
-                     VQL=flow_query,
-                )])
-
-              for followup_response in stub.Query(followup_request):
-                try:
-                    flow_results = json.loads(followup_response.Response)
-                except:
-                  pass
-              state = flow_results[0]['state']
-              if state == 1:
-                time.sleep(5)
-                break
-
-            # Grab the source from the artifact
-            source_query="SELECT * from source(client_id='"+ str(flow['request']['client_id']) + "', flow_id='" + str(flow['flow_id']) +  "', artifact='" + self.artifact + "')"
-            source_request = api_pb2.VQLCollectorArgs(
-            max_wait=1,
-            Query=[api_pb2.VQLRequest(
-                Name="TheHive-SourceQuery",
-                VQL=source_query,
-            )])
-            source_results=[]
-            for source_response in stub.Query(source_request):
-              try:
-                source_result = json.loads(source_response.Response)
-                source_results += source_result
-                self.report({'message': source_results })
-              except:
-                pass
+            query_results = json.loads(response.Response)
+            #flow=list(init_results[0].values())[0]
+            self.report({'message': query_results })
           except:
             pass