Create vulnerability.py
taylor_socfortress committed
Jul 10, 2023 at 16:46 UTC
2707d802c6451e790cf8eb9dfa127a8e48e97bab
1 file changed
+107
backend/app/services/WazuhManager/vulnerability.py
new
+107
@@ -0,0 +1,107 @@
1
+from typing import Dict, List, Optional, Any
2
+from loguru import logger
3
+from app.services.WazuhManager.universal import UniversalService
4
+import requests
5
+
6
+class WazuhHttpRequests:
7
+ """
8
+ Class to handle HTTP requests to the Wazuh API.
9
+ """
10
+ def __init__(self, connector_url: str, wazuh_auth_token: str) -> None:
11
+ """
12
+ Args:
13
+ connector_url (str): The URL of the Wazuh Manager.
14
+ wazuh_auth_token (str): The Wazuh API authentication token.
15
+ """
16
+ self.connector_url = connector_url
17
+ self.wazuh_auth_token = wazuh_auth_token
18
+ self.headers = {"Authorization": f"Bearer {wazuh_auth_token}"}
19
+
20
+ def get_request(self, endpoint: str, params: Optional[Dict[str, Any]] = None) -> Optional[Dict[str, Any]]:
21
+ """
22
+ Function to handle GET requests.
23
+
24
+ Args:
25
+ endpoint (str): The endpoint to make a GET request to.
26
+ params (Optional[Dict[str, Any]]): Any parameters to pass in the GET request.
27
+
28
+ Returns:
29
+ Optional[Dict[str, Any]]: The response from the GET request as a dictionary.
30
+ """
31
+ try:
32
+ response = requests.get(
33
+ f"{self.connector_url}/{endpoint}",
34
+ headers=self.headers,
35
+ params=params,
36
+ verify=False,
37
+ )
38
+ response.raise_for_status()
39
+ return response.json()
40
+
41
+ except Exception as e:
42
+ logger.error(f"GET request to {self.connector_url}/{endpoint} failed with error: {e}")
43
+ return None
44
+
45
+class VulnerabilityService:
46
+ """
47
+ A service class that encapsulates the logic for pulling API data from Wazuh Manager.
48
+ """
49
+ def __init__(self, universal_service: UniversalService) -> None:
50
+ """
51
+ Args:
52
+ universal_service (UniversalService): The UniversalService instance to use.
53
+ """
54
+ self.universal_service = universal_service
55
+ self.auth_token = universal_service.get_auth_token()
56
+ self.wazuh_http_requests = WazuhHttpRequests(self.universal_service.connector_url, self.auth_token)
57
+
58
+ def agent_vulnerabilities(self, agent_id: str) -> List[Dict[str, Any]]:
59
+ """
60
+ Get the vulnerabilities of an agent from Wazuh Manager.
61
+
62
+ Args:
63
+ agent_id (str): The id of the agent to get vulnerabilities for.
64
+
65
+ Returns:
66
+ List[Dict[str, Any]]: A list of processed vulnerabilities.
67
+ """
68
+ response = self.wazuh_http_requests.get_request(f"vulnerability/{agent_id}", params={"wait_for_complete": True})
69
+
70
+ if response is not None:
71
+ processed_vulnerabilities = self._process_agent_vulnerabilities(response)
72
+ return processed_vulnerabilities
73
+ return []
74
+
75
+ def _process_agent_vulnerabilities(self, response: Dict[str, Any]) -> List[Dict[str, Any]]:
76
+ """
77
+ Process the vulnerabilities of an agent from Wazuh Manager.
78
+
79
+ Args:
80
+ response (dict): The response from Wazuh Manager containing agent vulnerabilities.
81
+
82
+ Returns:
83
+ List[Dict[str, Any]]: A list of processed vulnerabilities.
84
+ """
85
+ vulnerabilities = response.get("data", {}).get("affected_items", [])
86
+ # Use list comprehension to create the processed_vulnerabilities list
87
+ processed_vulnerabilities = [
88
+ {
89
+ "severity": vuln.get("severity"),
90
+ "updated": vuln.get("updated"),
91
+ "version": vuln.get("version"),
92
+ "type": vuln.get("type"),
93
+ "name": vuln.get("name"),
94
+ "external_references": vuln.get("external_references"),
95
+ "condition": vuln.get("condition"),
96
+ "detection_time": vuln.get("detection_time"),
97
+ "cvss3_score": vuln.get("cvss3_score"),
98
+ "published": vuln.get("published"),
99
+ "architecture": vuln.get("architecture"),
100
+ "cve": vuln.get("cve"),
101
+ "status": vuln.get("status"),
102
+ "title": vuln.get("title"),
103
+ "cvss2_score": vuln.get("cvss2_score"),
104
+ }
105
+ for vuln in vulnerabilities
106
+ ]
107
+ return processed_vulnerabilities