Create universal.py
taylor_socfortress committed
Jul 10, 2023 at 16:44 UTC
b818d646d83e88963b0b98d30a04cf7b9c3aa8e5
1 file changed
+209
backend/app/services/Velociraptor/universal.py
new
+209
@@ -0,0 +1,209 @@
1
+from app.models.agents import (
2
+ AgentMetadata,
3
+ agent_metadata_schema,
4
+ agent_metadatas_schema,
5
+)
6
+from typing import Dict, List
7
+from app import db
8
+from datetime import datetime
9
+import requests
10
+from loguru import logger
11
+from elasticsearch7 import Elasticsearch
12
+from app.models.connectors import connector_factory, Connector
13
+import pyvelociraptor
14
+from pyvelociraptor import api_pb2
15
+from pyvelociraptor import api_pb2_grpc
16
+import grpc
17
+import json
18
+
19
+
20
+class UniversalService:
21
+ """
22
+ A service class that encapsulates the logic for polling messages from Velociraptor.
23
+ """
24
+
25
+ def __init__(self) -> None:
26
+ self.setup_velociraptor_connector("Velociraptor")
27
+ self.setup_grpc_channel_and_stub()
28
+
29
+ def setup_velociraptor_connector(self, connector_name: str):
30
+ """
31
+ Collects the details of the Velociraptor connector and sets them up.
32
+
33
+ Args:
34
+ connector_name (str): The name of the Velociraptor connector.
35
+ """
36
+ self.connector_url, self.connector_api_key = self.collect_velociraptor_details(connector_name)
37
+ self.config = pyvelociraptor.LoadConfigFile(self.connector_api_key)
38
+
39
+ def collect_velociraptor_details(self, connector_name: str):
40
+ """
41
+ Collects the details of the Velociraptor connector.
42
+
43
+ Args:
44
+ connector_name (str): The name of the Velociraptor connector.
45
+
46
+ Returns:
47
+ tuple: A tuple containing the connection URL, and api key.
48
+ """
49
+ connector_instance = connector_factory.create(connector_name, connector_name)
50
+ connection_successful = connector_instance.verify_connection()
51
+ if connection_successful:
52
+ connection_details = Connector.get_connector_info_from_db(connector_name)
53
+ return (
54
+ connection_details.get("connector_url"),
55
+ connection_details.get("connector_api_key"),
56
+ )
57
+ else:
58
+ return None, None
59
+
60
+ def setup_grpc_channel_and_stub(self):
61
+ """
62
+ Sets up the gRPC channel and stub for Velociraptor.
63
+ """
64
+ creds = grpc.ssl_channel_credentials(
65
+ root_certificates=self.config["ca_certificate"].encode("utf8"),
66
+ private_key=self.config["client_private_key"].encode("utf8"),
67
+ certificate_chain=self.config["client_cert"].encode("utf8"),
68
+ )
69
+ options = (("grpc.ssl_target_name_override", "VelociraptorServer"),)
70
+ self.channel = grpc.secure_channel(self.config["api_connection_string"], creds, options)
71
+ self.stub = api_pb2_grpc.APIStub(self.channel)
72
+
73
+ def create_vql_request(self, vql: str):
74
+ """
75
+ Creates a VQLCollectorArgs object with given VQL query.
76
+
77
+ Args:
78
+ vql (str): The VQL query.
79
+
80
+ Returns:
81
+ VQLCollectorArgs: The VQLCollectorArgs object with given VQL query.
82
+ """
83
+ return api_pb2.VQLCollectorArgs(
84
+ max_wait=1,
85
+ Query=[
86
+ api_pb2.VQLRequest(
87
+ Name="VQLRequest",
88
+ VQL=vql,
89
+ ),
90
+ ],
91
+ )
92
+
93
+ def execute_query(self, vql: str):
94
+ """
95
+ Executes a VQL query and returns the results.
96
+
97
+ Args:
98
+ vql (str): The VQL query to be executed.
99
+
100
+ Returns:
101
+ dict: A dictionary with the success status, a message, and potentially the results.
102
+ """
103
+ client_request = self.create_vql_request(vql)
104
+ try:
105
+ results = []
106
+ for response in self.stub.Query(client_request):
107
+ if response.Response:
108
+ results += json.loads(response.Response)
109
+ return {
110
+ "success": True,
111
+ "message": "Successfully executed query",
112
+ "results": results,
113
+ }
114
+ except Exception as e:
115
+ return {
116
+ "success": False,
117
+ "message": f"Failed to execute query: {e}",
118
+ }
119
+
120
+
121
+ def watch_flow_completion(self, flow_id: str):
122
+ """
123
+ Watch for the completion of a flow.
124
+
125
+ Args:
126
+ flow_id (str): The ID of the flow.
127
+
128
+ Returns:
129
+ dict: A dictionary with the success status and a message.
130
+ """
131
+ vql = f"SELECT * FROM watch_monitoring(artifact='System.Flow.Completion') WHERE FlowId='{flow_id}' LIMIT 1"
132
+ return self.execute_query(vql)
133
+
134
+ def read_collection_results(self, client_id: str, flow_id: str, artifact: str = 'Generic.Client.Info/BasicInformation'):
135
+ """
136
+ Read the results of a collection.
137
+
138
+ Args:
139
+ client_id (str): The client ID.
140
+ flow_id (str): The ID of the flow.
141
+ artifact (str, optional): The artifact. Defaults to 'Generic.Client.Info/BasicInformation'.
142
+
143
+ Returns:
144
+ dict: A dictionary with the success status, a message, and potentially the results.
145
+ """
146
+ vql = f"SELECT * FROM source(client_id='{client_id}', flow_id='{flow_id}', artifact='{artifact}')"
147
+ return self.execute_query(vql)
148
+
149
+ def get_client_id(self, client_name: str):
150
+ """
151
+ Get the client_id associated with a given client_name.
152
+
153
+ Args:
154
+ client_name (str): The asset name to search for.
155
+
156
+ Returns:
157
+ dict: A dictionary with the success status, a message, and potentially the client_id.
158
+ """
159
+ # Formulate queries
160
+ try:
161
+ vql_client_id = f"select client_id from clients(search='host:{client_name}')"
162
+ vql_last_seen_at = f"select last_seen_at from clients(search='host:{client_name}')"
163
+
164
+ # Get the last seen timestamp
165
+ last_seen_at = self._get_last_seen_timestamp(vql_last_seen_at)
166
+
167
+ # if last_seen_at is longer than 30 seconds from now, return False
168
+ if self._is_offline(last_seen_at):
169
+ return {
170
+ "success": False,
171
+ "message": f"{client_name} has not been seen in the last 30 seconds and may not be online with the Velociraptor server.",
172
+ "results": [
173
+ {"client_id": None}
174
+ ]
175
+ }
176
+
177
+ return self.execute_query(vql_client_id)
178
+ except Exception as e:
179
+ return {
180
+ "success": False,
181
+ "message": f"Failed to get Client ID for {client_name}: {e}",
182
+ "results": [
183
+ {"client_id": None}
184
+ ]
185
+ }
186
+
187
+ def _get_last_seen_timestamp(self, vql: str):
188
+ """
189
+ Executes the VQL query and returns the last_seen_at timestamp.
190
+
191
+ Args:
192
+ vql (str): The VQL query.
193
+
194
+ Returns:
195
+ float: The last_seen_at timestamp.
196
+ """
197
+ return self.execute_query(vql)["results"][0]["last_seen_at"]
198
+
199
+ def _is_offline(self, last_seen_at: float):
200
+ """
201
+ Determines if the client is offline based on the last_seen_at timestamp.
202
+
203
+ Args:
204
+ last_seen_at (float): The last_seen_at timestamp.
205
+
206
+ Returns:
207
+ bool: True if the client is offline, False otherwise.
208
+ """
209
+ return (datetime.now() - datetime.fromtimestamp(last_seen_at / 1000000)).total_seconds() > 30