Create connectors.py
taylor_socfortress committed
Jul 10, 2023 at 16:39 UTC
6a9b1cb7a960ea2e278aadf1b1ee58cabc477063
1 file changed
+483
backend/app/models/connectors.py
new
+483
@@ -0,0 +1,483 @@
1
+import importlib
2
+import json
3
+import os
4
+import pika
5
+from dataclasses import dataclass
6
+import requests
7
+from abc import ABC, abstractmethod
8
+from elasticsearch7 import Elasticsearch
9
+from loguru import logger
10
+from sqlalchemy.orm.exc import NoResultFound
11
+import pyvelociraptor
12
+from pyvelociraptor import api_pb2
13
+from pyvelociraptor import api_pb2_grpc
14
+from werkzeug.utils import secure_filename
15
+import grpc
16
+
17
+from sqlalchemy.exc import SQLAlchemyError
18
+from flask import current_app
19
+
20
+from app.models.models import Connectors, connectors_schema, ConnectorsAvailable
21
+
22
+
23
+def dynamic_import(module_name, class_name):
24
+ """
25
+ This function dynamically imports a module and returns a specific class from it.
26
+
27
+ :param module_name: A string that specifies the name of the module to import.
28
+ :param class_name: A string that specifies the name of the class to get from the module.
29
+ :return: The class specified by class_name from the module specified by module_name.
30
+ """
31
+ module = importlib.import_module(module_name)
32
+ class_ = getattr(module, class_name)
33
+ return class_
34
+
35
+
36
+@dataclass
37
+class Connector(ABC):
38
+ """
39
+ This abstract base class defines the interface for a connector. A connector is an object that
40
+ connects to a specific service or system and performs actions on it. The specific service or
41
+ system a connector connects to is defined by the connector's attributes.
42
+
43
+ :param attributes: A dictionary of attributes necessary for the connector to connect to the service or system.
44
+ """
45
+
46
+ attributes: dict
47
+
48
+ @abstractmethod
49
+ def verify_connection(self):
50
+ """
51
+ This abstract method should be implemented by all subclasses of Connector. It is meant to verify the
52
+ connection to the service or system the connector is designed to connect to.
53
+
54
+ :return: Depends on the implementation in the subclass.
55
+ """
56
+ pass
57
+
58
+ @staticmethod
59
+ def get_connector_info_from_db(connector_name):
60
+ """
61
+ This method retrieves connector information from the database.
62
+
63
+ :param connector_name: A string that specifies the name of the connector whose information is to be retrieved.
64
+ :return: A dictionary of the connector's attributes if the connector exists. Otherwise, it raises a NoResultFound exception.
65
+ Raises:
66
+ NoResultFound: If the connector_name is not found in the database.
67
+ """
68
+ connector = (
69
+ current_app.extensions["sqlalchemy"]
70
+ .db.session.query(Connectors)
71
+ .filter_by(connector_name=connector_name)
72
+ .first()
73
+ )
74
+ if connector:
75
+ attributes = {
76
+ col.name: getattr(connector, col.name)
77
+ for col in Connectors.__table__.columns
78
+ }
79
+ return attributes
80
+ else:
81
+ raise NoResultFound
82
+
83
+
84
+class WazuhIndexerConnector(Connector):
85
+ """
86
+ This class represents a connector for the Wazuh indexer service. It is a subclass of Connector.
87
+
88
+ :param connector_name: A string that specifies the name of the connector.
89
+ """
90
+
91
+ def __init__(self, connector_name):
92
+ super().__init__(attributes=self.get_connector_info_from_db(connector_name))
93
+
94
+ def verify_connection(self):
95
+ """
96
+ This method verifies the connection to the Wazuh indexer service.
97
+
98
+ :return: A dictionary containing the status of the connection attempt and information about the cluster's health.
99
+ """
100
+ logger.info(
101
+ f"Verifying the wazuh-indexer connection to {self.attributes['connector_url']}"
102
+ )
103
+ try:
104
+ es = Elasticsearch(
105
+ [self.attributes["connector_url"]],
106
+ http_auth=(
107
+ self.attributes["connector_username"],
108
+ self.attributes["connector_password"],
109
+ ),
110
+ verify_certs=False,
111
+ timeout=15,
112
+ max_retries=10,
113
+ retry_on_timeout=False,
114
+ )
115
+ cluster_health = es.cluster.health()
116
+ logger.info(f"Connection to {self.attributes['connector_url']} successful")
117
+ return {"connectionSuccessful": True}
118
+ except Exception as e:
119
+ logger.error(
120
+ f"Connection to {self.attributes['connector_url']} failed with error: {e}"
121
+ )
122
+ return {"connectionSuccessful": False, "clusterHealth": None}
123
+
124
+
125
+class GraylogConnector(Connector):
126
+ """
127
+ This class represents a connector for the Graylog service. It is a subclass of Connector.
128
+
129
+ :param connector_name: A string that specifies the name of the connector.
130
+ """
131
+
132
+ def __init__(self, connector_name):
133
+ super().__init__(attributes=self.get_connector_info_from_db(connector_name))
134
+
135
+ def verify_connection(self):
136
+ """
137
+ Verifies the connection to Graylog service.
138
+
139
+ Returns:
140
+ dict: A dictionary containing 'connectionSuccessful' status and 'roles' if the connection is successful.
141
+ """
142
+ logger.info(
143
+ f"Verifying the graylog connection to {self.attributes['connector_url']}"
144
+ )
145
+ try:
146
+ graylog_roles = requests.get(
147
+ f"{self.attributes['connector_url']}/api/authz/roles/user/{self.attributes['connector_username']}",
148
+ auth=(
149
+ self.attributes["connector_username"],
150
+ self.attributes["connector_password"],
151
+ ),
152
+ verify=False,
153
+ )
154
+ if graylog_roles.status_code == 200:
155
+ logger.info(
156
+ f"Connection to {self.attributes['connector_url']} successful"
157
+ )
158
+ return {"connectionSuccessful": True}
159
+ else:
160
+ logger.error(
161
+ f"Connection to {self.attributes['connector_url']} failed with error: {graylog_roles.text}"
162
+ )
163
+ return {"connectionSuccessful": False, "roles": None}
164
+ except Exception as e:
165
+ logger.error(
166
+ f"Connection to {self.attributes['connector_url']} failed with error: {e}"
167
+ )
168
+ return {"connectionSuccessful": False, "roles": None}
169
+
170
+
171
+class WazuhManagerConnector(Connector):
172
+ """
173
+ This class represents a connector for the Wazuh manager service. It is a subclass of Connector.
174
+
175
+ :param connector_name: A string that specifies the name of the connector.
176
+ """
177
+
178
+ def __init__(self, connector_name):
179
+ super().__init__(attributes=self.get_connector_info_from_db(connector_name))
180
+
181
+ def verify_connection(self):
182
+ """
183
+ Verifies the connection to Wazuh manager service.
184
+
185
+ Returns:
186
+ dict: A dictionary containing 'connectionSuccessful' status and 'authToken' if the connection is successful.
187
+ """
188
+ logger.info(
189
+ f"Verifying the wazuh-manager connection to {self.attributes['connector_url']}"
190
+ )
191
+ try:
192
+ wazuh_auth_token = requests.get(
193
+ f"{self.attributes['connector_url']}/security/user/authenticate",
194
+ auth=(
195
+ self.attributes["connector_username"],
196
+ self.attributes["connector_password"],
197
+ ),
198
+ verify=False,
199
+ )
200
+ if wazuh_auth_token.status_code == 200:
201
+ logger.debug("Wazuh Authentication Token successful")
202
+ wazuh_auth_token = wazuh_auth_token.json()
203
+ wazuh_auth_token = wazuh_auth_token["data"]["token"]
204
+ return {"connectionSuccessful": True, "authToken": wazuh_auth_token}
205
+ else:
206
+ logger.error(
207
+ f"Connection to {self.attributes['connector_url']} failed with error: {wazuh_auth_token.text}"
208
+ )
209
+ return {"connectionSuccessful": False, "authToken": None}
210
+ except Exception as e:
211
+ logger.error(
212
+ f"Connection to {self.attributes['connector_url']} failed with error: {e}"
213
+ )
214
+ return {"connectionSuccessful": False, "authToken": None}
215
+
216
+ def get_auth_token(self):
217
+ """
218
+ Returns the authentication token for the Wazuh manager service.
219
+
220
+ Returns:
221
+ str: Authentication token for the Wazuh manager service.
222
+ """
223
+ return self.verify_connection()["authToken"]
224
+
225
+
226
+class ShuffleConnector(Connector):
227
+ """
228
+ This class represents a connector for the Shuffle service. It is a subclass of Connector.
229
+
230
+ :param connector_name: A string that specifies the name of the connector.
231
+ """
232
+
233
+ def __init__(self, connector_name):
234
+ super().__init__(attributes=self.get_connector_info_from_db(connector_name))
235
+
236
+ def verify_connection(self):
237
+ """
238
+ Verifies the connection to Shuffle service.
239
+
240
+ Returns:
241
+ dict: A dictionary containing 'connectionSuccessful' status and 'apps' if the connection is successful.
242
+ """
243
+ logger.info(
244
+ f"Verifying the shuffle connection to {self.attributes['connector_url']}"
245
+ )
246
+ try:
247
+ headers = {
248
+ "Authorization": f"Bearer {self.attributes['connector_api_key']}"
249
+ }
250
+ shuffle_apps = requests.get(
251
+ f"{self.attributes['connector_url']}/api/v1/apps",
252
+ headers=headers,
253
+ verify=False,
254
+ )
255
+ if shuffle_apps.status_code == 200:
256
+ logger.info(
257
+ f"Connection to {self.attributes['connector_url']} successful"
258
+ )
259
+ return {"connectionSuccessful": True}
260
+ else:
261
+ logger.error(
262
+ f"Connection to {self.attributes['connector_url']} failed with error: {shuffle_apps.text}"
263
+ )
264
+ return {"connectionSuccessful": False}
265
+ except Exception as e:
266
+ logger.error(
267
+ f"Connection to {self.attributes['connector_url']} failed with error: {e}"
268
+ )
269
+ return {"connectionSuccessful": False}
270
+
271
+
272
+class DfirIrisConnector(Connector):
273
+ """
274
+ This class represents a connector for the DFIR IRIS service. It is a subclass of Connector.
275
+
276
+ :param connector_name: A string that specifies the name of the connector.
277
+ """
278
+
279
+ def __init__(self, connector_name):
280
+ super().__init__(attributes=self.get_connector_info_from_db(connector_name))
281
+
282
+ def verify_connection(self):
283
+ """
284
+ Verifies the connection to DFIR IRIS service.
285
+
286
+ Returns:
287
+ dict: A dictionary containing 'connectionSuccessful' status and 'response' if the connection is successful.
288
+ """
289
+ logger.info(
290
+ f"Verifying the dfir-iris connection to {self.attributes['connector_url']}"
291
+ )
292
+ try:
293
+ headers = {
294
+ "Authorization": f"Bearer {self.attributes['connector_api_key']}"
295
+ }
296
+ dfir_iris = requests.get(
297
+ f"{self.attributes['connector_url']}/api/ping",
298
+ headers=headers,
299
+ verify=False,
300
+ )
301
+ # See if 200 is returned
302
+ if dfir_iris.status_code == 200:
303
+ logger.info(
304
+ f"Connection to {self.attributes['connector_url']} successful"
305
+ )
306
+ return {"connectionSuccessful": True}
307
+ else:
308
+ logger.error(
309
+ f"Connection to {self.attributes['connector_url']} failed with error: {dfir_iris.text}"
310
+ )
311
+ return {"connectionSuccessful": False, "response": None}
312
+ except Exception as e:
313
+ logger.error(
314
+ f"Connection to {self.attributes['connector_url']} failed with error: {e}"
315
+ )
316
+ return {"connectionSuccessful": False, "response": None}
317
+
318
+
319
+class VelociraptorConnector(Connector):
320
+ """
321
+ A connector for the Velociraptor service, a subclass of Connector.
322
+
323
+ Args:
324
+ connector_name (str): The name of the connector.
325
+ """
326
+
327
+ def __init__(self, connector_name):
328
+ super().__init__(attributes=self.get_connector_info_from_db(connector_name))
329
+
330
+ def verify_connection(self):
331
+ """
332
+ Verifies the connection to Velociraptor service.
333
+
334
+ Returns:
335
+ dict: A dictionary containing 'connectionSuccessful' status and 'response' if the connection is successful.
336
+ """
337
+ try:
338
+ connector_api_key = self.attributes["connector_api_key"]
339
+
340
+ with open(connector_api_key, "r") as f:
341
+ api_key = f.read()
342
+
343
+ try:
344
+ config = pyvelociraptor.LoadConfigFile(connector_api_key)
345
+ creds = grpc.ssl_channel_credentials(
346
+ root_certificates=config["ca_certificate"].encode("utf8"),
347
+ private_key=config["client_private_key"].encode("utf8"),
348
+ certificate_chain=config["client_cert"].encode("utf8"),
349
+ )
350
+
351
+ options = (("grpc.ssl_target_name_override", "VelociraptorServer"),)
352
+
353
+ with grpc.secure_channel(
354
+ config["api_connection_string"], creds, options
355
+ ) as channel:
356
+ stub = api_pb2_grpc.APIStub(channel)
357
+ client_query = "SELECT * FROM info()"
358
+
359
+ client_request = api_pb2.VQLCollectorArgs(
360
+ max_wait=60,
361
+ Query=[
362
+ api_pb2.VQLRequest(
363
+ Name="ClientQuery",
364
+ VQL=client_query,
365
+ ),
366
+ ],
367
+ )
368
+
369
+ r = []
370
+ for response in stub.Query(client_request):
371
+ if response.Response:
372
+ r = r + json.loads(response.Response)
373
+ return {"connectionSuccessful": True}
374
+ except Exception as e:
375
+ logger.error(f"Failed to verify connection to Velociraptor: {e}")
376
+ return {"connectionSuccessful": False, "response": None}
377
+ except Exception as e:
378
+ logger.error(f"Failed to get connector_api_key from the database: {e}")
379
+ return {"connectionSuccessful": False, "response": None}
380
+
381
+
382
+class RabbitMQConnector(Connector):
383
+ """
384
+ A connector for the RabbitMQ service, a subclass of Connector.
385
+
386
+ Args:
387
+ connector_name (str): The name of the connector.
388
+ """
389
+
390
+ def __init__(self, connector_name):
391
+ super().__init__(attributes=self.get_connector_info_from_db(connector_name))
392
+
393
+ def verify_connection(self):
394
+ """
395
+ Verifies the connection to RabbitMQ service.
396
+ """
397
+ logger.info(
398
+ f"Verifying the rabbitmq connection to {self.attributes['connector_url']}"
399
+ )
400
+ try:
401
+ # For the connector_url, strip out the host and port and use that for the connection
402
+ # This is because the connection string is not in the format that pika expects
403
+ connector_host, connector_port = self.attributes["connector_url"].split(":")
404
+ connector_port = int(connector_port)
405
+
406
+ credentials = pika.PlainCredentials(
407
+ self.attributes["connector_username"],
408
+ self.attributes["connector_password"],
409
+ )
410
+ parameters = pika.ConnectionParameters(
411
+ connector_host,
412
+ connector_port,
413
+ credentials=credentials,
414
+ )
415
+ connection = pika.BlockingConnection(parameters)
416
+ if connection.is_open:
417
+ logger.info(
418
+ f"Connection to {self.attributes['connector_url']} successful"
419
+ )
420
+ return {"connectionSuccessful": True}
421
+ else:
422
+ logger.error(f"Connection to {self.attributes['connector_url']} failed")
423
+ return {"connectionSuccessful": False, "response": None}
424
+ except Exception as e:
425
+ logger.error(
426
+ f"Connection to {self.attributes['connector_url']} failed with error: {e}"
427
+ )
428
+ return {"connectionSuccessful": False, "response": None}
429
+
430
+
431
+class ConnectorFactory:
432
+ """
433
+ This class represents a factory for creating connector instances.
434
+
435
+ :param creators: A dictionary mapping connector keys to their corresponding creator names.
436
+ """
437
+
438
+ def __init__(self):
439
+ """
440
+ Initialize a new instance of the ConnectorFactory.
441
+ """
442
+ self._creators = {}
443
+
444
+ def register_creator(self, key, creator):
445
+ """
446
+ Register a new connector creator.
447
+
448
+ :param key: The key of the connector.
449
+ :param creator: The creator of the connector.
450
+ """
451
+ self._creators[key] = creator
452
+
453
+ def create(self, key, connector_name):
454
+ """
455
+ Create a new connector instance.
456
+
457
+ :param key: The key of the connector.
458
+ :param connector_name: The name of the connector.
459
+
460
+ :return: A new instance of the connector.
461
+
462
+ :raises ValueError: If the key is not found in the list of creators.
463
+ """
464
+ creator = self._creators.get(key)
465
+ if not creator:
466
+ raise ValueError(key)
467
+ # use dynamic_import to get the class and initialize it
468
+ connector_class = dynamic_import("app.models.connectors", creator)
469
+ return connector_class(connector_name)
470
+
471
+
472
+# Instantiate factory
473
+connector_factory = ConnectorFactory()
474
+
475
+
476
+# Register connector creators
477
+connector_factory.register_creator("Wazuh-Indexer", "WazuhIndexerConnector")
478
+connector_factory.register_creator("Graylog", "GraylogConnector")
479
+connector_factory.register_creator("Wazuh-Manager", "WazuhManagerConnector")
480
+connector_factory.register_creator("DFIR-IRIS", "DfirIrisConnector")
481
+connector_factory.register_creator("Velociraptor", "VelociraptorConnector")
482
+connector_factory.register_creator("RabbitMQ", "RabbitMQConnector")
483
+connector_factory.register_creator("Shuffle", "ShuffleConnector")