Velociraptor: Add upload functionality for flow results

responders/Velociraptor/README.md

@@ -6,4 +6,7 @@ The following options are required in the Velociraptor Responder configuration:
 
 - `velociraptor_client_config`: The path to the Velociraptor API client config.  
 (See the following for generating an API client config: https://www.velocidex.com/docs/user-interface/api/, and ensure the appropriate ACLs are granted to the API user).  
-- `velociraptor_artifact`: The name artifact you which to collect (as you would see it in the Velociraptor GUI)
+- `velociraptor_artifact`: The name artifact you which to collect (as you would see it in the Velociraptor GUI).
+- `upload_flow_results`: Upload flow results to TheHive case (bool).
+- `thehive_url`: URL of your TheHive installation (e.g. 'http://127.0.0.1:9000').
+- `thehive_apikey`: TheHive API key used to add flow results/file(s) to a case.

responders/Velociraptor/velociraptor_flow.json

@@ -24,7 +24,27 @@
       "multi": false,
       "required": true,
       "defaultValue": ""
+    },
+    {
+      "name": "upload_flow_results",
+      "description": "Upload the results of a flow as an observable",
+      "type": "boolean",
+      "multi": false,
+      "required": true
+    },
+    {
+      "name": "thehive_url",
+      "description": "URL pointing to your TheHive installation, e.g. 'http://127.0.0.1:9000'",
+      "type": "string",
+      "multi": false,
+      "required": true
+    },
+    {
+      "name": "thehive_apikey",
+      "description": "TheHive API key (used to add the downloaded file back to the alert/case)",
+      "type": "string",
+      "multi": false,
+      "required": true
     }
   ]
 }
-

responders/Velociraptor/velociraptor_flow.py

@@ -3,23 +3,31 @@
 import json
 import grpc
 import re
+import os
 import time
 import yaml
+from thehive4py.api import TheHiveApi
+from thehive4py.models import Case, CaseObservable
 import pyvelociraptor
 from pyvelociraptor import api_pb2
 from pyvelociraptor import api_pb2_grpc
 
+
 class Velociraptor(Responder):
   def __init__(self):
     Responder.__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.upload_flow_results = self.get_param('config.upload_flow_results', None, 'Upload decision missing!')
     self.observable_type = self.get_param('data.dataType', None, "Data type is empty")
     self.observable = self.get_param('data.data', None, 'Data missing!')
-
+    self.thehive_url = self.get_param('config.thehive_url', None, "TheHive URL missing!")
+    self.thehive_apikey = self.get_param('config.thehive_apikey', None, "TheHive API key missing!")
+
   def run(self):
     Responder.run(self)
+    case_id = self.get_param('data._parent')
     creds = grpc.ssl_channel_credentials(
         root_certificates=self.config["ca_certificate"].encode("utf8"),
         private_key=self.config["client_private_key"].encode("utf8"),
@@ -72,18 +80,18 @@ def run(self):
           try:
             init_results = json.loads(response.Response)
             flow=list(init_results[0].values())[0]
-            self.report({'message': init_results })
-
+
+            flow_id = str(flow['flow_id']) 
             # Define second query
-            flow_query = "SELECT * from flows(client_id='" + str(flow['request']['client_id']) + "', flow_id='" + str(flow['flow_id']) + "')"
+            flow_query = "SELECT * from flows(client_id='" + str(flow['request']['client_id']) + "', flow_id='" + flow_id + "')"
 
             state=0
 
             # Check to see if the flow has completed        
-            while (state == 0):
+            while (state != 2):
 
               followup_request = api_pb2.VQLCollectorArgs(
-                max_wait=1,
+                max_wait=10,
                 Query=[api_pb2.VQLRequest(
                      Name="TheHive-QueryForFlow",
                      VQL=flow_query,
@@ -95,29 +103,84 @@ def run(self):
                 except:
                   pass
               state = flow_results[0]['state']
-              if state == 1:
+              global artifact_results
+              artifact_results = flow_results[0]['artifacts_with_results']
+              self.report({'message': state })
+              if state == 2:
                 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
+            for artifact in artifact_results:
+              source_query="SELECT * from source(client_id='"+ str(flow['request']['client_id']) + "', flow_id='" + flow_id +  "', artifact='" + artifact + "')"
+              source_request = api_pb2.VQLCollectorArgs(
+              max_wait=10,
+              Query=[api_pb2.VQLRequest(
+                  Name="TheHive-SourceQuery",
+                  VQL=source_query,
+              )])
+              for source_response in stub.Query(source_request):
+                try:
+                  source_result = json.loads(source_response.Response)
+                  source_results += source_result
+                except:
+                  pass
+            self.report({'message': source_results })
+
+            if self.upload_flow_results is True: 
+              # Create flow download
+              vfs_query = "SELECT create_flow_download(client_id='"+ str(flow['request']['client_id']) + "', flow_id='" + str(flow['flow_id']) +  "', wait='true') as VFSPath from scope()" 
+              vfs_request = api_pb2.VQLCollectorArgs(
+                max_wait=10,
+                Query=[api_pb2.VQLRequest(
+                    Name="TheHive-VFSQuery",
+                    VQL=vfs_query,
+              )])
+              for vfs_response in stub.Query(vfs_request):
+                  try:
+                    vfs_result = json.loads(vfs_response.Response)[0]['VFSPath']
+                  except:
+                    pass
+              # "Artifact" plugin.
+              offset = 0
+
+              file_request = api_pb2.VFSFileBuffer(
+                  vfs_path=vfs_result,
+                  length=10000,
+                  offset=offset,
+              )
+
+              res = stub.VFSGetBuffer(file_request)
+
+              if len(res.data) == 0:
+                 break
+
+
+
+              f = open("/tmp/" + self.artifact + "_" + client_id + "_" + flow_id + "_" + case_id + ".zip",'wb')
+              f.write(res.data)
+              offset+=len(res.data)
+
+              #Upload file to TheHive
+              api = TheHiveApi(self.thehive_url, self.thehive_apikey, cert=False)
+
+              description = "Velociraptor flow for artifact" + self.artifact + "for client " + client_id + " via flow " + flow_id + "." 
+              filepath = '/tmp/' + self.artifact + "_" + client_id + "_" + flow_id + "_" + case_id + ".zip"
+              file_observable = CaseObservable(dataType='file',
+                   data=[filepath],
+                   tlp=self.get_param('data.tlp'),
+                   ioc=True,
+                   tags=['src:Velociraptor', client_id ],
+                   message=description
+              )
+
+              response = api.create_case_observable(case_id, file_observable)
+              if response.status_code != 201:
+                self.error({'message': str(response.status_code) + " " + response.text})
+              os.remove(filepath)
           except:
             pass
-
   def operations(self, raw):
       global client_id
       return [self.build_operation('AddTagToArtifact', tag=client_id)]