Create metrics.py
taylor_socfortress committed
Jul 10, 2023 at 16:43 UTC
dda81e8bcccc8a4ae57f2cb36cfc29d472fdd372
1 file changed
+227
backend/app/services/Graylog/metrics.py
new
+227
@@ -0,0 +1,227 @@
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 app.models.connectors import connector_factory, Connector, GraylogConnector
12
+from app.services.Graylog.universal import UniversalService
13
+
14
+
15
+class MetricsService:
16
+ """
17
+ A service class that encapsulates the logic for pulling metrics from Graylog.
18
+ """
19
+
20
+ METRIC_NAMES: Dict[str, str] = {
21
+ "org.graylog2.throughput.input.1-sec-rate": "input_1_sec_rate",
22
+ "org.graylog2.throughput.output.1-sec-rate": "output_1_sec_rate",
23
+ "org.graylog2.buffers.input.usage": "input_usage",
24
+ "org.graylog2.buffers.output.usage": "output_usage",
25
+ "org.graylog2.buffers.process.usage": "processor_usage",
26
+ "org.graylog2.throughput.input": "total_input",
27
+ "org.graylog2.throughput.output": "total_output",
28
+ }
29
+
30
+ HEADERS: Dict[str, str] = {"X-Requested-By": "CoPilot"}
31
+
32
+ def collect_uncommitted_journal_size(self):
33
+ """
34
+ Collects the journal size of uncommitted messages from Graylog.
35
+
36
+ Returns:
37
+ list: A list containing the metrics.
38
+ """
39
+ (
40
+ connector_url,
41
+ connector_username,
42
+ connector_password,
43
+ ) = UniversalService().collect_graylog_details("Graylog")
44
+ if (
45
+ connector_url is None
46
+ or connector_username is None
47
+ or connector_password is None
48
+ ):
49
+ return {"message": "Failed to collect Graylog details", "success": False}
50
+ else:
51
+ journal_size = self._collect_metrics_uncommitted_journal_size(
52
+ connector_url, connector_username, connector_password
53
+ )
54
+
55
+ if journal_size["success"] is False:
56
+ return journal_size
57
+ return journal_size
58
+
59
+ def collect_throughput_metrics(self):
60
+ """
61
+ Collects the following Graylog Metrics:
62
+ - Input Usage
63
+ - Output Usage
64
+ - Processor Usage
65
+ - Input 1 Seconds Rate
66
+ - Output 1 Seconds Rate
67
+ - Total Input
68
+ - Total Output
69
+
70
+ Returns:
71
+ list: A list containing the metrics.
72
+ """
73
+ (
74
+ connector_url,
75
+ connector_username,
76
+ connector_password,
77
+ ) = UniversalService().collect_graylog_details("Graylog")
78
+ if (
79
+ connector_url is None
80
+ or connector_username is None
81
+ or connector_password is None
82
+ ):
83
+ return {"message": "Failed to collect Graylog details", "success": False}
84
+ else:
85
+ throughput_usage = self._collect_metrics_throughput_usage(
86
+ connector_url, connector_username, connector_password
87
+ )
88
+
89
+ if throughput_usage["success"] is False:
90
+ return throughput_usage
91
+ return throughput_usage
92
+
93
+ def _collect_metrics_uncommitted_journal_size(
94
+ self, connector_url: str, connector_username: str, connector_password: str
95
+ ):
96
+ """
97
+ Collects the journal size of uncommitted messages from Graylog.
98
+
99
+ Args:
100
+ connector_url (str): The URL of the Graylog connector.
101
+ connector_username (str): The username of the Graylog connector.
102
+ connector_password (str): The password of the Graylog connector.
103
+
104
+ Returns:
105
+ int: The journal size.
106
+ """
107
+ try:
108
+ logger.info("Collecting journal size from Graylog")
109
+ headers = {"X-Requested-By": "CoPilot"}
110
+ # Get the Graylog Journal Size
111
+ uncommitted_journal_size_response = requests.get(
112
+ f"{connector_url}/api/system/journal",
113
+ headers=headers,
114
+ auth=(connector_username, connector_password),
115
+ verify=False,
116
+ )
117
+ uncommitted_journal_size = uncommitted_journal_size_response.json()
118
+
119
+ logger.info(
120
+ f"Received {uncommitted_journal_size} uncommitted journal entries from Graylog"
121
+ )
122
+ return {
123
+ "message": "Successfully retrieved journal size",
124
+ "success": True,
125
+ "uncommitted_journal_entries": uncommitted_journal_size.get(
126
+ "uncommitted_journal_entries", 0
127
+ ),
128
+ }
129
+ except Exception as e:
130
+ logger.error(f"Failed to collect journal size from Graylog: {e}")
131
+ return {
132
+ "message": "Failed to collect journal size from Graylog",
133
+ "success": False,
134
+ }
135
+
136
+ def _collect_metrics_throughput_usage(
137
+ self, connector_url: str, connector_username: str, connector_password: str
138
+ ) -> Dict[str, object]:
139
+ """
140
+ Collects throughput usage from Graylog.
141
+
142
+ Args:
143
+ connector_url (str): The URL of the Graylog connector.
144
+ connector_username (str): The username of the Graylog connector.
145
+ connector_password (str): The password of the Graylog connector.
146
+
147
+ Returns:
148
+ dict: A dictionary containing the throughput usage.
149
+ """
150
+ logger.info("Collecting throughput usage from Graylog")
151
+
152
+ try:
153
+ throughput_metrics = self._make_throughput_api_call(
154
+ connector_url, self.HEADERS, connector_username, connector_password
155
+ )
156
+ return self._parse_throughput_metrics(throughput_metrics)
157
+ except Exception as e:
158
+ logger.error(f"Failed to collect throughput usage from Graylog: {e}")
159
+ return {
160
+ "message": "Failed to collect throughput usage from Graylog",
161
+ "success": False,
162
+ }
163
+
164
+ def _make_throughput_api_call(
165
+ self,
166
+ connector_url: str,
167
+ headers: Dict[str, str],
168
+ connector_username: str,
169
+ connector_password: str,
170
+ ) -> Dict[str, object]:
171
+ """
172
+ Makes Throughput API call to Graylog.
173
+
174
+ Args:
175
+ connector_url (str): The URL of the Graylog connector.
176
+ headers (dict): The headers for the request.
177
+ connector_username (str): The username of the Graylog connector.
178
+ connector_password (str): The password of the Graylog connector.
179
+
180
+ Returns:
181
+ dict: The dictionary containing throughput metrics.
182
+ """
183
+ throughput = requests.get(
184
+ f"{connector_url}/api/system/metrics",
185
+ headers=headers,
186
+ auth=(connector_username, connector_password),
187
+ verify=False,
188
+ )
189
+ throughput_json = throughput.json()
190
+ throughput_metrics = throughput_json["gauges"]
191
+ input_output_throughput = throughput_json["counters"]
192
+
193
+ throughput_metrics.update(input_output_throughput) # Merge the two dictionaries
194
+ return throughput_metrics
195
+
196
+ def _parse_throughput_metrics(
197
+ self, throughput_metrics: Dict[str, object]
198
+ ) -> Dict[str, object]:
199
+ """
200
+ Parses throughput metrics.
201
+
202
+ Args:
203
+ throughput_metrics (dict): The dictionary containing throughput metrics.
204
+
205
+ Returns:
206
+ dict: The dictionary with parsed throughput metrics.
207
+ """
208
+ throughput_metrics_list = []
209
+ results = {}
210
+
211
+ for metric, data in throughput_metrics.items():
212
+ if metric in self.METRIC_NAMES:
213
+ value = data["value"] if "value" in data else data["count"]
214
+ throughput_metrics_list.append({"metric": metric, "value": value})
215
+
216
+ variable_name = self.METRIC_NAMES.get(metric)
217
+ if variable_name is not None:
218
+ results[variable_name] = value
219
+
220
+ logger.info(
221
+ f"Received throughput usage from Graylog: {throughput_metrics_list}"
222
+ )
223
+ return {
224
+ "message": "Successfully retrieved throughput usage",
225
+ "success": True,
226
+ "throughput_metrics": throughput_metrics_list,
227
+ }