Create agents.py
taylor_socfortress committed
Jul 10, 2023 at 16:38 UTC
640681e2077010c13e5e0af66717579a0910ddaf
1 file changed
+69
backend/app/models/agents.py
new
+69
@@ -0,0 +1,69 @@
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 for agent metadata which stores the agent ID, IP address, hostname, OS, last seen timestamp,
11
+# and boolean for critical assest.
12
+# Path: backend\app\models.py
13
+class AgentMetadata(db.Model):
14
+ id = db.Column(db.Integer, primary_key=True)
15
+ agent_id = db.Column(db.String(100))
16
+ ip_address = db.Column(db.String(100))
17
+ os = db.Column(db.String(100))
18
+ hostname = db.Column(db.String(100))
19
+ critical_asset = db.Column(db.Boolean, default=False)
20
+ last_seen = db.Column(db.DateTime)
21
+
22
+ def __init__(self, agent_id, ip_address, os, hostname, critical_asset, last_seen):
23
+ self.agent_id = agent_id
24
+ self.ip_address = ip_address
25
+ self.os = os
26
+ self.hostname = hostname
27
+ self.critical_asset = critical_asset
28
+ self.last_seen = last_seen
29
+
30
+ def __repr__(self):
31
+ return f"<AgentMetadata {self.agent_id}>"
32
+
33
+ def mark_as_critical(self):
34
+ """
35
+ Marks the agent as a critical asset.
36
+ """
37
+ self.critical_asset = True
38
+ db.session.commit()
39
+
40
+ def mark_as_non_critical(self):
41
+ """
42
+ Marks the agent as a non-critical asset.
43
+ """
44
+ self.critical_asset = False
45
+ db.session.commit()
46
+
47
+ def commit_wazuh_agent_to_db(self):
48
+ """
49
+ Commits the agent to the database.
50
+ """
51
+ db.session.add(self)
52
+ db.session.commit()
53
+
54
+
55
+class AgentMetadataSchema(ma.Schema):
56
+ class Meta:
57
+ fields = (
58
+ "id",
59
+ "agent_id",
60
+ "ip_address",
61
+ "os",
62
+ "hostname",
63
+ "critical_asset",
64
+ "last_seen",
65
+ )
66
+
67
+
68
+agent_metadata_schema = AgentMetadataSchema()
69
+agent_metadatas_schema = AgentMetadataSchema(many=True)