@cryptotaxi247 / CoPilot / commits / 9c55d6e2

Influxdb (#21)

* added InfluxDB connector * Influxdb Checks route and connector * precommit fixes

taylor_socfortress committed Jul 13, 2023 at 14:58 UTC 9c55d6e2358a543fee9ab76dcfdc2be4b1fbce6d
6 files changed +280
backend/app/__init__.py
+2
@@ -49,6 +49,7 @@ from app.routes.alerts import bp as alerts_bp
49 from app.routes.connectors import bp as connectors_bp
50 from app.routes.dfir_iris import bp as dfir_iris_bp
51 from app.routes.graylog import bp as graylog_bp
52 +from app.routes.influxdb import bp as influxdb_bp
53 from app.routes.rules import bp as rules_bp
54 from app.routes.shuffle import bp as shuffle_bp
55 from app.routes.sublime import bp as sublime_bp
@@ -65,3 +66,4 @@ app.register_blueprint(shuffle_bp) # Register the shuffle blueprint
66 app.register_blueprint(velociraptor_bp) # Register the velociraptor blueprint
67 app.register_blueprint(dfir_iris_bp) # Register the dfir_iris blueprint
68 app.register_blueprint(sublime_bp) # Register the sublime blueprint
69 +app.register_blueprint(influxdb_bp) # Register the influxdb blueprint
backend/app/models/connectors.py
+50
@@ -425,6 +425,55 @@ class SublimeConnector(Connector):
425 return {"connectionSuccessful": False, "response": None}
426
427
428 +class InfluxDBConnector(Connector):
429 + """
430 + A connector for the InfluxDB service, a subclass of Connector.
431 +
432 + Args:
433 + connector_name (str): The name of the connector.
434 + """
435 +
436 + def __init__(self, connector_name: str):
437 + super().__init__(attributes=self.get_connector_info_from_db(connector_name))
438 +
439 + def verify_connection(self) -> Dict[str, Any]:
440 + """
441 + Verifies the connection to InfluxDB service.
442 +
443 + Returns:
444 + dict: A dictionary containing 'connectionSuccessful' status and 'response' if the connection is successful.
445 + """
446 + logger.info(
447 + f"Verifying the InfluxDB connection to {self.attributes['connector_url']}",
448 + )
449 + try:
450 + headers = {
451 + "Authorization": f"Token {self.attributes['connector_api_key']}",
452 + "Content-Type": "application/json",
453 + "Accept": "application/json",
454 + }
455 + influxdb = requests.get(
456 + f"{self.attributes['connector_url']}/api/v2/buckets",
457 + headers=headers,
458 + verify=False,
459 + )
460 + if influxdb.status_code == 200:
461 + logger.info(
462 + f"Connection to {self.attributes['connector_url']} successful",
463 + )
464 + return {"connectionSuccessful": True}
465 + else:
466 + logger.error(
467 + f"Connection to {self.attributes['connector_url']} failed with error: {influxdb.text}",
468 + )
469 + return {"connectionSuccessful": False, "response": None}
470 + except Exception as e:
471 + logger.error(
472 + f"Connection to {self.attributes['connector_url']} failed with error: {e}",
473 + )
474 + return {"connectionSuccessful": False, "response": None}
475 +
476 +
477 class RabbitMQConnector(Connector):
478 """
479 A connector for the RabbitMQ service, a subclass of Connector.
@@ -528,3 +577,4 @@ connector_factory.register_creator("Velociraptor", "VelociraptorConnector")
577 connector_factory.register_creator("RabbitMQ", "RabbitMQConnector")
578 connector_factory.register_creator("Shuffle", "ShuffleConnector")
579 connector_factory.register_creator("Sublime", "SublimeConnector")
580 +connector_factory.register_creator("InfluxDB", "InfluxDBConnector")
backend/app/routes/influxdb.py new
+23
@@ -0,0 +1,23 @@
1 +from flask import Blueprint
2 +from flask import jsonify
3 +from loguru import logger
4 +
5 +from app.services.InfluxDB.checks import InfluxDBChecksService
6 +
7 +bp = Blueprint("influxdb", __name__)
8 +
9 +
10 +@bp.route("/influxdb/checks", methods=["GET"])
11 +def get_checks() -> jsonify:
12 + """
13 + Endpoint to retrive list of InfluxDB checks.
14 + Requires API token of `admin` user.
15 +
16 + Returns:
17 + jsonify: A JSON response containing the list of all checks from InfluxDB.
18 +
19 + """
20 + logger.info("Received request to get all InfluxDB checks")
21 + service = InfluxDBChecksService.from_connector_details("InfluxDB")
22 + checks = service.collect_checks()
23 + return jsonify(checks)
backend/app/services/InfluxDB/checks.py new
+122
@@ -0,0 +1,122 @@
1 +from typing import Dict
2 +from typing import List
3 +from typing import Union
4 +
5 +import requests
6 +from loguru import logger
7 +
8 +from app.services.InfluxDB.universal import UniversalService
9 +
10 +
11 +class InvalidPayloadError(Exception):
12 + """
13 + Exception to be raised when the payload is invalid.
14 + """
15 +
16 + pass
17 +
18 +
19 +class InfluxDBSession:
20 + """
21 + Handles the session and connection to the InfluxDB server.
22 +
23 + Attributes:
24 + session (requests.Session): The session object for making HTTP requests.
25 + connector_url (str): The base URL for the InfluxDB API.
26 + """
27 +
28 + def __init__(self, connector_url: str, connector_api_key: str):
29 + """
30 + The constructor for InfluxDBSession class.
31 +
32 + Args:
33 + connector_url (str): The base URL for the InfluxDB API.
34 + connector_api_key (str): The API key for the InfluxDB API.
35 + """
36 + self.session = requests.Session()
37 + self.session.headers.update(
38 + {"Authorization": f"Bearer {connector_api_key}", "Content-Type": "application/json"},
39 + )
40 + self.connector_url = connector_url
41 +
42 + def send_request(self, url: str, params: Dict = None, verify: bool = False) -> requests.Response:
43 + """
44 + Sends a GET request to a specific URL.
45 +
46 + Args:
47 + url (str): The URL to send the GET request to.
48 + params (Dict, optional): The params to send with the GET request. Defaults to None.
49 + verify (bool, optional): Whether to verify the SSL certificate. Defaults to False.
50 +
51 + Returns:
52 + requests.Response: The response object from the GET request.
53 + """
54 + return self.session.get(url, params=params, verify=verify)
55 +
56 +
57 +class InfluxDBChecksService:
58 + """
59 + Handles operations related to InfluxDB alerts.
60 +
61 + Attributes:
62 + session (InfluxDBSession): The session object for making HTTP requests.
63 + connector_url (str): The base URL for the InfluxDB API.
64 + connector_api_key (str): The API key for the InfluxDB API.
65 + """
66 +
67 + def __init__(self, session: InfluxDBSession, connector_url: str, connector_api_key: str):
68 + """
69 + The constructor for InfluxDBChecksService class.
70 +
71 + Args:
72 + session (InfluxDBSession): The session object for making HTTP requests.
73 + connector_url (str): The base URL for the InfluxDB API.
74 + connector_api_key (str): The API key for the InfluxDB API.
75 + """
76 + self.session = session
77 + self.connector_url = connector_url
78 + self.connector_api_key = connector_api_key
79 +
80 + @classmethod
81 + def from_connector_details(cls, connector_name: str) -> "InfluxDBChecksService":
82 + """
83 + Creates an instance of InfluxDBChecksService using connector details.
84 +
85 + Args:
86 + connector_name (str): The name of the connector.
87 +
88 + Returns:
89 + InfluxDBChecksService: An instance of the class.
90 + """
91 + connector_url, connector_api_key = UniversalService().collect_influxdb_details(connector_name)
92 + session = InfluxDBSession(connector_url, connector_api_key)
93 + return cls(session, connector_url, connector_api_key)
94 +
95 + def collect_checks(self) -> List[Dict[str, Union[str, int]]]:
96 + """
97 + Collects all checks from InfluxDB.
98 +
99 + Returns:
100 + List[Dict[str, Union[str, int]]]: A list of all checks from InfluxDB.
101 + """
102 + logger.info("Collecting checks from InfluxDB")
103 + url = f"{self.connector_url}/api/v2/checks"
104 + params = {"orgID": "a1b203a448a55d31"}
105 + response = self.session.send_request(url=url, params=params)
106 + if response.status_code != 200:
107 + logger.error("Failed to collect checks from InfluxDB")
108 + logger.error(response.text)
109 + raise Exception("Failed to collect checks from InfluxDB")
110 + checks = []
111 + for check in response.json().get("checks"):
112 + checks.append(
113 + {
114 + "check_id": check.get("id"),
115 + "check_name": check.get("name"),
116 + "check_type": check.get("type"),
117 + "check_status": check.get("status"),
118 + "check_last_triggered": check.get("latestCompleted"),
119 + },
120 + )
121 + logger.info("Successfully collected checks from InfluxDB")
122 + return {"success": True, "message": "Checks received", "checks": checks}
backend/app/services/InfluxDB/universal.py new
+36
@@ -0,0 +1,36 @@
1 +from app.models.connectors import Connector
2 +from app.models.connectors import connector_factory
3 +
4 +
5 +class UniversalService:
6 + """
7 + A service class that encapsulates the logic for polling messages from InfluxDB.
8 + """
9 +
10 + def __init__(self) -> None:
11 + self.collect_influxdb_details("InfluxDB")
12 + (
13 + self.connector_url,
14 + self.connector_api_key,
15 + ) = self.collect_influxdb_details("InfluxDB")
16 +
17 + def collect_influxdb_details(self, connector_name: str):
18 + """
19 + Collects the details of the InfluxDB connector.
20 +
21 + Args:
22 + connector_name (str): The name of the InfluxDB connector.
23 +
24 + Returns:
25 + tuple: A tuple containing the connection URL, and api key.
26 + """
27 + connector_instance = connector_factory.create(connector_name, connector_name)
28 + connection_successful = connector_instance.verify_connection()
29 + if connection_successful:
30 + connection_details = Connector.get_connector_info_from_db(connector_name)
31 + return (
32 + connection_details.get("connector_url"),
33 + connection_details.get("connector_api_key"),
34 + )
35 + else:
36 + return None, None
backend/app/static/swagger.json
+47
@@ -81,6 +81,14 @@
81 "description": "Find out more",
82 "url": "http://swagger.io"
83 }
84 + },
85 + {
86 + "name": "InfluxDB",
87 + "description": "Everything about InfluxDB",
88 + "externalDocs": {
89 + "description": "Find out more",
90 + "url": "http://swagger.io"
91 + }
92 }
93 ],
94 "paths": {
@@ -2018,6 +2026,45 @@
2026 "operationId": "createAlert",
2027 "tags": ["Sublime"]
2028 }
2029 + },
2030 + "/influxdb/checks": {
2031 + "get": {
2032 + "summary": "Get all checks",
2033 + "description": "Endpoint to get all checks.",
2034 + "responses": {
2035 + "200": {
2036 + "description": "Successful operation",
2037 + "content": {
2038 + "application/json": {
2039 + "schema": {
2040 + "type": "object",
2041 + "properties": {
2042 + "checks": {
2043 + "type": "array",
2044 + "items": {
2045 + "type": "object",
2046 + "description": "Check details"
2047 + }
2048 + }
2049 + }
2050 + }
2051 + }
2052 + }
2053 + },
2054 + "default": {
2055 + "description": "Unexpected error",
2056 + "content": {
2057 + "application/json": {
2058 + "schema": {
2059 + "$ref": "#/components/schemas/Error"
2060 + }
2061 + }
2062 + }
2063 + }
2064 + },
2065 + "operationId": "getAllChecks",
2066 + "tags": ["InfluxDB"]
2067 + }
2068 }
2069 },
2070 "components": {