Create universal.py
taylor_socfortress committed
Jul 10, 2023 at 16:46 UTC
087a8fa078e7c727186cd480660f637849a07ae1
1 file changed
+78
backend/app/services/WazuhManager/universal.py
new
+78
@@ -0,0 +1,78 @@
1
+from loguru import logger
2
+import requests
3
+from app.models.connectors import connector_factory, Connector
4
+
5
+class UniversalService:
6
+ """
7
+ A service class that encapsulates the logic for polling messages from the Wazuh-Manager API.
8
+ """
9
+
10
+ def __init__(self) -> None:
11
+ (
12
+ self.connector_url,
13
+ self.connector_username,
14
+ self.connector_password,
15
+ ) = self.collect_wazuhmanager_details("Wazuh-Manager")
16
+
17
+ def collect_wazuhmanager_details(self, connector_name: str):
18
+ connector_instance = connector_factory.create(connector_name, connector_name)
19
+ if connector_instance.verify_connection():
20
+ connection_details = Connector.get_connector_info_from_db(connector_name)
21
+ return (
22
+ connection_details.get("connector_url"),
23
+ connection_details.get("connector_username"),
24
+ connection_details.get("connector_password"),
25
+ )
26
+ else:
27
+ logger.error(f"Connection to {connector_name} failed.")
28
+ return None, None, None
29
+
30
+ def get_auth_token(self):
31
+ """
32
+ Gets the authentication token from the Wazuh-Manager API.
33
+
34
+ Returns:
35
+ str: The authentication token.
36
+ """
37
+ try:
38
+ response = requests.get(
39
+ f"{self.connector_url}/security/user/authenticate",
40
+ auth=(self.connector_username, self.connector_password),
41
+ verify=False,
42
+ )
43
+ response.raise_for_status()
44
+ except requests.exceptions.RequestException as e:
45
+ logger.error(f"Failed to get auth token: {e}")
46
+ return None
47
+ auth_token = response.json()["data"]["token"]
48
+ logger.info(f"Authentication token: {auth_token}")
49
+ return auth_token
50
+
51
+ def restart_service(self):
52
+ """
53
+ Restart the Wazuh Manager service.
54
+
55
+ Returns:
56
+ json: A JSON response containing the updated agent information.
57
+ """
58
+ headers = {"Authorization": f"Bearer {self.get_auth_token()}"}
59
+ try:
60
+ response = requests.put(
61
+ f"{self.connector_url}/manager/restart",
62
+ headers=headers,
63
+ verify=False,
64
+ )
65
+ if response.status_code == 200:
66
+ logger.info(f"Wazuh Manager service restarted")
67
+ return {"message": "Wazuh Manager service restarted", "success": True}
68
+ else:
69
+ logger.error(
70
+ f"Wazuh Manager service restart failed with error: {response.text}"
71
+ )
72
+ return {
73
+ "message": "Wazuh Manager service restart failed",
74
+ "success": False,
75
+ }
76
+ except Exception as e:
77
+ logger.error(f"Wazuh Manager service restart failed with error: {e}")
78
+ return {"message": "Wazuh Manager service restart failed", "success": False}