@cryptotaxi247 / CoPilot / commits / 70a202e1

Create enabled_rule.py

taylor_socfortress committed Jul 10, 2023 at 16:46 UTC 70a202e150a5463cbacf4c8a41c670e824bb93e8
1 file changed +290
backend/app/services/WazuhManager/enabled_rule.py new
+290
@@ -0,0 +1,290 @@
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 +
13 +class WazuhHttpRequests:
14 + """
15 + Class to handle HTTP requests to the Wazuh API.
16 + """
17 + def __init__(self, connector_url: str, wazuh_auth_token: str) -> None:
18 + """
19 + Args:
20 + connector_url (str): The URL of the Wazuh Manager.
21 + wazuh_auth_token (str): The Wazuh API authentication token.
22 + """
23 + self.connector_url = connector_url
24 + self.wazuh_auth_token = wazuh_auth_token
25 + self.headers = {"Authorization": f"Bearer {wazuh_auth_token}"}
26 +
27 + def get_request(self, endpoint: str, params: Optional[Dict[str, str]] = None) -> Dict[str, Union[str, bool]]:
28 + """
29 + Function to handle GET requests.
30 +
31 + Args:
32 + endpoint (str): The endpoint to make a GET request to.
33 + params (Optional[Dict[str, str]]): Any parameters to pass in the GET request.
34 +
35 + Returns:
36 + Dict[str, Union[str, bool]]: A dictionary with the requested data or error message.
37 + """
38 + try:
39 + logger.info(f"GET request to {endpoint}")
40 + response = requests.get(
41 + f"{self.connector_url}/{endpoint}",
42 + headers=self.headers,
43 + params=params,
44 + verify=False,
45 + )
46 + response.raise_for_status()
47 + logger.info(f"Respones: {response.json()}")
48 + return {"data": response.json(), "success": True}
49 +
50 + except Exception as e:
51 + logger.error(f"GET request to {endpoint} failed: {e}")
52 + return {"message": f"GET request to {endpoint} failed: {e}", "success": False}
53 +
54 + def put_request(self, endpoint: str, data: str, params: Optional[Dict[str, str]] = None) -> Dict[str, bool]:
55 + """
56 + Function to handle PUT requests.
57 +
58 + Args:
59 + endpoint (str): The endpoint to make a PUT request to.
60 + data (str): Data to be updated on the PUT request.
61 + params (Optional[Dict[str, str]]): Any parameters to pass in the PUT request.
62 +
63 + Returns:
64 + Dict[str, bool]: A dictionary indicating the success of the operation.
65 + """
66 + try:
67 + headers = self.headers.copy()
68 + headers.update({"Content-Type": "application/octet-stream"})
69 +
70 + response = requests.put(
71 + f"{self.connector_url}/{endpoint}",
72 + headers=headers,
73 + params=params,
74 + data=data,
75 + verify=False,
76 + )
77 + response.raise_for_status()
78 + return {"message": f"Successfully updated {endpoint}", "success": True}
79 +
80 + except Exception as e:
81 + logger.error(f"Failed to update {endpoint}: {e}")
82 + return {"message": f"Failed to update {endpoint}: {e}", "success": False}
83 +
84 +
85 +class EnableRuleService:
86 + """
87 + A service class that encapsulates the logic for handling rule enabling related operations in Wazuh Manager.
88 + """
89 + def __init__(self, universal_service: UniversalService) -> None:
90 + """
91 + Args:
92 + universal_service (UniversalService): The UniversalService instance to use.
93 + """
94 + self.universal_service = universal_service
95 + self.auth_token = universal_service.get_auth_token()
96 + self.wazuh_http_requests = WazuhHttpRequests(self.universal_service.connector_url, self.auth_token)
97 +
98 + def enable_rule(self, request: Dict[str, str]) -> Dict[str, Union[str, bool]]:
99 + """
100 + Enable a rule in the Wazuh Manager.
101 +
102 + Args:
103 + request (Dict[str, str]): The request to enable a rule in Wazuh Manager.
104 +
105 + Returns:
106 + Dict[str, Union[str, bool]]: A dictionary containing status of the operation.
107 + """
108 + try:
109 + self._validate_request(request)
110 + rule_id = request["rule_id"]
111 + filename = self._fetch_filename(rule_id)
112 + logger.info(f"Getting file content of {filename}")
113 + file_content = self._fetch_file_content(filename)
114 + previous_level = self._get_previous_level(rule_id)
115 + updated_file_content = self._set_level_previous(file_content, rule_id, previous_level)
116 + xml_content = self._json_to_xml(updated_file_content)
117 + self._delete_rule_from_db(rule_id)
118 + self._put_updated_rule(filename, xml_content)
119 + UniversalService().restart_service()
120 + return {
121 + "message": f"Rule {rule_id} successfully enabled in file {filename}.",
122 + "success": True,
123 + }
124 + except Exception as e:
125 + return {"message": str(e), "success": False}
126 +
127 + def _validate_request(self, request: Dict[str, str]) -> str:
128 + """
129 + Validate the request to enable a rule in Wazuh Manager and return rule_id.
130 +
131 + Args:
132 + request (Dict[str, str]): The request to enable a rule in Wazuh Manager.
133 +
134 + Raises:
135 + ValueError: If the request is missing rule_id.
136 +
137 + Returns:
138 + str: rule_id.
139 + """
140 + logger.info(f"Validating enable rule request: {request}")
141 + if "rule_id" not in request:
142 + raise ValueError("Request missing rule_id")
143 + return request["rule_id"]
144 +
145 + def _fetch_filename(self, rule_id: str) -> str:
146 + """
147 + Get the filename of the rule to be enabled.
148 +
149 + Args:
150 + rule_id (str): The id of the rule to be enabled.
151 +
152 + Raises:
153 + RuntimeError: If the filename cannot be obtained.
154 +
155 + Returns:
156 + str: The filename of the rule to be enabled.
157 + """
158 + filename_data = self.wazuh_http_requests.get_request("rules", {"rule_ids": rule_id})
159 + if not filename_data["success"]:
160 + raise ValueError(filename_data["message"])
161 + return filename_data["data"]["data"]["affected_items"][0]["filename"]
162 +
163 + def _fetch_file_content(self, filename: str) -> Any:
164 + """
165 + Get the content of the rule file.
166 +
167 + Args:
168 + filename (str): The filename of the rule to be enabled.
169 +
170 + Raises:
171 + RuntimeError: If the file content cannot be obtained.
172 +
173 + Returns:
174 + Any: The content of the rule file.
175 + """
176 + file_content_data = self.wazuh_http_requests.get_request(f"rules/files/{filename}")
177 + if not file_content_data["success"]:
178 + raise ValueError(file_content_data["message"])
179 + return file_content_data["data"]["data"]["affected_items"][0]["group"]
180 +
181 + def _get_previous_level(self, rule_id: str) -> str:
182 + """
183 + Get the previous level of the rule from the `disabled_rules` table.
184 +
185 + Args:
186 + rule_id (str):The rule id to be enabled.
187 +
188 + Raises:
189 + ValueError: If the rule was not previously disabled.
190 +
191 + Returns:
192 + str: The previous level of the rule.
193 + """
194 + disabled_rule = DisabledRules.query.filter_by(rule_id=rule_id).first()
195 + if not disabled_rule:
196 + raise ValueError(f"Rule {rule_id} is not disabled.")
197 + return disabled_rule.previous_level
198 +
199 + def _set_level_previous(self, file_content: Any, rule_id: str, previous_level: str) -> Any:
200 + """
201 + Set the level of the rule to be enabled to the previous level.
202 +
203 + Args:
204 + file_content (Any): The content of the rule to be enabled.
205 + rule_id (str): The id of the rule to be enabled.
206 + previous_level (str): The previous level of the rule to be enabled.
207 +
208 + Returns:
209 + Any: The content of the rule with the level set to the previous level.
210 + """
211 + logger.info(
212 + f"Setting rule {rule_id} level to {previous_level} for file_content: {file_content}"
213 + )
214 + # If 'file_content' is a dictionary (representing a single group), make it a list of one group
215 + if isinstance(file_content, dict):
216 + file_content = [file_content]
217 +
218 + for group_block in file_content:
219 + rule_block = group_block.get("rule", None)
220 + if not rule_block:
221 + continue
222 + if isinstance(rule_block, dict):
223 + rule_block = [rule_block]
224 +
225 + for rule in rule_block:
226 + if rule["@id"] == rule_id:
227 + # Set the rule level to the previous level.
228 + rule["@level"] = previous_level
229 + break
230 +
231 + return file_content
232 +
233 + def _json_to_xml(self, file_content: Any) -> str:
234 + """
235 + Convert the rule content from JSON to XML.
236 +
237 + Args:
238 + file_content (Any): The content of the rule to be enabled.
239 +
240 + Raises:
241 + Exception: If the JSON to XML conversion fails.
242 +
243 + Returns:
244 + str: The content of the rule to be enabled in XML format.
245 + """
246 + logger.info(f"Converting file_content to XML: {file_content}")
247 +
248 + xml_content_list = []
249 + for group in file_content:
250 + xml_dict = {"group": group}
251 + xml_content = xmltodict.unparse(xml_dict, pretty=True)
252 + # Remove the `<?xml version="1.0" encoding="utf-8"?>` from the
253 + # beginning of the XML string.
254 + xml_content = xml_content.replace(
255 + '<?xml version="1.0" encoding="utf-8"?>', ""
256 + )
257 + xml_content_list.append(xml_content)
258 +
259 + # Concatenate all XML strings
260 + xml_content = "\n".join(xml_content_list)
261 + # Remove top and bottom line breaks
262 + xml_content = xml_content.strip()
263 +
264 + return xml_content
265 +
266 + def _delete_rule_from_db(self, rule_id: str):
267 + """
268 + Delete the rule from the `disabled_rules` table.
269 +
270 + Args:
271 + rule_id (str): The rule id to be deleted.
272 + """
273 + disabled_rule = DisabledRules.query.filter_by(rule_id=rule_id).first()
274 + db.session.delete(disabled_rule)
275 + db.session.commit()
276 +
277 + def _put_updated_rule(self, filename: str, xml_content: str):
278 + """
279 + PUT the updated rule to the Wazuh Manager.
280 +
281 + Args:
282 + filename (str): The filename of the rule.
283 + xml_content (str): The XML content of the rule.
284 +
285 + Raises:
286 + RuntimeError: If the PUT operation fails.
287 + """
288 + response = self.wazuh_http_requests.put_request(f"rules/files/{filename}", xml_content, params={"overwrite": "true"})
289 + if not response["success"]:
290 + raise RuntimeError(f"Could not PUT rule {filename}")