improves pandas representation

3 files changed, 22 insertions, 12 deletions

diff --git a/sec_certs/certificate.py b/sec_certs/certificate.py
index 3a05cbf3..20e73585 100644
--- a/sec_certs/certificate.py
+++ b/sec_certs/certificate.py
@@ -897,8 +897,10 @@ class CommonCriteriaCert(Certificate, ComplexSerializableType):
         def from_dict(cls, dct: Dict[str, str]):
             return cls(*tuple(dct.values()))
 
-    pandas_serialization_vars = ['dgst', 'name', 'status', 'category', 'manufacturer', 'scheme', 'security_level', 'not_valid_before',
-                                 'not_valid_after', 'report_link', 'st_link', 'src', 'manufacturer_web']
+    pandas_columns = ['dgst', 'name', 'status', 'category', 'manufacturer', 'scheme', 'security_level',
+                      'not_valid_before', 'not_valid_after', 'report_link', 'st_link',
+                      'manufacturer_web', 'extracted_versions', 'cpe_matches', 'verified_cpe_matches',
+                      'related_cves']
 
     def __init__(self, status: str, category: str, name: str, manufacturer: str, scheme: str,
                  security_level: Union[str, set], not_valid_before: date,
@@ -950,7 +952,11 @@ class CommonCriteriaCert(Certificate, ComplexSerializableType):
         return self.manufacturer + ' ' + self.name + ' dgst: ' + self.dgst
 
     def to_pandas_tuple(self):
-        return tuple(getattr(self, i) for i in self.pandas_serialization_vars)
+        return self.dgst, self.name, self.status, self.category, self.manufacturer, self.scheme, self.security_level,\
+               self.not_valid_before, self.not_valid_after, self.report_link, self.st_link, self.manufacturer_web, \
+               self.heuristics.extracted_versions, self.heuristics.cpe_matches, self.heuristics.verified_cpe_matches, \
+               self.heuristics.related_cves
+
 
     def merge(self, other: 'CommonCriteriaCert'):
         """
diff --git a/sec_certs/dataset.py b/sec_certs/dataset.py
index 868fc552..60738b92 100644
--- a/sec_certs/dataset.py
+++ b/sec_certs/dataset.py
@@ -154,9 +154,11 @@ class CCDataset(Dataset, ComplexSerializableType):
         def from_dict(cls, dct: Dict[str, bool]):
             return cls(*tuple(dct.values()))
 
-    def __init__(self, certs: Dict[str, 'Certificate'], root_dir: Path, name: str = 'dataset name',
+    def __init__(self, certs: Dict[str, 'CommonCriteriaCert'], root_dir: Path, name: str = 'dataset name',
                  description: str = 'dataset_description', state: Optional[DatasetInternalState] = None):
         super().__init__(certs, root_dir, name, description)
+        self.certs: Dict[str, 'CommonCriteriaCert'] = certs  # semantically redundant, but the type hint is needed
+
         if state is None:
             state = self.DatasetInternalState()
         self.state = state
@@ -166,8 +168,16 @@ class CCDataset(Dataset, ComplexSerializableType):
 
     def to_pandas(self):
         tuples = [x.to_pandas_tuple() for x in self.certs.values()]
-        cols = CommonCriteriaCert.pandas_serialization_vars
-        return pd.DataFrame(tuples, columns=cols).set_index('dgst')
+        cols = CommonCriteriaCert.pandas_columns
+
+        df = pd.DataFrame(tuples, columns=cols)
+        df = df.set_index('dgst')
+
+        df.not_valid_before = pd.to_datetime(df.not_valid_before, infer_datetime_format=True)
+        df.not_valid_after = pd.to_datetime(df.not_valid_after, infer_datetime_format=True)
+        df = df.astype({'category': 'category', 'status': 'category', 'scheme': 'category'})
+
+        return df
 
     @classmethod
     def from_dict(cls, dct: Dict):
diff --git a/sec_certs/helpers.py b/sec_certs/helpers.py
index a8711f73..87f6e25e 100644
--- a/sec_certs/helpers.py
+++ b/sec_certs/helpers.py
@@ -491,12 +491,6 @@ def extract_keywords(filepath: Path) -> Tuple[int, Optional[Dict[str, str]]]:
     return constants.RETURNCODE_OK, result
 
 
-def match_certs(certs: Dict[str, str], cpes: List[str]):
-    results = {}
-    for dgst, cert_name in certs.items():
-        results[dgst] = process.extract(cert_name, cpes, scorer=fuzz.token_set_ratio, limit=10)
-    return results
-
 def analyze_matched_algs(data: Dict):
     pd_data = pd.Series(data)
     pd_data.hist(bins=50)