From 63e8967d9a75203c830c90e23983e741c30e7d1f Mon Sep 17 00:00:00 2001
From: Paul Tikken <paul.tikken@gmail.com>
Date: Fri, 13 Oct 2023 14:25:14 +0000
Subject: [PATCH] moved general functions to data_source_connection class

---
 CveXplore/objects/capec.py | 16 ----------------
 CveXplore/objects/cpe.py   | 24 +++++++++++-------------
 CveXplore/objects/cves.py  | 22 +++++++++++-----------
 CveXplore/objects/cwe.py   | 16 ----------------
 CveXplore/objects/via4.py  | 16 ----------------
 5 files changed, 22 insertions(+), 72 deletions(-)

diff --git a/CveXplore/objects/capec.py b/CveXplore/objects/capec.py
index 78ecffe2..a7f17d1c 100644
--- a/CveXplore/objects/capec.py
+++ b/CveXplore/objects/capec.py
@@ -49,22 +49,6 @@ def iter_related_capecs(self):
 
                     yield capec_doc
 
-    def to_dict(self):
-        """
-        Method to convert the entire object to a dictionary
-
-        :return: Data from object
-        :rtype: dict
-        """
-
-        return {k: v for (k, v) in self.__dict__.items() if not k.startswith("_")}
-
-    def __eq__(self, other):
-        return self.__dict__ == other.__dict__
-
-    def __ne__(self, other):
-        return self.__dict__ != other.__dict__
-
     def __repr__(self):
         """String representation of object"""
         return "<< Capec:{} >>".format(self.id)
diff --git a/CveXplore/objects/cpe.py b/CveXplore/objects/cpe.py
index a032a1ee..5bb6b462 100644
--- a/CveXplore/objects/cpe.py
+++ b/CveXplore/objects/cpe.py
@@ -57,21 +57,19 @@ def iter_cves_matching_cpe(self, vuln_prod_search: bool = False):
             else:
                 yield None
 
-    def to_dict(self):
+    def to_cve_summary(self) -> dict:
         """
-        Method to convert the entire object to a dictionary
-
-        :return: Data from object
-        :rtype: dict
+        Method to request all cve's from the database based on this cpe object
         """
-
-        return {k: v for (k, v) in self.__dict__.items() if not k.startswith("_")}
-
-    def __eq__(self, other):
-        return self.__dict__ == other.__dict__
-
-    def __ne__(self, other):
-        return self.__dict__ != other.__dict__
+        all_cves = list(self.iter_cves_matching_cpe())
+        data_cves = [
+            d.to_dict("id", "cvss", "cvss3", "published", "modified", "summary")
+            for d in all_cves
+        ]
+        data_cpe = self.to_dict()
+        data_cpe["cvecount"] = len(data_cves)
+        data_cpe["cves"] = data_cves
+        return data_cpe
 
     def __repr__(self):
         """String representation of object"""
diff --git a/CveXplore/objects/cves.py b/CveXplore/objects/cves.py
index 12775313..b897f583 100644
--- a/CveXplore/objects/cves.py
+++ b/CveXplore/objects/cves.py
@@ -81,15 +81,21 @@ def iter_capec(self):
             for each in self.capec:
                 yield each
 
-    def to_dict(self):
+    def to_dict(self, *print_keys: str) -> dict:
         """
         Method to convert the entire object to a dictionary
-
-        :return: Data from object
-        :rtype: dict
         """
 
-        full_dict = {k: v for (k, v) in self.__dict__.items() if not k.startswith("_")}
+        if len(print_keys) != 0:
+            full_dict = {
+                k: v
+                for (k, v) in self.__dict__.items()
+                if not k.startswith("_") and k in print_keys
+            }
+        else:
+            full_dict = {
+                k: v for (k, v) in self.__dict__.items() if not k.startswith("_")
+            }
 
         for k, v in full_dict.items():
             if isinstance(v, DatasourceConnection):
@@ -106,12 +112,6 @@ def to_dict(self):
 
         return full_dict
 
-    def __eq__(self, other):
-        return self.__dict__ == other.__dict__
-
-    def __ne__(self, other):
-        return self.__dict__ != other.__dict__
-
     def __repr__(self):
         """String representation of object"""
         return "<< Cves:{} >>".format(self.id)
diff --git a/CveXplore/objects/cwe.py b/CveXplore/objects/cwe.py
index 063ab326..976f7623 100644
--- a/CveXplore/objects/cwe.py
+++ b/CveXplore/objects/cwe.py
@@ -48,22 +48,6 @@ def iter_related_capecs(self):
         for each in related_capecs:
             yield each
 
-    def to_dict(self):
-        """
-        Method to convert the entire object to a dictionary
-
-        :return: Data from object
-        :rtype: dict
-        """
-
-        return {k: v for (k, v) in self.__dict__.items() if not k.startswith("_")}
-
-    def __eq__(self, other):
-        return self.__dict__ == other.__dict__
-
-    def __ne__(self, other):
-        return self.__dict__ != other.__dict__
-
     def __repr__(self):
         """String representation of object"""
         return "<< Cwe:{} >>".format(self.id)
diff --git a/CveXplore/objects/via4.py b/CveXplore/objects/via4.py
index 0a9bb731..0c355dd4 100644
--- a/CveXplore/objects/via4.py
+++ b/CveXplore/objects/via4.py
@@ -16,22 +16,6 @@ def __init__(self, **kwargs):
         for each in kwargs:
             setattr(self, each, kwargs[each])
 
-    def to_dict(self):
-        """
-        Method to convert the entire object to a dictionary
-
-        :return: Data from object
-        :rtype: dict
-        """
-
-        return {k: v for (k, v) in self.__dict__.items() if not k.startswith("_")}
-
-    def __eq__(self, other):
-        return self.__dict__ == other.__dict__
-
-    def __ne__(self, other):
-        return self.__dict__ != other.__dict__
-
     def __repr__(self):
         """String representation of object"""
         return "<< Via4:{} >>".format(self.id)