Create models.py
taylor_socfortress committed
Jul 10, 2023 at 16:39 UTC
8f09eb00a56251b3a73e59d2f99d181ac479eb97
1 file changed
+311
backend/app/models/models.py
new
+311
@@ -0,0 +1,311 @@
1
+from datetime import datetime
2
+
3
+from loguru import logger
4
+from sqlalchemy.dialects.postgresql import JSONB # Add this line
5
+
6
+from app import db
7
+from app import ma
8
+
9
+
10
+class ConnectorsAvailable(db.Model):
11
+ id = db.Column(db.Integer, primary_key=True)
12
+ connector_name = db.Column(db.String(100), unique=True)
13
+ connector_description = db.Column(db.String(100))
14
+ connector_supports = db.Column(db.String(100))
15
+ connector_configured = db.Column(db.Boolean, default=False)
16
+ connector_verified = db.Column(db.Boolean, default=False)
17
+
18
+ def __init__(self, connector_name, connector_supports):
19
+ self.connector_name = connector_name
20
+ self.connector_supports = connector_supports
21
+
22
+ def __repr__(self):
23
+ return f"<ConnectorsAvailble {self.connector_name}>"
24
+
25
+
26
+class ConnectorsAvailableSchema(ma.Schema):
27
+ class Meta:
28
+ fields = (
29
+ "id",
30
+ "connector_name",
31
+ "connector_description",
32
+ "connector_supports",
33
+ "connector_configured",
34
+ "connector_verified",
35
+ )
36
+
37
+
38
+connector_available_schema = ConnectorsAvailableSchema()
39
+connectors_available_schema = ConnectorsAvailableSchema(many=True)
40
+
41
+
42
+# Class for the connector which will store the endpoint url, connector name, connector type, connector last updated,
43
+# username and password
44
+class Connectors(db.Model):
45
+ id = db.Column(db.Integer, primary_key=True)
46
+ connector_name = db.Column(db.String(100), unique=True)
47
+ connector_type = db.Column(db.String(100))
48
+ connector_url = db.Column(db.String(100))
49
+ connector_last_updated = db.Column(db.DateTime, default=datetime.utcnow)
50
+ connector_username = db.Column(db.String(100))
51
+ connector_password = db.Column(db.String(100))
52
+ connector_api_key = db.Column(db.String(100))
53
+
54
+ def __init__(
55
+ self,
56
+ connector_name,
57
+ connector_type,
58
+ connector_url,
59
+ connector_username,
60
+ connector_password,
61
+ connector_api_key,
62
+ ):
63
+ self.connector_name = connector_name
64
+ self.connector_type = connector_type
65
+ self.connector_url = connector_url
66
+ self.connector_username = connector_username
67
+ self.connector_password = connector_password
68
+ # If the `connector_name` is `shuffle` or `dfir-irs` then set the `connector_api_key`. Otherwise set it to
69
+ # `None`
70
+ if (
71
+ connector_name.lower() == "shuffle"
72
+ or connector_name.lower() == "dfir-irs"
73
+ or connector_name.lower() == "velociraptor"
74
+ ):
75
+ logger.info(f"Setting the API key for {connector_name}")
76
+ self.connector_api_key = connector_api_key
77
+ else:
78
+ logger.info(f"Not setting the API key for {connector_name}")
79
+ self.connector_api_key = None
80
+
81
+ def __repr__(self):
82
+ return f"<Connectors {self.connector_name}>"
83
+
84
+
85
+class ConnectorsSchema(ma.Schema):
86
+ class Meta:
87
+ fields = (
88
+ "id",
89
+ "connector_name",
90
+ "connector_type",
91
+ "connector_url",
92
+ "connector_last_updated",
93
+ "connector_username",
94
+ "connector_password",
95
+ "connector_api_key",
96
+ )
97
+
98
+
99
+connector_schema = ConnectorsSchema()
100
+connectors_schema = ConnectorsSchema(many=True)
101
+
102
+
103
+# Class for the disabled rule IDs which will store the rule ID, previous configuration, new configuration, reason for
104
+# disabling, date disabled, and the length of time the rule will be disabled for
105
+# Path: backend\app\models.py
106
+# class DisabledRules(db.Model):
107
+# id = db.Column(db.Integer, primary_key=True)
108
+# rule_id = db.Column(db.String(100))
109
+# previous_level = db.Column(db.String(1000))
110
+# new_level = db.Column(db.String(1000))
111
+# reason_for_disabling = db.Column(db.String(100))
112
+# date_disabled = db.Column(db.DateTime, default=datetime.utcnow)
113
+# length_of_time = db.Column(db.Integer)
114
+
115
+# def __init__(
116
+# self,
117
+# rule_id,
118
+# previous_level,
119
+# new_level,
120
+# reason_for_disabling,
121
+# length_of_time,
122
+# ):
123
+# self.rule_id = rule_id
124
+# self.previous_level = previous_level
125
+# self.new_level = new_level
126
+# self.reason_for_disabling = reason_for_disabling
127
+# self.length_of_time = length_of_time
128
+
129
+# def __repr__(self):
130
+# return f"<DisabledRules {self.rule_id}>"
131
+
132
+
133
+# class DisabledRulesSchema(ma.Schema):
134
+# class Meta:
135
+# fields = (
136
+# "id",
137
+# "rule_id",
138
+# "previous_level",
139
+# "new_level",
140
+# "reason_for_disabling",
141
+# "date_disabled",
142
+# "length_of_time",
143
+# )
144
+
145
+
146
+# disabled_rule_schema = DisabledRulesSchema()
147
+# disabled_rules_schema = DisabledRulesSchema(many=True)
148
+
149
+
150
+# Class for Wazuh Indexer allocation which stores disk stats and the host.
151
+# Generate timestamp for each entry and invoke every 5 minutes.
152
+# Path: backend\app\models.py
153
+class WazuhIndexerAllocation(db.Model):
154
+ id = db.Column(db.Integer, primary_key=True)
155
+ node = db.Column(db.String(100))
156
+ disk_used = db.Column(db.Float)
157
+ disk_available = db.Column(db.Float)
158
+ disk_total = db.Column(db.Float)
159
+ disk_percent = db.Column(db.Float)
160
+ timestamp = db.Column(db.DateTime, default=datetime.utcnow)
161
+
162
+ def __init__(
163
+ self,
164
+ node,
165
+ disk_used,
166
+ disk_available,
167
+ disk_total,
168
+ disk_percent,
169
+ ):
170
+ self.node = node
171
+ self.disk_used = disk_used
172
+ self.disk_available = disk_available
173
+ self.disk_total = disk_total
174
+ self.disk_percent = disk_percent
175
+
176
+ def __repr__(self):
177
+ return f"<WazuhIndexerAllocation {self.node}>"
178
+
179
+
180
+class WazuhIndexerAllocationSchema(ma.Schema):
181
+ class Meta:
182
+ fields = (
183
+ "id",
184
+ "node",
185
+ "disk_used",
186
+ "disk_available",
187
+ "disk_total",
188
+ "disk_percent",
189
+ "timestamp",
190
+ )
191
+
192
+
193
+wazuh_indexer_allocation_schema = WazuhIndexerAllocationSchema()
194
+wazuh_indexer_allocations_schema = WazuhIndexerAllocationSchema(many=True)
195
+
196
+
197
+# Class for Graylog allocation which stores throughput metrics
198
+# Generate timestamp for each entry and invoke every 5 minutes.
199
+# Path: backend\app\models.py
200
+class GraylogMetricsAllocation(db.Model):
201
+ id = db.Column(db.Integer, primary_key=True)
202
+ input_usage = db.Column(db.Float)
203
+ output_usage = db.Column(db.Float)
204
+ processor_usage = db.Column(db.Float)
205
+ input_1_sec_rate = db.Column(db.Float)
206
+ output_1_sec_rate = db.Column(db.Float)
207
+ total_input = db.Column(db.Float)
208
+ total_output = db.Column(db.Float)
209
+ timestamp = db.Column(db.DateTime, default=datetime.utcnow)
210
+
211
+ def __init__(
212
+ self,
213
+ input_usage,
214
+ output_usage,
215
+ processor_usage,
216
+ input_1_sec_rate,
217
+ output_1_sec_rate,
218
+ total_input,
219
+ total_output,
220
+ ):
221
+ self.input_usage = input_usage
222
+ self.output_usage = output_usage
223
+ self.processor_usage = processor_usage
224
+ self.input_1_sec_rate = input_1_sec_rate
225
+ self.output_1_sec_rate = output_1_sec_rate
226
+ self.total_input = total_input
227
+ self.total_output = total_output
228
+
229
+ def __repr__(self):
230
+ return f"<GraylogMetricsAllocation {self.id}>"
231
+
232
+
233
+class GraylogMetricsAllocationSchema(ma.Schema):
234
+ class Meta:
235
+ fields = (
236
+ "id",
237
+ "input_usage",
238
+ "output_usage",
239
+ "processor_usage",
240
+ "input_1_sec_rate",
241
+ "output_1_sec_rate",
242
+ "total_input",
243
+ "total_output",
244
+ "timestamp",
245
+ )
246
+
247
+
248
+graylog_metrics_allocation_schema = GraylogMetricsAllocationSchema()
249
+graylog_metrics_allocations_schema = GraylogMetricsAllocationSchema(many=True)
250
+
251
+
252
+# Class for cases which stores the case ID, case name, list of agents
253
+# Path: backend\app\models.py
254
+class Case(db.Model):
255
+ id = db.Column(db.Integer, primary_key=True)
256
+ case_id = db.Column(db.Integer)
257
+ case_name = db.Column(db.String(100))
258
+ agents = db.Column(db.String(1000))
259
+
260
+ def __init__(self, case_id, case_name, agents):
261
+ self.case_id = case_id
262
+ self.case_name = case_name
263
+ self.agents = agents
264
+
265
+ def __repr__(self):
266
+ return f"<Case {self.case_id}>"
267
+
268
+
269
+class CaseSchema(ma.Schema):
270
+ class Meta:
271
+ fields = (
272
+ "id",
273
+ "case_id",
274
+ "case_name",
275
+ "agents",
276
+ )
277
+
278
+
279
+case_schema = CaseSchema()
280
+cases_schema = CaseSchema(many=True)
281
+
282
+
283
+# Class for artifacts collected which stores the artifact name, artificat results (json), hostname
284
+# Path: backend\app\models.py
285
+class Artifact(db.Model):
286
+ id = db.Column(db.Integer, primary_key=True)
287
+ artifact_name = db.Column(db.String(100))
288
+ artifact_results = db.Column(JSONB)
289
+ hostname = db.Column(db.String(100))
290
+
291
+ def __init__(self, artifact_name, artifact_results, hostname):
292
+ self.artifact_name = artifact_name
293
+ self.artifact_results = artifact_results
294
+ self.hostname = hostname
295
+
296
+ def __repr__(self):
297
+ return f"<Artifact {self.artifact_name}>"
298
+
299
+
300
+class ArtifactSchema(ma.Schema):
301
+ class Meta:
302
+ fields = (
303
+ "id",
304
+ "artifact_name",
305
+ "artifact_results",
306
+ "hostname",
307
+ )
308
+
309
+
310
+artifact_schema = ArtifactSchema()
311
+artifacts_schema = ArtifactSchema(many=True)