Influx alerts (#22)
* restructure checks * influxdb alerts model and endpoint * influxdb POST alert
taylor_socfortress committed
Jul 13, 2023 at 20:13 UTC
ec2cacd600d32e1b53793b37965b7e081a581802
8 files changed
+467
-42
backend/app/__init__.py
+1
@@ -35,6 +35,7 @@ from app.models import artifacts # noqa: F401
35
from app.models import cases # noqa: F401
36
from app.models import connectors # noqa: F401
37
from app.models import graylog # noqa: F401
38
+from app.models import influxdb_alerts # noqa: F401
39
from app.models import models # noqa: F401
40
from app.models import rules # noqa: F401
41
from app.models import smtp # noqa: F401
backend/app/models/influxdb_alerts.py
new
+61
@@ -0,0 +1,61 @@
1
+from datetime import datetime
2
+
3
+from sqlalchemy import Column
4
+from sqlalchemy import DateTime
5
+from sqlalchemy import Integer
6
+from sqlalchemy import String
7
+
8
+from app import db
9
+from app import ma
10
+
11
+
12
+# Path: backend\app\models.py
13
+class InfluxDBAlerts(db.Model):
14
+ """
15
+ Class for InfluxDB Alerts which stores the message ID, and timestamp.
16
+ This class inherits from SQLAlchemy's Model class.
17
+ """
18
+
19
+ id: Column[Integer] = db.Column(db.Integer, primary_key=True)
20
+ check_name: Column[String] = db.Column(db.String(1000))
21
+ message: Column[String] = db.Column(db.String(1000))
22
+ timestamp: Column[DateTime] = db.Column(db.DateTime, default=datetime.utcnow)
23
+
24
+ def __init__(self, check_name: str, message: str):
25
+ """
26
+ Initialize a new instance of the InfluxDB Alerts class.
27
+
28
+ :param check_name: The Message ID of the alert.
29
+ """
30
+ self.check_name = check_name
31
+ self.message = message
32
+
33
+ def __repr__(self) -> str:
34
+ """
35
+ Returns a string representation of the Alert instance.
36
+
37
+ :return: A string representation of the Check Name.
38
+ """
39
+ return f"<Alert {self.check_name}>"
40
+
41
+
42
+class InfluxDBAlertsSchema(ma.Schema):
43
+ """
44
+ Schema for serializing and deserializing instances of the InfluxDB Alerts class.
45
+ """
46
+
47
+ class Meta:
48
+ """
49
+ Meta class defines the fields to be serialized/deserialized.
50
+ """
51
+
52
+ fields: tuple = (
53
+ "id",
54
+ "check_name",
55
+ "message",
56
+ "timestamp",
57
+ )
58
+
59
+
60
+InfluxDB_alert_schema: InfluxDBAlertsSchema = InfluxDBAlertsSchema()
61
+InfluxDB_alerts_schema: InfluxDBAlertsSchema = InfluxDBAlertsSchema(many=True)
backend/app/routes/influxdb.py
+44
@@ -1,7 +1,12 @@
1
+from typing import Any
2
+from typing import Dict
3
+
4
from flask import Blueprint
5
from flask import jsonify
6
+from flask import request
7
from loguru import logger
8
9
+from app.services.InfluxDB.alerts import InfluxDBAlertsService
10
from app.services.InfluxDB.checks import InfluxDBChecksService
11
12
bp = Blueprint("influxdb", __name__)
@@ -21,3 +26,42 @@ def get_checks() -> jsonify:
26
service = InfluxDBChecksService.from_connector_details("InfluxDB")
27
checks = service.collect_checks()
28
return jsonify(checks)
29
+
30
+
31
+@bp.route("/influxdb/checks/<check_id>", methods=["GET"])
32
+def get_check_query(check_id: str) -> jsonify:
33
+ """
34
+ Endpoint to retrive query of InfluxDB check.
35
+ Requires API token of `admin` user.
36
+
37
+ Returns:
38
+ jsonify: A JSON response containing the query of the check from InfluxDB.
39
+
40
+ """
41
+ logger.info("Received request to get InfluxDB check query")
42
+ service = InfluxDBChecksService.from_connector_details("InfluxDB")
43
+ check_query = service.collect_check_query(check_id)
44
+ return jsonify(check_query)
45
+
46
+
47
+@bp.route("/influxdb/alert", methods=["POST"])
48
+def put_alert() -> jsonify:
49
+ """
50
+ Endpoint to store alert in the `influxdb_alerts` table.
51
+ Invoked by the InfluxDB alert webhook which is configured in the InfluxDB UI.
52
+
53
+ Returns:
54
+ jsonify: A JSON response containing if the alert was stored successfully.
55
+ """
56
+ logger.info("Received request to store InfluxDB alert")
57
+ data: Dict[str, Any] = request.get_json()
58
+ logger.info(data)
59
+ service = InfluxDBAlertsService.from_connector_details("InfluxDB")
60
+
61
+ try:
62
+ check_name, message = service.validate_payload(data=data)
63
+ service.store_alerts(check_name=check_name, message=message)
64
+ return jsonify({"message": "Successfully stored alert.", "success": True}), 200
65
+ except Exception as e:
66
+ logger.error(f"Received invalid payload. {e}")
67
+ return jsonify({"message": "Invalid payload.", "success": False}), 400
backend/app/services/InfluxDB/alerts.py
new
+139
@@ -0,0 +1,139 @@
1
+from typing import Dict
2
+
3
+import requests
4
+from loguru import logger
5
+
6
+from app import db
7
+from app.models.influxdb_alerts import InfluxDBAlerts
8
+from app.services.InfluxDB.universal import UniversalService
9
+
10
+# OrgID for the InfluxDB server
11
+ORG_ID = "a1b203a448a55d31"
12
+
13
+
14
+class InvalidPayloadError(Exception):
15
+ """
16
+ Custom exception to be raised when the payload is invalid.
17
+ Inherits from the base Exception class.
18
+ """
19
+
20
+ pass
21
+
22
+
23
+class ChecksCollectionError(Exception):
24
+ """
25
+ Custom exception to be raised when there is a failure in collecting checks from InfluxDB.
26
+ Inherits from the base Exception class.
27
+ """
28
+
29
+ pass
30
+
31
+
32
+class InfluxDBSession:
33
+ """
34
+ Class to handle the session and connection to the InfluxDB server.
35
+
36
+ Attributes:
37
+ session: requests.Session object for making HTTP requests.
38
+ connector_url: string representing the base URL for the InfluxDB API.
39
+ """
40
+
41
+ def __init__(self, connector_url: str, connector_api_key: str):
42
+ """
43
+ Initializes InfluxDBSession with the connector URL and API key.
44
+
45
+ Args:
46
+ connector_url: string representing the base URL for the InfluxDB API.
47
+ connector_api_key: string representing the API key for the InfluxDB API.
48
+ """
49
+ self.session = requests.Session()
50
+ self.session.headers.update(
51
+ {"Authorization": f"Bearer {connector_api_key}", "Content-Type": "application/json"},
52
+ )
53
+ self.connector_url = connector_url
54
+
55
+ def send_request(self, url: str, params: Dict = None, verify: bool = False) -> requests.Response:
56
+ """
57
+ Sends a GET request to a specific URL.
58
+
59
+ Args:
60
+ url: string representing the URL to send the GET request to.
61
+ params: dictionary representing the params to send with the GET request. Defaults to None.
62
+ verify: boolean representing whether to verify the SSL certificate. Defaults to False.
63
+
64
+ Returns:
65
+ Response object from the GET request.
66
+ """
67
+ return self.session.get(url, params=params, verify=verify)
68
+
69
+
70
+class InfluxDBAlertsService:
71
+ """
72
+ Class to handle operations related to InfluxDB alerts.
73
+
74
+ Attributes:
75
+ session: InfluxDBSession object for making HTTP requests.
76
+ connector_url: string representing the base URL for the InfluxDB API.
77
+ connector_api_key: string representing the API key for the InfluxDB API.
78
+ """
79
+
80
+ def __init__(self, session: InfluxDBSession, connector_url: str, connector_api_key: str):
81
+ """
82
+ Initializes InfluxDBAlertsService with the session, connector URL, and API key.
83
+
84
+ Args:
85
+ session: InfluxDBSession object for making HTTP requests.
86
+ connector_url: string representing the base URL for the InfluxDB API.
87
+ connector_api_key: string representing the API key for the InfluxDB API.
88
+ """
89
+ self.session = session
90
+ self.connector_url = connector_url
91
+ self.connector_api_key = connector_api_key
92
+
93
+ @classmethod
94
+ def from_connector_details(cls, connector_name: str) -> "InfluxDBAlertsService":
95
+ """
96
+ Creates an instance of InfluxDBAlertsService using connector details.
97
+
98
+ Args:
99
+ connector_name: string representing the name of the connector.
100
+
101
+ Returns:
102
+ An instance of the InfluxDBAlertsService class.
103
+ """
104
+ connector_url, connector_api_key = UniversalService().collect_influxdb_details(connector_name)
105
+ session = InfluxDBSession(connector_url, connector_api_key)
106
+ return cls(session, connector_url, connector_api_key)
107
+
108
+ def validate_payload(self, data: Dict[str, object]) -> str:
109
+ """
110
+ Validates the payload received from the Sublime alert webhook.
111
+
112
+ Args:
113
+ data (Dict[str, object]): The data received from the webhook.
114
+
115
+ Returns:
116
+ str: The message ID from the payload.
117
+
118
+ Raises:
119
+ InvalidPayloadError: If the payload is invalid.
120
+ """
121
+ try:
122
+ check_name = data["_check_name"]
123
+ message = data["_message"]
124
+ return check_name, message
125
+ except KeyError:
126
+ raise InvalidPayloadError("Invalid payload.")
127
+
128
+ def store_alerts(self, check_name: str, message: str) -> None:
129
+ """
130
+ Stores the alerts in the database.
131
+
132
+ Args:
133
+ check_name: string representing the name of the check.
134
+ message: string representing the message of the alert.
135
+ """
136
+ logger.info("Storing alerts in the database.")
137
+ influxbd_alert = InfluxDBAlerts(check_name=check_name, message=message)
138
+ db.session.add(influxbd_alert)
139
+ db.session.commit()
backend/app/services/InfluxDB/checks.py
+62
-24
@@ -7,10 +7,23 @@ from loguru import logger
7
8
from app.services.InfluxDB.universal import UniversalService
9
10
+# OrgID for the InfluxDB server
11
+ORG_ID = "a1b203a448a55d31"
12
+
13
14
class InvalidPayloadError(Exception):
15
"""
13
- Exception to be raised when the payload is invalid.
16
+ Custom exception to be raised when the payload is invalid.
17
+ Inherits from the base Exception class.
18
+ """
19
+
20
+ pass
21
+
22
+
23
+class ChecksCollectionError(Exception):
24
+ """
25
+ Custom exception to be raised when there is a failure in collecting checks from InfluxDB.
26
+ Inherits from the base Exception class.
27
"""
28
29
pass
@@ -18,20 +31,20 @@ class InvalidPayloadError(Exception):
31
32
class InfluxDBSession:
33
"""
21
- Handles the session and connection to the InfluxDB server.
34
+ Class to handle the session and connection to the InfluxDB server.
35
36
Attributes:
24
- session (requests.Session): The session object for making HTTP requests.
25
- connector_url (str): The base URL for the InfluxDB API.
37
+ session: requests.Session object for making HTTP requests.
38
+ connector_url: string representing the base URL for the InfluxDB API.
39
"""
40
41
def __init__(self, connector_url: str, connector_api_key: str):
42
"""
30
- The constructor for InfluxDBSession class.
43
+ Initializes InfluxDBSession with the connector URL and API key.
44
45
Args:
33
- connector_url (str): The base URL for the InfluxDB API.
34
- connector_api_key (str): The API key for the InfluxDB API.
46
+ connector_url: string representing the base URL for the InfluxDB API.
47
+ connector_api_key: string representing the API key for the InfluxDB API.
48
"""
49
self.session = requests.Session()
50
self.session.headers.update(
@@ -44,34 +57,34 @@ class InfluxDBSession:
57
Sends a GET request to a specific URL.
58
59
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.
60
+ url: string representing the URL to send the GET request to.
61
+ params: dictionary representing the params to send with the GET request. Defaults to None.
62
+ verify: boolean representing whether to verify the SSL certificate. Defaults to False.
63
64
Returns:
52
- requests.Response: The response object from the GET request.
65
+ Response object from the GET request.
66
"""
67
return self.session.get(url, params=params, verify=verify)
68
69
70
class InfluxDBChecksService:
71
"""
59
- Handles operations related to InfluxDB alerts.
72
+ Class to handle operations related to InfluxDB alerts.
73
74
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.
75
+ session: InfluxDBSession object for making HTTP requests.
76
+ connector_url: string representing the base URL for the InfluxDB API.
77
+ connector_api_key: string representing the API key for the InfluxDB API.
78
"""
79
80
def __init__(self, session: InfluxDBSession, connector_url: str, connector_api_key: str):
81
"""
69
- The constructor for InfluxDBChecksService class.
82
+ Initializes InfluxDBChecksService with the session, connector URL, and API key.
83
84
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.
85
+ session: InfluxDBSession object for making HTTP requests.
86
+ connector_url: string representing the base URL for the InfluxDB API.
87
+ connector_api_key: string representing the API key for the InfluxDB API.
88
"""
89
self.session = session
90
self.connector_url = connector_url
@@ -83,10 +96,10 @@ class InfluxDBChecksService:
96
Creates an instance of InfluxDBChecksService using connector details.
97
98
Args:
86
- connector_name (str): The name of the connector.
99
+ connector_name: string representing the name of the connector.
100
101
Returns:
89
- InfluxDBChecksService: An instance of the class.
102
+ An instance of the InfluxDBChecksService class.
103
"""
104
connector_url, connector_api_key = UniversalService().collect_influxdb_details(connector_name)
105
session = InfluxDBSession(connector_url, connector_api_key)
@@ -97,16 +110,19 @@ class InfluxDBChecksService:
110
Collects all checks from InfluxDB.
111
112
Returns:
100
- List[Dict[str, Union[str, int]]]: A list of all checks from InfluxDB.
113
+ A list of dictionaries, where each dictionary represents a check from InfluxDB. Each dictionary includes
114
+ 'check_id', 'check_name', 'check_type', 'check_status', 'check_last_triggered' keys.
115
+ Additionally, a 'success' key is included in the returned list to indicate if the checks were successfully
116
+ collected, and a 'message' key is included to provide additional information.
117
"""
118
logger.info("Collecting checks from InfluxDB")
119
url = f"{self.connector_url}/api/v2/checks"
104
- params = {"orgID": "a1b203a448a55d31"}
120
+ params = {"orgID": ORG_ID} # Using the extracted variable
121
response = self.session.send_request(url=url, params=params)
122
if response.status_code != 200:
123
logger.error("Failed to collect checks from InfluxDB")
124
logger.error(response.text)
109
- raise Exception("Failed to collect checks from InfluxDB")
125
+ raise ChecksCollectionError("Failed to collect checks from InfluxDB") # Using the new exception
126
checks = []
127
for check in response.json().get("checks"):
128
checks.append(
@@ -120,3 +136,25 @@ class InfluxDBChecksService:
136
)
137
logger.info("Successfully collected checks from InfluxDB")
138
return {"success": True, "message": "Checks received", "checks": checks}
139
+
140
+ def collect_check_query(self, check_id: str) -> Dict[str, Union[str, int]]:
141
+ """
142
+ Collects the query for a specific check from InfluxDB.
143
+
144
+ Args:
145
+ check_id: string representing the ID of the check.
146
+
147
+ Returns:
148
+ A dictionary representing the query for the check. The dictionary includes 'query' and 'success' keys.
149
+ Additionally, a 'message' key is included to provide additional information.
150
+ """
151
+ logger.info(f"Collecting query for check {check_id} from InfluxDB")
152
+ url = f"{self.connector_url}/api/v2/checks/{check_id}"
153
+ params = {"orgID": ORG_ID}
154
+ response = self.session.send_request(url=url, params=params)
155
+ if response.status_code != 200:
156
+ logger.error(f"Failed to collect query for check {check_id} from InfluxDB")
157
+ logger.error(response.text)
158
+ raise ChecksCollectionError(f"Failed to collect query for check {check_id} from InfluxDB")
159
+ logger.info(f"Successfully collected query for check {check_id} from InfluxDB")
160
+ return {"success": True, "message": "Query received", "query": response.json()}
backend/app/services/InfluxDB/universal.py
+10
-18
@@ -2,28 +2,20 @@ from app.models.connectors import Connector
2
from app.models.connectors import connector_factory
3
4
5
+class ConnectionFailedError(Exception): # New exception class for connection failure
6
+ """Exception to be raised when connection to InfluxDB fails."""
7
+
8
+ pass
9
+
10
+
11
class UniversalService:
6
- """
7
- A service class that encapsulates the logic for polling messages from InfluxDB.
8
- """
12
+ """A service class that encapsulates the logic for polling messages from InfluxDB."""
13
14
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")
15
+ (self.connector_url, self.connector_api_key) = self.collect_influxdb_details("InfluxDB") # Removed redundant call
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
- """
18
+ """Collects the details of the InfluxDB connector."""
19
connector_instance = connector_factory.create(connector_name, connector_name)
20
connection_successful = connector_instance.verify_connection()
21
if connection_successful:
@@ -33,4 +25,4 @@ class UniversalService:
25
connection_details.get("connector_api_key"),
26
)
27
else:
36
- return None, None
28
+ raise ConnectionFailedError(f"Failed to connect to {connector_name}") # Raise an exception instead of returning None, None
backend/app/static/swagger.json
+116
@@ -2065,6 +2065,122 @@
2065
"operationId": "getAllChecks",
2066
"tags": ["InfluxDB"]
2067
}
2068
+ },
2069
+ "/influxdb/checks/{check_id}": {
2070
+ "get": {
2071
+ "summary": "Get check query",
2072
+ "description": "Endpoint to get check query.",
2073
+ "parameters": [
2074
+ {
2075
+ "name": "check_id",
2076
+ "in": "path",
2077
+ "description": "The ID of the check.",
2078
+ "required": true,
2079
+ "schema": {
2080
+ "type": "string"
2081
+ }
2082
+ }
2083
+ ],
2084
+ "responses": {
2085
+ "200": {
2086
+ "description": "Successful operation",
2087
+ "content": {
2088
+ "application/json": {
2089
+ "schema": {
2090
+ "type": "object",
2091
+ "properties": {
2092
+ "query": {
2093
+ "type": "string",
2094
+ "description": "The query of the check."
2095
+ }
2096
+ }
2097
+ }
2098
+ }
2099
+ }
2100
+ },
2101
+ "default": {
2102
+ "description": "Unexpected error",
2103
+ "content": {
2104
+ "application/json": {
2105
+ "schema": {
2106
+ "$ref": "#/components/schemas/Error"
2107
+ }
2108
+ }
2109
+ }
2110
+ }
2111
+ },
2112
+ "operationId": "getCheckQuery",
2113
+ "tags": ["InfluxDB"]
2114
+ }
2115
+ },
2116
+ "/influxdb/alert": {
2117
+ "post": {
2118
+ "tags": ["InfluxDB"],
2119
+ "summary": "Create alert",
2120
+ "description": "Endpoint to create alert.",
2121
+ "requestBody": {
2122
+ "content": {
2123
+ "application/json": {
2124
+ "schema": {
2125
+ "type": "object",
2126
+ "properties": {
2127
+ "_check_name": {
2128
+ "type": "string",
2129
+ "description": "The name of the check."
2130
+ },
2131
+ "_message": {
2132
+ "type": "string",
2133
+ "description": "The alert message."
2134
+ }
2135
+ },
2136
+ "required": ["_check_name", "_message"]
2137
+ }
2138
+ }
2139
+ }
2140
+ },
2141
+ "responses": {
2142
+ "200": {
2143
+ "description": "Successfully stored alert.",
2144
+ "content": {
2145
+ "application/json": {
2146
+ "schema": {
2147
+ "type": "object",
2148
+ "properties": {
2149
+ "message": {
2150
+ "type": "string",
2151
+ "example": "Successfully stored alert."
2152
+ },
2153
+ "success": {
2154
+ "type": "boolean",
2155
+ "example": true
2156
+ }
2157
+ }
2158
+ }
2159
+ }
2160
+ }
2161
+ },
2162
+ "400": {
2163
+ "description": "Invalid payload.",
2164
+ "content": {
2165
+ "application/json": {
2166
+ "schema": {
2167
+ "type": "object",
2168
+ "properties": {
2169
+ "message": {
2170
+ "type": "string",
2171
+ "example": "Invalid payload."
2172
+ },
2173
+ "success": {
2174
+ "type": "boolean",
2175
+ "example": false
2176
+ }
2177
+ }
2178
+ }
2179
+ }
2180
+ }
2181
+ }
2182
+ }
2183
+ }
2184
}
2185
},
2186
"components": {
backend/migrations/versions/cc9a9e5057ac_add_influxdb_alerts_model.py
new
+34
@@ -0,0 +1,34 @@
1
+"""Add Influxdb Alerts Model.
2
+
3
+Revision ID: cc9a9e5057ac
4
+Revises: 16a3bb6544b1
5
+Create Date: 2023-07-13 19:52:44.366350
6
+
7
+"""
8
+import sqlalchemy as sa
9
+from alembic import op
10
+
11
+# revision identifiers, used by Alembic.
12
+revision = "cc9a9e5057ac"
13
+down_revision = "16a3bb6544b1"
14
+branch_labels = None
15
+depends_on = None
16
+
17
+
18
+def upgrade():
19
+ # ### commands auto generated by Alembic - please adjust! ###
20
+ op.create_table(
21
+ "influx_db_alerts",
22
+ sa.Column("id", sa.Integer(), nullable=False),
23
+ sa.Column("check_name", sa.String(length=1000), nullable=True),
24
+ sa.Column("message", sa.String(length=1000), nullable=True),
25
+ sa.Column("timestamp", sa.DateTime(), nullable=True),
26
+ sa.PrimaryKeyConstraint("id"),
27
+ )
28
+ # ### end Alembic commands ###
29
+
30
+
31
+def downgrade():
32
+ # ### commands auto generated by Alembic - please adjust! ###
33
+ op.drop_table("influx_db_alerts")
34
+ # ### end Alembic commands ###