@cryptotaxi247 / CoPilot / commits / f6d3a874

added docstrings and typehints to models (#3)

taylor_socfortress committed Jul 11, 2023 at 07:02 UTC f6d3a87423cdc26aa6174633c8f572fc2fda9db8
9 files changed +444 -160
backend/app/models/agents.py
+55 -18
@@ -1,25 +1,49 @@
1 -# from datetime import datetime
1 +from datetime import datetime
2
3 -# from loguru import logger
4 -# from sqlalchemy.dialects.postgresql import JSONB # Add this line
3 +from sqlalchemy import Boolean
4 +from sqlalchemy import Column
5 +from sqlalchemy import DateTime
6 +from sqlalchemy import Integer
7 +from sqlalchemy import String
8
9 from app import db
10 from app import ma
11
12
10 -# Class for agent metadata which stores the agent ID, IP address, hostname, OS, last seen timestamp,
11 -# and boolean for critical assest.
13 # Path: backend\app\models.py
14 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):
15 + """
16 + Class for agent metadata which stores the agent ID, IP address, hostname, OS, last seen timestamp,
17 + and boolean for critical asset. This class inherits from SQLAlchemy's Model class.
18 + """
19 +
20 + id: Column[Integer] = db.Column(db.Integer, primary_key=True)
21 + agent_id: Column[String] = db.Column(db.String(100))
22 + ip_address: Column[String] = db.Column(db.String(100))
23 + os: Column[String] = db.Column(db.String(100))
24 + hostname: Column[String] = db.Column(db.String(100))
25 + critical_asset: Column[Boolean] = db.Column(db.Boolean, default=False)
26 + last_seen: Column[DateTime] = db.Column(db.DateTime)
27 +
28 + def __init__(
29 + self,
30 + agent_id: str,
31 + ip_address: str,
32 + os: str,
33 + hostname: str,
34 + critical_asset: bool,
35 + last_seen: datetime,
36 + ):
37 + """
38 + Initialize a new instance of the AgentMetadata class.
39 +
40 + :param agent_id: Unique ID for the agent.
41 + :param ip_address: IP address of the agent.
42 + :param os: Operating system of the agent.
43 + :param hostname: Hostname of the agent.
44 + :param critical_asset: Boolean value indicating if the agent is a critical asset.
45 + :param last_seen: Timestamp of when the agent was last seen.
46 + """
47 self.agent_id = agent_id
48 self.ip_address = ip_address
49 self.os = os
@@ -27,7 +51,12 @@ class AgentMetadata(db.Model):
51 self.critical_asset = critical_asset
52 self.last_seen = last_seen
53
30 - def __repr__(self):
54 + def __repr__(self) -> str:
55 + """
56 + Returns a string representation of the AgentMetadata instance.
57 +
58 + :return: A string representation of the agent ID.
59 + """
60 return f"<AgentMetadata {self.agent_id}>"
61
62 def mark_as_critical(self):
@@ -53,8 +82,16 @@ class AgentMetadata(db.Model):
82
83
84 class AgentMetadataSchema(ma.Schema):
85 + """
86 + Schema for serializing and deserializing instances of the AgentMetadata class.
87 + """
88 +
89 class Meta:
57 - fields = (
90 + """
91 + Meta class defines the fields to be serialized/deserialized.
92 + """
93 +
94 + fields: tuple = (
95 "id",
96 "agent_id",
97 "ip_address",
@@ -65,5 +102,5 @@ class AgentMetadataSchema(ma.Schema):
102 )
103
104
68 -agent_metadata_schema = AgentMetadataSchema()
69 -agent_metadatas_schema = AgentMetadataSchema(many=True)
105 +agent_metadata_schema: AgentMetadataSchema = AgentMetadataSchema()
106 +agent_metadatas_schema: AgentMetadataSchema = AgentMetadataSchema(many=True)
backend/app/models/artifacts.py
+40 -13
@@ -1,29 +1,56 @@
1 -from datetime import datetime
1 +from sqlalchemy import Column
2 +from sqlalchemy import Integer
3 +from sqlalchemy import String
4 +from sqlalchemy.dialects.postgresql import TEXT # Add this line
5
6 from app import db
7 from app import ma
5 -from sqlalchemy.dialects.postgresql import TEXT # Add this line
8
7 -# Class for artifacts collected which stores the artifact name, artificat results (json), hostname
9 +
10 # Path: backend\app\models.py
11 class Artifact(db.Model):
10 - id = db.Column(db.Integer, primary_key=True)
11 - artifact_name = db.Column(db.String(100))
12 - artifact_results = db.Column(TEXT)
13 - hostname = db.Column(db.String(100))
14 -
15 - def __init__(self, artifact_name, artifact_results, hostname):
12 + """
13 + Class for artifacts collected which stores the artifact name, artifact results (json), and hostname.
14 + This class inherits from SQLAlchemy's Model class.
15 + """
16 +
17 + id: Column[Integer] = db.Column(db.Integer, primary_key=True)
18 + artifact_name: Column[String] = db.Column(db.String(100))
19 + artifact_results: Column[TEXT] = db.Column(TEXT)
20 + hostname: Column[String] = db.Column(db.String(100))
21 +
22 + def __init__(self, artifact_name: str, artifact_results: str, hostname: str):
23 + """
24 + Initialize a new instance of the Artifact class.
25 +
26 + :param artifact_name: The name of the artifact.
27 + :param artifact_results: The results of the artifact, stored as a JSON string.
28 + :param hostname: The hostname where the artifact was collected.
29 + """
30 self.artifact_name = artifact_name
31 self.artifact_results = artifact_results
32 self.hostname = hostname
33
20 - def __repr__(self):
34 + def __repr__(self) -> str:
35 + """
36 + Returns a string representation of the Artifact instance.
37 +
38 + :return: A string representation of the artifact name.
39 + """
40 return f"<Artifact {self.artifact_name}>"
41
42
43 class ArtifactSchema(ma.Schema):
44 + """
45 + Schema for serializing and deserializing instances of the Artifact class.
46 + """
47 +
48 class Meta:
26 - fields = (
49 + """
50 + Meta class defines the fields to be serialized/deserialized.
51 + """
52 +
53 + fields: tuple = (
54 "id",
55 "artifact_name",
56 "artifact_results",
@@ -31,5 +58,5 @@ class ArtifactSchema(ma.Schema):
58 )
59
60
34 -artifact_schema = ArtifactSchema()
35 -artifacts_schema = ArtifactSchema(many=True)
61 +artifact_schema: ArtifactSchema = ArtifactSchema()
62 +artifacts_schema: ArtifactSchema = ArtifactSchema(many=True)
backend/app/models/cases.py
+39 -12
@@ -1,28 +1,55 @@
1 -from datetime import datetime
1 +from sqlalchemy import Column
2 +from sqlalchemy import Integer
3 +from sqlalchemy import String
4
5 from app import db
6 from app import ma
7
6 -# Class for cases which stores the case ID, case name, list of agents
8 +
9 # Path: backend\app\models.py
10 class Case(db.Model):
9 - id = db.Column(db.Integer, primary_key=True)
10 - case_id = db.Column(db.Integer)
11 - case_name = db.Column(db.String(100))
12 - agents = db.Column(db.String(1000))
13 -
14 - def __init__(self, case_id, case_name, agents):
11 + """
12 + Class for cases which stores the case ID, case name, and a list of agents.
13 + This class inherits from SQLAlchemy's Model class.
14 + """
15 +
16 + id: Column[Integer] = db.Column(db.Integer, primary_key=True)
17 + case_id: Column[Integer] = db.Column(db.Integer)
18 + case_name: Column[String] = db.Column(db.String(100))
19 + agents: Column[String] = db.Column(db.String(1000))
20 +
21 + def __init__(self, case_id: int, case_name: str, agents: str):
22 + """
23 + Initialize a new instance of the Case class.
24 +
25 + :param case_id: The ID of the case.
26 + :param case_name: The name of the case.
27 + :param agents: A comma-separated string of agents associated with the case.
28 + """
29 self.case_id = case_id
30 self.case_name = case_name
31 self.agents = agents
32
19 - def __repr__(self):
33 + def __repr__(self) -> str:
34 + """
35 + Returns a string representation of the Case instance.
36 +
37 + :return: A string representation of the case ID.
38 + """
39 return f"<Case {self.case_id}>"
40
41
42 class CaseSchema(ma.Schema):
43 + """
44 + Schema for serializing and deserializing instances of the Case class.
45 + """
46 +
47 class Meta:
25 - fields = (
48 + """
49 + Meta class defines the fields to be serialized/deserialized.
50 + """
51 +
52 + fields: tuple = (
53 "id",
54 "case_id",
55 "case_name",
@@ -30,5 +57,5 @@ class CaseSchema(ma.Schema):
57 )
58
59
33 -case_schema = CaseSchema()
34 -cases_schema = CaseSchema(many=True)
60 +case_schema: CaseSchema = CaseSchema()
61 +cases_schema: CaseSchema = CaseSchema(many=True)
backend/app/models/connectors.py
+21 -21
@@ -3,6 +3,8 @@ import json
3 from abc import ABC
4 from abc import abstractmethod
5 from dataclasses import dataclass
6 +from typing import Any
7 +from typing import Dict
8
9 import grpc
10 import pika
@@ -17,10 +19,8 @@ from sqlalchemy.orm.exc import NoResultFound
19
20 from app.models.models import Connectors
21
20 -# from werkzeug.utils import secure_filename
22
22 -
23 -def dynamic_import(module_name, class_name):
23 +def dynamic_import(module_name: str, class_name: str) -> Any:
24 """
25 This function dynamically imports a module and returns a specific class from it.
26
@@ -43,10 +43,10 @@ class Connector(ABC):
43 :param attributes: A dictionary of attributes necessary for the connector to connect to the service or system.
44 """
45
46 - attributes: dict
46 + attributes: Dict[str, Any]
47
48 @abstractmethod
49 - def verify_connection(self):
49 + def verify_connection(self) -> Dict[str, Any]:
50 """
51 This abstract method should be implemented by all subclasses of Connector. It is meant to verify the
52 connection to the service or system the connector is designed to connect to.
@@ -56,7 +56,7 @@ class Connector(ABC):
56 pass
57
58 @staticmethod
59 - def get_connector_info_from_db(connector_name):
59 + def get_connector_info_from_db(connector_name: str) -> Dict[str, Any]:
60 """
61 This method retrieves connector information from the database.
62
@@ -89,10 +89,10 @@ class WazuhIndexerConnector(Connector):
89 :param connector_name: A string that specifies the name of the connector.
90 """
91
92 - def __init__(self, connector_name):
92 + def __init__(self, connector_name: str):
93 super().__init__(attributes=self.get_connector_info_from_db(connector_name))
94
95 - def verify_connection(self):
95 + def verify_connection(self) -> Dict[str, Any]:
96 """
97 This method verifies the connection to the Wazuh indexer service.
98
@@ -131,10 +131,10 @@ class GraylogConnector(Connector):
131 :param connector_name: A string that specifies the name of the connector.
132 """
133
134 - def __init__(self, connector_name):
134 + def __init__(self, connector_name: str):
135 super().__init__(attributes=self.get_connector_info_from_db(connector_name))
136
137 - def verify_connection(self):
137 + def verify_connection(self) -> Dict[str, Any]:
138 """
139 Verifies the connection to Graylog service.
140
@@ -177,10 +177,10 @@ class WazuhManagerConnector(Connector):
177 :param connector_name: A string that specifies the name of the connector.
178 """
179
180 - def __init__(self, connector_name):
180 + def __init__(self, connector_name: str):
181 super().__init__(attributes=self.get_connector_info_from_db(connector_name))
182
183 - def verify_connection(self):
183 + def verify_connection(self) -> Dict[str, Any]:
184 """
185 Verifies the connection to Wazuh manager service.
186
@@ -215,7 +215,7 @@ class WazuhManagerConnector(Connector):
215 )
216 return {"connectionSuccessful": False, "authToken": None}
217
218 - def get_auth_token(self):
218 + def get_auth_token(self) -> str:
219 """
220 Returns the authentication token for the Wazuh manager service.
221
@@ -232,10 +232,10 @@ class ShuffleConnector(Connector):
232 :param connector_name: A string that specifies the name of the connector.
233 """
234
235 - def __init__(self, connector_name):
235 + def __init__(self, connector_name: str):
236 super().__init__(attributes=self.get_connector_info_from_db(connector_name))
237
238 - def verify_connection(self):
238 + def verify_connection(self) -> Dict[str, Any]:
239 """
240 Verifies the connection to Shuffle service.
241
@@ -278,10 +278,10 @@ class DfirIrisConnector(Connector):
278 :param connector_name: A string that specifies the name of the connector.
279 """
280
281 - def __init__(self, connector_name):
281 + def __init__(self, connector_name: str):
282 super().__init__(attributes=self.get_connector_info_from_db(connector_name))
283
284 - def verify_connection(self):
284 + def verify_connection(self) -> Dict[str, Any]:
285 """
286 Verifies the connection to DFIR IRIS service.
287
@@ -326,10 +326,10 @@ class VelociraptorConnector(Connector):
326 connector_name (str): The name of the connector.
327 """
328
329 - def __init__(self, connector_name):
329 + def __init__(self, connector_name: str):
330 super().__init__(attributes=self.get_connector_info_from_db(connector_name))
331
332 - def verify_connection(self):
332 + def verify_connection(self) -> Dict[str, Any]:
333 """
334 Verifies the connection to Velociraptor service.
335
@@ -391,10 +391,10 @@ class RabbitMQConnector(Connector):
391 connector_name (str): The name of the connector.
392 """
393
394 - def __init__(self, connector_name):
394 + def __init__(self, connector_name: str):
395 super().__init__(attributes=self.get_connector_info_from_db(connector_name))
396
397 - def verify_connection(self):
397 + def verify_connection(self) -> Dict[str, Any]:
398 """
399 Verifies the connection to RabbitMQ service.
400 """
backend/app/models/graylog.py
+60 -22
@@ -1,32 +1,53 @@
1 from datetime import datetime
2
3 +from sqlalchemy import Column
4 +from sqlalchemy import DateTime
5 +from sqlalchemy import Float
6 +from sqlalchemy import Integer
7 +
8 from app import db
9 from app import ma
10
6 -# Class for Graylog allocation which stores throughput metrics
7 -# Generate timestamp for each entry and invoke every 5 minutes.
11 +
12 # Path: backend\app\models.py
13 class GraylogMetricsAllocation(db.Model):
10 - id = db.Column(db.Integer, primary_key=True)
11 - input_usage = db.Column(db.Float)
12 - output_usage = db.Column(db.Float)
13 - processor_usage = db.Column(db.Float)
14 - input_1_sec_rate = db.Column(db.Float)
15 - output_1_sec_rate = db.Column(db.Float)
16 - total_input = db.Column(db.Float)
17 - total_output = db.Column(db.Float)
18 - timestamp = db.Column(db.DateTime, default=datetime.utcnow)
14 + """
15 + Class for Graylog metrics allocation which stores throughput metrics.
16 + The timestamp is generated for each entry and it is invoked every 5 minutes.
17 + This class inherits from SQLAlchemy's Model class.
18 + """
19 +
20 + id: Column[Integer] = db.Column(db.Integer, primary_key=True)
21 + input_usage: Column[Float] = db.Column(db.Float)
22 + output_usage: Column[Float] = db.Column(db.Float)
23 + processor_usage: Column[Float] = db.Column(db.Float)
24 + input_1_sec_rate: Column[Float] = db.Column(db.Float)
25 + output_1_sec_rate: Column[Float] = db.Column(db.Float)
26 + total_input: Column[Float] = db.Column(db.Float)
27 + total_output: Column[Float] = db.Column(db.Float)
28 + timestamp: Column[DateTime] = db.Column(db.DateTime, default=datetime.utcnow)
29
30 def __init__(
31 self,
22 - input_usage,
23 - output_usage,
24 - processor_usage,
25 - input_1_sec_rate,
26 - output_1_sec_rate,
27 - total_input,
28 - total_output,
32 + input_usage: float,
33 + output_usage: float,
34 + processor_usage: float,
35 + input_1_sec_rate: float,
36 + output_1_sec_rate: float,
37 + total_input: float,
38 + total_output: float,
39 ):
40 + """
41 + Initialize a new instance of the GraylogMetricsAllocation class.
42 +
43 + :param input_usage: The input usage value.
44 + :param output_usage: The output usage value.
45 + :param processor_usage: The processor usage value.
46 + :param input_1_sec_rate: The input per second rate.
47 + :param output_1_sec_rate: The output per second rate.
48 + :param total_input: The total input value.
49 + :param total_output: The total output value.
50 + """
51 self.input_usage = input_usage
52 self.output_usage = output_usage
53 self.processor_usage = processor_usage
@@ -35,13 +56,26 @@ class GraylogMetricsAllocation(db.Model):
56 self.total_input = total_input
57 self.total_output = total_output
58
38 - def __repr__(self):
59 + def __repr__(self) -> str:
60 + """
61 + Returns a string representation of the GraylogMetricsAllocation instance.
62 +
63 + :return: A string representation of the instance's id.
64 + """
65 return f"<GraylogMetricsAllocation {self.id}>"
66
67
68 class GraylogMetricsAllocationSchema(ma.Schema):
69 + """
70 + Schema for serializing and deserializing instances of the GraylogMetricsAllocation class.
71 + """
72 +
73 class Meta:
44 - fields = (
74 + """
75 + Meta class defines the fields to be serialized/deserialized.
76 + """
77 +
78 + fields: tuple = (
79 "id",
80 "input_usage",
81 "output_usage",
@@ -54,5 +88,9 @@ class GraylogMetricsAllocationSchema(ma.Schema):
88 )
89
90
57 -graylog_metrics_allocation_schema = GraylogMetricsAllocationSchema()
58 -graylog_metrics_allocations_schema = GraylogMetricsAllocationSchema(many=True)
91 +graylog_metrics_allocation_schema: GraylogMetricsAllocationSchema = (
92 + GraylogMetricsAllocationSchema()
93 +)
94 +graylog_metrics_allocations_schema: GraylogMetricsAllocationSchema = (
95 + GraylogMetricsAllocationSchema(many=True)
96 +)
backend/app/models/models.py
+108 -34
@@ -1,30 +1,66 @@
1 from datetime import datetime
2
3 from loguru import logger
4 +from sqlalchemy import Boolean
5 +from sqlalchemy import Column
6 +from sqlalchemy import DateTime
7 +from sqlalchemy import Integer
8 +from sqlalchemy import String
9
10 from app import db
11 from app import ma
12
13
14 class ConnectorsAvailable(db.Model):
10 - id = db.Column(db.Integer, primary_key=True)
11 - connector_name = db.Column(db.String(100), unique=True)
12 - connector_description = db.Column(db.String(100))
13 - connector_supports = db.Column(db.String(100))
14 - connector_configured = db.Column(db.Boolean, default=False)
15 - connector_verified = db.Column(db.Boolean, default=False)
16 -
17 - def __init__(self, connector_name, connector_supports):
15 + """
16 + Class representing the available connectors in the application.
17 + This class inherits from SQLAlchemy's Model class.
18 +
19 + :ivar id: Unique integer ID of the connector.
20 + :ivar connector_name: Name of the connector.
21 + :ivar connector_description: Description of the connector.
22 + :ivar connector_supports: Information on what the connector supports.
23 + :ivar connector_configured: Boolean indicating whether the connector is configured or not.
24 + :ivar connector_verified: Boolean indicating whether the connector is verified or not.
25 + """
26 +
27 + id: Column[Integer] = db.Column(db.Integer, primary_key=True)
28 + connector_name: Column[String] = db.Column(db.String(100), unique=True)
29 + connector_description: Column[String] = db.Column(db.String(100))
30 + connector_supports: Column[String] = db.Column(db.String(100))
31 + connector_configured: Column[Boolean] = db.Column(db.Boolean, default=False)
32 + connector_verified: Column[Boolean] = db.Column(db.Boolean, default=False)
33 +
34 + def __init__(self, connector_name: str, connector_supports: str):
35 + """
36 + Initialize a new instance of the ConnectorsAvailable class.
37 +
38 + :param connector_name: The name of the connector.
39 + :param connector_supports: Information on what the connector supports.
40 + """
41 self.connector_name = connector_name
42 self.connector_supports = connector_supports
43
21 - def __repr__(self):
44 + def __repr__(self) -> str:
45 + """
46 + Returns a string representation of the ConnectorsAvailable instance.
47 +
48 + :return: A string representation of the connector name.
49 + """
50 return f"<ConnectorsAvailble {self.connector_name}>"
51
52
53 class ConnectorsAvailableSchema(ma.Schema):
54 + """
55 + Schema for serializing and deserializing instances of the ConnectorsAvailable class.
56 + """
57 +
58 class Meta:
27 - fields = (
59 + """
60 + Meta class defines the fields to be serialized/deserialized.
61 + """
62 +
63 + fields: tuple = (
64 "id",
65 "connector_name",
66 "connector_description",
@@ -34,38 +70,63 @@ class ConnectorsAvailableSchema(ma.Schema):
70 )
71
72
37 -connector_available_schema = ConnectorsAvailableSchema()
38 -connectors_available_schema = ConnectorsAvailableSchema(many=True)
73 +connector_available_schema: ConnectorsAvailableSchema = ConnectorsAvailableSchema()
74 +connectors_available_schema: ConnectorsAvailableSchema = ConnectorsAvailableSchema(
75 + many=True,
76 +)
77
78
41 -# Class for the connector which will store the endpoint url, connector name, connector type, connector last updated,
42 -# username and password
79 class Connectors(db.Model):
44 - id = db.Column(db.Integer, primary_key=True)
45 - connector_name = db.Column(db.String(100), unique=True)
46 - connector_type = db.Column(db.String(100))
47 - connector_url = db.Column(db.String(100))
48 - connector_last_updated = db.Column(db.DateTime, default=datetime.utcnow)
49 - connector_username = db.Column(db.String(100))
50 - connector_password = db.Column(db.String(100))
51 - connector_api_key = db.Column(db.String(100))
80 + """
81 + Class for the connector which will store the endpoint url, connector name, connector type, connector last updated,
82 + username and password.
83 +
84 + :ivar id: Unique integer ID of the connector.
85 + :ivar connector_name: Name of the connector.
86 + :ivar connector_type: Type of the connector.
87 + :ivar connector_url: URL of the connector.
88 + :ivar connector_last_updated: Timestamp when the connector was last updated.
89 + :ivar connector_username: Username for the connector.
90 + :ivar connector_password: Password for the connector.
91 + :ivar connector_api_key: API key for the connector.
92 + """
93 +
94 + id: Column[Integer] = db.Column(db.Integer, primary_key=True)
95 + connector_name: Column[String] = db.Column(db.String(100), unique=True)
96 + connector_type: Column[String] = db.Column(db.String(100))
97 + connector_url: Column[String] = db.Column(db.String(100))
98 + connector_last_updated: Column[DateTime] = db.Column(
99 + db.DateTime, default=datetime.utcnow,
100 + )
101 + connector_username: Column[String] = db.Column(db.String(100))
102 + connector_password: Column[String] = db.Column(db.String(100))
103 + connector_api_key: Column[String] = db.Column(db.String(100))
104
105 def __init__(
106 self,
55 - connector_name,
56 - connector_type,
57 - connector_url,
58 - connector_username,
59 - connector_password,
60 - connector_api_key,
107 + connector_name: str,
108 + connector_type: str,
109 + connector_url: str,
110 + connector_username: str,
111 + connector_password: str,
112 + connector_api_key: str,
113 ):
114 + """
115 + Initialize a new instance of the Connectors class.
116 +
117 + :param connector_name: The name of the connector.
118 + :param connector_type: The type of the connector.
119 + :param connector_url: The URL of the connector.
120 + :param connector_username: The username for the connector.
121 + :param connector_password: The password for the connector.
122 + :param connector_api_key: The API key for the connector.
123 + """
124 self.connector_name = connector_name
125 self.connector_type = connector_type
126 self.connector_url = connector_url
127 self.connector_username = connector_username
128 self.connector_password = connector_password
67 - # If the `connector_name` is `shuffle` or `dfir-irs` then set the `connector_api_key`. Otherwise set it to
68 - # `None`
129 +
130 if (
131 connector_name.lower() == "shuffle"
132 or connector_name.lower() == "dfir-irs"
@@ -77,13 +138,26 @@ class Connectors(db.Model):
138 logger.info(f"Not setting the API key for {connector_name}")
139 self.connector_api_key = None
140
80 - def __repr__(self):
141 + def __repr__(self) -> str:
142 + """
143 + Returns a string representation of the Connectors instance.
144 +
145 + :return: A string representation of the connector name.
146 + """
147 return f"<Connectors {self.connector_name}>"
148
149
150 class ConnectorsSchema(ma.Schema):
151 + """
152 + Schema for serializing and deserializing instances of the Connectors class.
153 + """
154 +
155 class Meta:
86 - fields = (
156 + """
157 + Meta class defines the fields to be serialized/deserialized.
158 + """
159 +
160 + fields: tuple = (
161 "id",
162 "connector_name",
163 "connector_type",
@@ -95,8 +169,8 @@ class ConnectorsSchema(ma.Schema):
169 )
170
171
98 -connector_schema = ConnectorsSchema()
99 -connectors_schema = ConnectorsSchema(many=True)
172 +connector_schema: ConnectorsSchema = ConnectorsSchema()
173 +connectors_schema: ConnectorsSchema = ConnectorsSchema(many=True)
174
175
176 # Class for the disabled rule IDs which will store the rule ID, previous configuration, new configuration, reason for
backend/app/models/rules.py
+57 -21
@@ -1,45 +1,81 @@
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
6 -# from loguru import logger
7 -# from sqlalchemy.dialects.postgresql import JSONB # Add this line
8 -
11
10 -# Class for the disabled rule IDs which will store the rule ID, previous configuration, new configuration, reason for
11 -# disabling, date disabled, and the length of time the rule will be disabled for
12 # Path: backend\app\rules.py
13 class DisabledRules(db.Model):
14 - id = db.Column(db.Integer, primary_key=True)
15 - rule_id = db.Column(db.String(100))
16 - previous_level = db.Column(db.String(1000))
17 - new_level = db.Column(db.String(1000))
18 - reason_for_disabling = db.Column(db.String(100))
19 - date_disabled = db.Column(db.DateTime, default=datetime.utcnow)
20 - length_of_time = db.Column(db.Integer)
14 + """
15 + Class for disabled rules which stores the rule ID, previous configuration, new configuration, reason for
16 + disabling, date disabled, and the length of time the rule will be disabled for.
17 + This class inherits from SQLAlchemy's Model class.
18 +
19 + :ivar id: Unique integer ID of the rule.
20 + :ivar rule_id: ID of the rule.
21 + :ivar previous_level: Previous level configuration of the rule.
22 + :ivar new_level: New level configuration of the rule.
23 + :ivar reason_for_disabling: Reason for disabling the rule.
24 + :ivar date_disabled: Date when the rule was disabled.
25 + :ivar length_of_time: Length of time the rule will be disabled for.
26 + """
27 +
28 + id: Column[Integer] = db.Column(db.Integer, primary_key=True)
29 + rule_id: Column[String] = db.Column(db.String(100))
30 + previous_level: Column[String] = db.Column(db.String(1000))
31 + new_level: Column[String] = db.Column(db.String(1000))
32 + reason_for_disabling: Column[String] = db.Column(db.String(100))
33 + date_disabled: Column[DateTime] = db.Column(db.DateTime, default=datetime.utcnow)
34 + length_of_time: Column[Integer] = db.Column(db.Integer)
35
36 def __init__(
37 self,
24 - rule_id,
25 - previous_level,
26 - new_level,
27 - reason_for_disabling,
28 - length_of_time,
38 + rule_id: str,
39 + previous_level: str,
40 + new_level: str,
41 + reason_for_disabling: str,
42 + length_of_time: int,
43 ):
44 + """
45 + Initialize a new instance of the DisabledRules class.
46 +
47 + :param rule_id: The ID of the rule.
48 + :param previous_level: The previous level configuration of the rule.
49 + :param new_level: The new level configuration of the rule.
50 + :param reason_for_disabling: The reason for disabling the rule.
51 + :param length_of_time: The length of time the rule will be disabled for.
52 + """
53 self.rule_id = rule_id
54 self.previous_level = previous_level
55 self.new_level = new_level
56 self.reason_for_disabling = reason_for_disabling
57 self.length_of_time = length_of_time
58
36 - def __repr__(self):
59 + def __repr__(self) -> str:
60 + """
61 + Returns a string representation of the DisabledRules instance.
62 +
63 + :return: A string representation of the rule ID.
64 + """
65 return f"<DisabledRules {self.rule_id}>"
66
67
68 class DisabledRulesSchema(ma.Schema):
69 + """
70 + Schema for serializing and deserializing instances of the DisabledRules class.
71 + """
72 +
73 class Meta:
42 - fields = (
74 + """
75 + Meta class defines the fields to be serialized/deserialized.
76 + """
77 +
78 + fields: tuple = (
79 "id",
80 "rule_id",
81 "previous_level",
@@ -50,5 +86,5 @@ class DisabledRulesSchema(ma.Schema):
86 )
87
88
53 -disabled_rule_schema = DisabledRulesSchema()
54 -disabled_rules_schema = DisabledRulesSchema(many=True)
89 +disabled_rule_schema: DisabledRulesSchema = DisabledRulesSchema()
90 +disabled_rules_schema: DisabledRulesSchema = DisabledRulesSchema(many=True)
backend/app/models/wazuh_indexer.py
+63 -18
@@ -1,41 +1,82 @@
1 from datetime import datetime
2
3 +from sqlalchemy import Column
4 +from sqlalchemy import DateTime
5 +from sqlalchemy import Float
6 +from sqlalchemy import Integer
7 +from sqlalchemy import String
8 +
9 from app import db
10 from app import ma
11
6 -# Class for Wazuh Indexer allocation which stores disk stats and the host.
7 -# Generate timestamp for each entry and invoke every 5 minutes.
12 +
13 # Path: backend\app\models.py
14 class WazuhIndexerAllocation(db.Model):
10 - id = db.Column(db.Integer, primary_key=True)
11 - node = db.Column(db.String(100))
12 - disk_used = db.Column(db.Float)
13 - disk_available = db.Column(db.Float)
14 - disk_total = db.Column(db.Float)
15 - disk_percent = db.Column(db.Float)
16 - timestamp = db.Column(db.DateTime, default=datetime.utcnow)
15 + """
16 + Class for Wazuh indexer allocation which stores disk stats and the host.
17 + The timestamp is generated for each entry and it is invoked every 5 minutes.
18 + This class inherits from SQLAlchemy's Model class.
19 +
20 + :ivar id: Unique integer ID for the allocation.
21 + :ivar node: The node or host for the allocation.
22 + :ivar disk_used: The amount of disk used.
23 + :ivar disk_available: The amount of disk available.
24 + :ivar disk_total: The total amount of disk.
25 + :ivar disk_percent: The percent of disk used.
26 + :ivar timestamp: The timestamp when the allocation was created.
27 + """
28 +
29 + id: Column[Integer] = db.Column(db.Integer, primary_key=True)
30 + node: Column[String] = db.Column(db.String(100))
31 + disk_used: Column[Float] = db.Column(db.Float)
32 + disk_available: Column[Float] = db.Column(db.Float)
33 + disk_total: Column[Float] = db.Column(db.Float)
34 + disk_percent: Column[Float] = db.Column(db.Float)
35 + timestamp: Column[DateTime] = db.Column(db.DateTime, default=datetime.utcnow)
36
37 def __init__(
38 self,
20 - node,
21 - disk_used,
22 - disk_available,
23 - disk_total,
24 - disk_percent,
39 + node: str,
40 + disk_used: float,
41 + disk_available: float,
42 + disk_total: float,
43 + disk_percent: float,
44 ):
45 + """
46 + Initialize a new instance of the WazuhIndexerAllocation class.
47 +
48 + :param node: The node or host for the allocation.
49 + :param disk_used: The amount of disk used.
50 + :param disk_available: The amount of disk available.
51 + :param disk_total: The total amount of disk.
52 + :param disk_percent: The percent of disk used.
53 + """
54 self.node = node
55 self.disk_used = disk_used
56 self.disk_available = disk_available
57 self.disk_total = disk_total
58 self.disk_percent = disk_percent
59
32 - def __repr__(self):
60 + def __repr__(self) -> str:
61 + """
62 + Returns a string representation of the WazuhIndexerAllocation instance.
63 +
64 + :return: A string representation of the node.
65 + """
66 return f"<WazuhIndexerAllocation {self.node}>"
67
68
69 class WazuhIndexerAllocationSchema(ma.Schema):
70 + """
71 + Schema for serializing and deserializing instances of the WazuhIndexerAllocation class.
72 + """
73 +
74 class Meta:
38 - fields = (
75 + """
76 + Meta class defines the fields to be serialized/deserialized.
77 + """
78 +
79 + fields: tuple = (
80 "id",
81 "node",
82 "disk_used",
@@ -46,5 +87,9 @@ class WazuhIndexerAllocationSchema(ma.Schema):
87 )
88
89
49 -wazuh_indexer_allocation_schema = WazuhIndexerAllocationSchema()
50 -wazuh_indexer_allocations_schema = WazuhIndexerAllocationSchema(many=True)
90 +wazuh_indexer_allocation_schema: WazuhIndexerAllocationSchema = (
91 + WazuhIndexerAllocationSchema()
92 +)
93 +wazuh_indexer_allocations_schema: WazuhIndexerAllocationSchema = (
94 + WazuhIndexerAllocationSchema(many=True)
95 +)
backend/app/routes/agents.py
+1 -1
@@ -40,7 +40,7 @@ def get_agent(agent_id):
40 json: A JSON response containing the details of the agent.
41 """
42 service = AgentService()
43 - agent = service.get_agent(agent_id)
43 + agent = service.get_agent(agent_id=agent_id)
44 return agent
45
46