Review code

files/gvm_library.py

@@ -11,6 +11,7 @@ import yaml
 from functools import reduce
 import os
 import git
+import pandas as pd
 
 # GVM Xpath Constants
 GVM_XPATH_ID = '@id'
@@ -461,16 +462,18 @@ class Task:
         with open(filename, "wb") as fh:
             fh.write(base64.b64decode(code))
 
-    def get_report(self) -> List[Tuple[str,str,str,str]]:
+    def get_report(self) -> pd.DataFrame:
         res = self.client.get_report(self.report_id,
                                      report_format_id=ReportFormats.anonymous_xml,
                                      ignore_pagination=True,
                                      details="1")
-        o_ids: list[str] = res.xpath('report/report/results/result/nvt/@oid')
-        severities: list[str] = res.xpath('report/report/results/result/nvt/severities/@score')
+        oids: list[str] = res.xpath('report/report/results/result/nvt/@oid')
+        sev: list[str] = res.xpath('report/report/results/result/nvt/severities/@score')
         treats: list[str] = res.xpath('report/report/results/result/threat/text()')
         ports: list[str] = res.xpath('report/report/results/result/port/text()')
-        return [ResultReport(o,s,t,p) for o,s,t,p in zip(o_ids, severities, treats, ports)]
+        df = pd.DataFrame({"oids": oids, "sev": sev, "threats": treats, "ports": ports})    
+        df['sev'] = df['sev'].astype(float)
+        return df
     
 class GVMClient():
     """
@@ -605,8 +608,8 @@ class ReportManager():
             else:
                 self.os_oids = os_oids
                 try:
-                    self.os_all_oids = list(reduce(lambda x,y: x + y,
-                                                   os_oids.values()))
+                    self.os_all_oids = tuple(set(reduce(lambda x,y: x + y,
+                                                   os_oids.values())))
                     logging.debug("Imported os security oids")
                     logging.debug(pretty_json(self.os_all_oids))
 
@@ -614,7 +617,7 @@ class ReportManager():
                     logging.warning("Impossible extract oids from imported yaml")
                     self.os_all_oids = []
 
-    def extract_oids(self, lines: List[str]) -> Set[str]:
+    def extract_oids(self, lines: List[str]) -> Tuple[str]:
         oids: List[str] = list()
         for line in lines:
             line = line.strip()
@@ -622,7 +625,7 @@ class ReportManager():
             if len(v_line := line.split(" ")[0]) > 0:
                 oids.append(v_line)
 
-        return set(oids)
+        return tuple(set(oids))
         
     def import_security_oids(self) -> None:
         user = os.environ.get(self.SS_SEC_USER)
@@ -656,7 +659,7 @@ class ReportManager():
             logging.debug("known oids")
             logging.debug(pretty_json(self.known_oids))
 
-    def import_report(self, host: str, report: List[ResultReport]):
+    def import_report(self, host: str, report: pd.DataFrame) -> None:
         self.imported_oids[host] = report
 
     def init_glob_vars(self):
@@ -695,7 +698,7 @@ class ReportManager():
             self.report[self.REPORT_DEPLOYMENT] = {self.REPORT_SEVERITY: r.severity,
                                                    self.REPORT_THREAT: r.threat}
 
-    def classify_reports(self) -> None:
+    def classify_reports(self) -> None: ######################################<<<<<<<<<<<<<<<<<<########
 
         # Init global aggregated variables
         self.init_glob_vars()