Create disabled_rule.py
taylor_socfortress committed
Jul 10, 2023 at 16:46 UTC
5e175a4b85e8b65ab5c6ab254bd9801f4d55b03b
1 file changed
+179
backend/app/services/WazuhManager/disabled_rule.py
new
+179
@@ -0,0 +1,179 @@
1
+from typing import Dict, Optional, Union, List, Any, Tuple
2
+from loguru import logger
3
+from app.services.WazuhManager.universal import UniversalService
4
+import requests
5
+from app.models.rules import DisabledRules
6
+from app.models.connectors import connector_factory, Connector
7
+from app import db
8
+import xmltodict
9
+import xml.etree.ElementTree as ET
10
+import json
11
+
12
+class WazuhHttpRequests:
13
+ """
14
+ Class to handle HTTP requests to the Wazuh API.
15
+ """
16
+ def __init__(self, connector_url: str, wazuh_auth_token: str) -> None:
17
+ """
18
+ Args:
19
+ connector_url (str): The URL of the Wazuh Manager.
20
+ wazuh_auth_token (str): The Wazuh API authentication token.
21
+ """
22
+ self.connector_url = connector_url
23
+ self.wazuh_auth_token = wazuh_auth_token
24
+ self.headers = {"Authorization": f"Bearer {wazuh_auth_token}"}
25
+
26
+ def get_request(self, endpoint: str, params: Optional[Dict[str, str]] = None) -> Dict[str, Union[str, bool]]:
27
+ """
28
+ Function to handle GET requests.
29
+
30
+ Args:
31
+ endpoint (str): The endpoint to make a GET request to.
32
+ params (Optional[Dict[str, str]]): Any parameters to pass in the GET request.
33
+
34
+ Returns:
35
+ Dict[str, Union[str, bool]]: A dictionary with the requested data or error message.
36
+ """
37
+ try:
38
+ response = requests.get(
39
+ f"{self.connector_url}/{endpoint}",
40
+ headers=self.headers,
41
+ params=params,
42
+ verify=False,
43
+ )
44
+ response.raise_for_status()
45
+ return {"data": response.json(), "success": True}
46
+
47
+ except Exception as e:
48
+ logger.error(f"GET request to {endpoint} failed: {e}")
49
+ return {"message": f"GET request to {endpoint} failed: {e}", "success": False}
50
+
51
+ def put_request(self, endpoint: str, data: str, params: Optional[Dict[str, str]] = None) -> Dict[str, bool]:
52
+ """
53
+ Function to handle PUT requests.
54
+
55
+ Args:
56
+ endpoint (str): The endpoint to make a PUT request to.
57
+ data (str): Data to be updated on the PUT request.
58
+ params (Optional[Dict[str, str]]): Any parameters to pass in the PUT request.
59
+
60
+ Returns:
61
+ Dict[str, bool]: A dictionary indicating the success of the operation.
62
+ """
63
+ try:
64
+ headers = self.headers.copy()
65
+ headers.update({"Content-Type": "application/octet-stream"})
66
+
67
+ response = requests.put(
68
+ f"{self.connector_url}/{endpoint}",
69
+ headers=headers,
70
+ params=params,
71
+ data=data,
72
+ verify=False,
73
+ )
74
+ response.raise_for_status()
75
+ return {"message": f"Successfully updated {endpoint}", "success": True}
76
+
77
+ except Exception as e:
78
+ logger.error(f"Failed to update {endpoint}: {e}")
79
+ return {"message": f"Failed to update {endpoint}: {e}", "success": False}
80
+
81
+
82
+class DisableRuleService:
83
+ """
84
+ A service class that encapsulates the logic for handling rule disabling related operations in Wazuh Manager.
85
+ """
86
+ def __init__(self, universal_service: UniversalService) -> None:
87
+ """
88
+ Args:
89
+ universal_service (UniversalService): The UniversalService instance to use.
90
+ """
91
+ self.universal_service = universal_service
92
+ self.auth_token = universal_service.get_auth_token()
93
+ self.wazuh_http_requests = WazuhHttpRequests(self.universal_service.connector_url, self.auth_token)
94
+
95
+ def disable_rule(self, request: Dict[str, Union[str, int]]) -> Dict[str, Union[str, bool]]:
96
+ try:
97
+ self._validate_request(request)
98
+ rule_id = request["rule_id"]
99
+ filename = self._fetch_filename(rule_id)
100
+ file_content = self._fetch_file_content(filename)
101
+ previous_level, updated_file_content = self._set_level_1(file_content, rule_id)
102
+ xml_content = self._convert_to_xml(updated_file_content)
103
+ self._store_disabled_rule_info(rule_id, previous_level, request["reason"], request["length_of_time"])
104
+ self._upload_updated_rule(filename, xml_content)
105
+ UniversalService().restart_service()
106
+ return {"message": f"Rule {rule_id} successfully disabled in file {filename}.", "success": True}
107
+
108
+ except Exception as e:
109
+ logger.error(str(e))
110
+ return {"message": str(e), "success": False}
111
+
112
+ def _validate_request(self, request: Dict[str, Union[str, int]]):
113
+ logger.info(f"Validating disable rule request: {request}")
114
+ if "rule_id" not in request:
115
+ raise ValueError("Request missing rule_id")
116
+ if "reason" not in request:
117
+ raise ValueError("Request missing reason")
118
+ if "length_of_time" not in request:
119
+ raise ValueError("Request missing length_of_time")
120
+ request["length_of_time"] = int(request["length_of_time"])
121
+
122
+ def _fetch_filename(self, rule_id: str) -> str:
123
+ filename_data = self.wazuh_http_requests.get_request("rules", {"rule_ids": rule_id})
124
+ if not filename_data["success"]:
125
+ raise ValueError(filename_data["message"])
126
+ return filename_data["data"]["data"]["affected_items"][0]["filename"]
127
+
128
+ def _fetch_file_content(self, filename: str) -> Union[Dict[str, str], List[Dict[str, str]]]:
129
+ file_content_data = self.wazuh_http_requests.get_request(f"rules/files/{filename}")
130
+ if not file_content_data["success"]:
131
+ raise ValueError(file_content_data["message"])
132
+ return file_content_data["data"]["data"]["affected_items"][0]["group"]
133
+
134
+ def _set_level_1(self, file_content: Union[Dict[str, str], List[Dict[str, str]]], rule_id: str) -> Tuple[str, Union[Dict[str, str], List[Dict[str, str]]]]:
135
+ logger.info(f"Setting rule {rule_id} level to 1 for file_content: {file_content}")
136
+ previous_level = None
137
+ if isinstance(file_content, dict):
138
+ file_content = [file_content]
139
+ for group_block in file_content:
140
+ rule_block = group_block.get("rule", None)
141
+ if not rule_block:
142
+ continue
143
+ if isinstance(rule_block, dict):
144
+ rule_block = [rule_block]
145
+ for rule in rule_block:
146
+ if rule["@id"] == rule_id:
147
+ previous_level = rule["@level"]
148
+ rule["@level"] = "1"
149
+ break
150
+ return previous_level, file_content
151
+
152
+ def _convert_to_xml(self, updated_file_content: Union[Dict[str, str], List[Dict[str, str]]]) -> str:
153
+ logger.info(f"Received updated_file_content: {updated_file_content}")
154
+ xml_content_list = []
155
+ for group in updated_file_content:
156
+ xml_dict = {"group": group}
157
+ xml_content = xmltodict.unparse(xml_dict, pretty=True)
158
+ xml_content = xml_content.replace('<?xml version="1.0" encoding="utf-8"?>', "")
159
+ xml_content_list.append(xml_content)
160
+ xml_content = "\n".join(xml_content_list)
161
+ xml_content = xml_content.strip()
162
+ return xml_content
163
+
164
+ def _store_disabled_rule_info(self, rule_id: str, previous_level: str, reason: str, length_of_time: str):
165
+ disabled_rule = DisabledRules(
166
+ rule_id=rule_id,
167
+ previous_level=previous_level,
168
+ new_level="1",
169
+ reason_for_disabling=reason,
170
+ length_of_time=length_of_time,
171
+ )
172
+ db.session.add(disabled_rule)
173
+ db.session.commit()
174
+
175
+ def _upload_updated_rule(self, filename: str, xml_content: str):
176
+ response = self.wazuh_http_requests.put_request(f"rules/files/{filename}", xml_content, {"overwrite": "true"})
177
+ if not response["success"]:
178
+ raise ValueError(response["message"])
179
+