Create artifacts.py
taylor_socfortress committed
Jul 10, 2023 at 16:44 UTC
981c11996a4c3515ef15b23120133ffec7024c75
1 file changed
+121
backend/app/services/Velociraptor/artifacts.py
new
+121
@@ -0,0 +1,121 @@
1
+from typing import Dict
2
+from loguru import logger
3
+from pyvelociraptor import api_pb2
4
+from werkzeug.utils import secure_filename
5
+from app.services.Velociraptor.universal import UniversalService
6
+import json
7
+
8
+
9
+class ArtifactsService:
10
+ """
11
+ A service class that encapsulates the logic for pulling artifacts from Velociraptor.
12
+ """
13
+
14
+ def __init__(self):
15
+ self.universal_service = UniversalService()
16
+
17
+ def _create_query(self, query: str):
18
+ """
19
+ Create a query string.
20
+
21
+ Args:
22
+ query (str): The query to be executed.
23
+
24
+ Returns:
25
+ str: The created query string.
26
+ """
27
+ return query
28
+
29
+ def _get_artifact_key(self, client_id: str, artifact: str):
30
+ """
31
+ Construct the artifact key.
32
+
33
+ Args:
34
+ client_id (str): The ID of the client.
35
+ artifact (str): The name of the artifact.
36
+
37
+ Returns:
38
+ str: The constructed artifact key.
39
+ """
40
+ return f"collect_client(client_id='{client_id}', artifacts=['{artifact}'])"
41
+
42
+ def collect_artifacts(self):
43
+ """
44
+ Collect the artifacts from Velociraptor.
45
+
46
+ Returns:
47
+ dict: A dictionary with the success status, a message, and potentially the artifacts.
48
+ """
49
+ query = self._create_query("SELECT name FROM artifact_definitions()")
50
+ return self.universal_service.execute_query(query)
51
+
52
+ def collect_artifacts_prefixed(self, prefix: str):
53
+ """
54
+ Collect the artifacts from Velociraptor that have a name beginning with a specific prefix.
55
+
56
+ Args:
57
+ prefix (str): The prefix to filter the artifacts.
58
+
59
+ Returns:
60
+ dict: A dictionary with the success status, a message, and potentially the artifacts.
61
+ """
62
+ artifacts_response = self.collect_artifacts()
63
+ if not artifacts_response["success"]:
64
+ return artifacts_response
65
+
66
+ filtered_artifacts = [
67
+ artifact
68
+ for artifact in artifacts_response["results"]
69
+ if artifact["name"].startswith(prefix)
70
+ ]
71
+
72
+ return {
73
+ "success": True,
74
+ "message": f"Successfully collected {prefix} artifacts",
75
+ "artifacts": filtered_artifacts,
76
+ }
77
+
78
+ def collect_artifacts_linux(self):
79
+ return self.collect_artifacts_prefixed("Linux.")
80
+
81
+ def collect_artifacts_windows(self):
82
+ return self.collect_artifacts_prefixed("Windows.")
83
+
84
+ def collect_artifacts_macos(self):
85
+ return self.collect_artifacts_prefixed("MacOS.")
86
+
87
+ def run_artifact_collection(self, client_id: str, artifact: str):
88
+ """
89
+ Run an artifact collection on a specific client.
90
+
91
+ Args:
92
+ client_id (str): The ID of the client.
93
+ artifact (str): The name of the artifact.
94
+
95
+ Returns:
96
+ dict: A dictionary with the success status, a message, and potentially the results.
97
+ """
98
+ try:
99
+ query = self._create_query(
100
+ f"SELECT collect_client(client_id='{client_id}', artifacts=['{artifact}']) FROM scope()"
101
+ )
102
+ flow = self.universal_service.execute_query(query)
103
+ logger.info(f"Successfully ran artifact collection on {flow}")
104
+
105
+ artifact_key = self._get_artifact_key(client_id, artifact)
106
+ flow_id = flow["results"][0][artifact_key]["flow_id"]
107
+ logger.info(f"Successfully ran artifact collection on {flow_id}")
108
+
109
+ completed = self.universal_service.watch_flow_completion(flow_id)
110
+ logger.info(f"Successfully watched flow completion on {completed}")
111
+
112
+ results = self.universal_service.read_collection_results(
113
+ client_id, flow_id, artifact
114
+ )
115
+ return results
116
+ except Exception as err:
117
+ logger.error(f"Failed to run artifact collection: {err}")
118
+ return {
119
+ "message": "Failed to run artifact collection",
120
+ "success": False,
121
+ }