simplified (de)serialization using explicit _type attribute

7 files changed, 58 insertions, 52 deletions

diff --git a/oop_demo.py b/oop_demo.py
index 455c0ad7..f8ae5508 100644
--- a/oop_demo.py
+++ b/oop_demo.py
@@ -1,4 +1,5 @@
-from sec_certs.dataset import CCDataset, DatasetJSONEncoder, DatasetJSONDecoder
+from sec_certs.dataset import CCDataset
+from sec_certs.serialization import CustomJSONEncoder, CustomJSONDecoder
 from pathlib import Path
 from datetime import datetime
 import logging
@@ -18,11 +19,11 @@ def main():
 
     # Dump dataset into JSON
     with open('./debug_dataset/cc_full_dataset.json', 'w') as handle:
-        json.dump(dset, handle, cls=DatasetJSONEncoder, indent=4)
+        json.dump(dset, handle, cls=CustomJSONEncoder, indent=4)
 
     # Load dataset from JSON
     with open('./debug_dataset/cc_full_dataset.json', 'r') as handle:
-        new_dset = json.load(handle, cls=DatasetJSONDecoder)
+        new_dset = json.load(handle, cls=CustomJSONDecoder)
 
     assert dset == new_dset
 
diff --git a/sec_certs/certificate.py b/sec_certs/certificate.py
index 9ec073c2..2f8254bf 100644
--- a/sec_certs/certificate.py
+++ b/sec_certs/certificate.py
@@ -22,13 +22,12 @@ class Certificate(ABC):
     def dgst(self):
         raise NotImplementedError('Not meant to be implemented')
 
-    @abstractmethod
-    def to_dict(self) -> dict:
-        raise NotImplementedError('Not meant to be implemented')
-
     def __eq__(self, other: 'Certificate') -> bool:
         return self.dgst == other.dgst
 
+    def to_dict(self):
+        return self.__dict__
+
     @classmethod
     @abstractmethod
     def from_dict(cls, dct: dict) -> 'Certificate':
@@ -40,9 +39,6 @@ class Certificate(ABC):
 
 
 class FIPSCertificate(Certificate):
-    def to_dict(self) -> dict:
-        pass
-
     @classmethod
     def from_dict(cls, dct: dict) -> 'FIPSCertificate':
         return FIPSCertificate()
@@ -78,6 +74,9 @@ class CommonCriteriaCert(Certificate):
         def from_dict(cls, dct):
             return cls(*tuple(dct.values()))
 
+        def __lt__(self, other):
+            return self.maintainance_date < other.maintainance_date
+
     @dataclass(eq=True, frozen=True)
     class ProtectionProfile:
         """
@@ -93,6 +92,9 @@ class CommonCriteriaCert(Certificate):
         def to_dict(self):
             return self.__dict__
 
+        def __lt__(self, other):
+            return self.pp_name < other.pp_name
+
         @classmethod
         def from_dict(cls, dct):
             return cls(*tuple(dct.values()))
@@ -159,7 +161,6 @@ class CommonCriteriaCert(Certificate):
     def from_dict(cls, dct: dict) -> 'CommonCriteriaCert':
         dct['maintainance_updates'] = set(dct['maintainance_updates'])
         dct['protection_profiles'] = set(dct['protection_profiles'])
-
         args = tuple(dct.values())
 
         return cls(*args)
diff --git a/sec_certs/dataset.py b/sec_certs/dataset.py
index 69fcb2c9..25a59dba 100644
--- a/sec_certs/dataset.py
+++ b/sec_certs/dataset.py
@@ -1,4 +1,4 @@
-from datetime import datetime, date
+from datetime import datetime
 from .certificate import CommonCriteriaCert, Certificate
 from abc import ABC, abstractmethod
 from . import helpers as helpers
@@ -9,42 +9,6 @@ from bs4 import BeautifulSoup
 import locale
 import logging
 from typing import Dict
-import json
-
-
-class DatasetJSONEncoder(json.JSONEncoder):
-    def default(self, obj):
-        if isinstance(obj, Certificate):
-            return obj.to_dict()
-        if isinstance(obj, set):
-            return sorted(list(obj))
-        if isinstance(obj, date):
-            return str(obj)
-        if isinstance(obj, Path):
-            return str(obj)
-        if isinstance(obj, CommonCriteriaCert.ProtectionProfile):
-            return obj.to_dict()
-        if isinstance(obj, CommonCriteriaCert.MaintainanceReport):
-            return obj.to_dict()
-        if isinstance(obj, Dataset):
-            return obj.to_dict()
-
-        return super().default(obj)
-
-
-class DatasetJSONDecoder(json.JSONDecoder):
-    def __init__(self, *args, **kwargs):
-        json.JSONDecoder.__init__(self, object_hook=self.object_hook, *args, **kwargs)
-
-    def object_hook(self, obj):
-        if 'root_dir' in obj:  # TODO: This is a heavy simplification
-            return CCDataset.from_dict(obj)
-        if 'pp_name' in obj and 'pp_link' in obj:
-            return CommonCriteriaCert.ProtectionProfile.from_dict(obj)
-        if 'maintainance_date' in obj and 'maintainance_title' in obj and 'maintainance_report_link' in obj and 'maintainance_st_link':
-            return CommonCriteriaCert.MaintainanceReport.from_dict(obj)
-        if 'category' in obj:  # TODO: This is heavy simplification.
-            return CommonCriteriaCert.from_dict(obj)
 
 
 class Dataset(ABC):
@@ -220,7 +184,8 @@ class CCDataset(Dataset):
 
         n_all = len(df_base)
         n_deduplicated = len(df_base.drop_duplicates(subset=['dgst']))
-        logging.warning(f'The CSV {file} contains {n_all - n_deduplicated} duplicates by the primary key.')
+        if n_dup := n_all - n_deduplicated > 0:
+            logging.warning(f'The CSV {file} contains {n_dup} duplicates by the primary key.')
 
         df_base = df_base.drop_duplicates(subset=['dgst'])
         df_main = df_main.drop_duplicates()
diff --git a/sec_certs/serialization.py b/sec_certs/serialization.py
new file mode 100644
index 00000000..2fb77110
--- /dev/null
+++ b/sec_certs/serialization.py
@@ -0,0 +1,31 @@
+import json
+from datetime import date
+from pathlib import Path
+from .dataset import CCDataset
+from .certificate import CommonCriteriaCert
+
+serializable_complex_types = (CCDataset, CommonCriteriaCert, CommonCriteriaCert.MaintainanceReport, CommonCriteriaCert.ProtectionProfile)
+serializable_complex_types_dict = {x.__name__: x for x in serializable_complex_types}
+
+
+class CustomJSONEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, serializable_complex_types):
+            return {**{'_type': type(obj).__name__}, **obj.to_dict()}
+        if isinstance(obj, set):
+            return sorted(list(obj))
+        if isinstance(obj, date):
+            return str(obj)
+        if isinstance(obj, Path):
+            return str(obj)
+        return super().default(obj)
+
+
+class CustomJSONDecoder(json.JSONDecoder):
+    def __init__(self, *args, **kwargs):
+        json.JSONDecoder.__init__(self, object_hook=self.object_hook, *args, **kwargs)
+
+    def object_hook(self, obj):
+        if '_type' in obj and obj['_type'] in serializable_complex_types_dict.keys():
+            complex_type = obj.pop('_type')
+            return serializable_complex_types_dict[complex_type].from_dict(obj)
diff --git a/test/data/test_cc_oop/fictional_cert.json b/test/data/test_cc_oop/fictional_cert.json
index 44a422db..1aeaa0db 100644
--- a/test/data/test_cc_oop/fictional_cert.json
+++ b/test/data/test_cc_oop/fictional_cert.json
@@ -1,4 +1,5 @@
 {
+    "_type": "CommonCriteriaCert",
     "category": "Sample category",
     "name": "Sample certificate name",
     "manufacturer": "Sample manufacturer",
@@ -15,12 +16,14 @@
     "manufacturer_web": "http://path.to/manufacturer/web",
     "protection_profiles": [
         {
+            "_type": "ProtectionProfile",
             "pp_name": "sample_pp",
             "pp_link": "http://sample.pp"
         }
     ],
     "maintainance_updates": [
         {
+            "_type": "MaintainanceReport",
             "maintainance_date": "1900-01-01",
             "maintainance_title": "Sample maintainance",
             "maintainance_report_link": "https://maintainance.up",
diff --git a/test/data/test_cc_oop/toy_dataset.json b/test/data/test_cc_oop/toy_dataset.json
index 49d612b5..24392ad8 100644
--- a/test/data/test_cc_oop/toy_dataset.json
+++ b/test/data/test_cc_oop/toy_dataset.json
@@ -1,4 +1,5 @@
 {
+    "_type": "CCDataset",
     "root_dir": "/fictional/path/to/dataset",
     "timestamp": "2020-11-16 17:04:14.770153",
     "sha256_digest": "not implemented",
@@ -7,6 +8,7 @@
     "n_certs": 2,
     "certs": [
         {
+            "_type": "CommonCriteriaCert",
             "category": "Access Control Devices and Systems",
             "name": "NetIQ Identity Manager 4.7",
             "manufacturer": "NetIQ Corporation",
@@ -26,6 +28,7 @@
             "maintainance_updates": []
         },
         {
+            "_type": "CommonCriteriaCert",
             "category": "Access Control Devices and Systems",
             "name": "Magic SSO V4.0",
             "manufacturer": "Dreamsecurity Co., Ltd.",
@@ -40,6 +43,7 @@
             "manufacturer_web": "https://www.dreamsecurity.com/",
             "protection_profiles": [
                 {
+                    "_type": "ProtectionProfile",
                     "pp_name": "Korean National Protection Profile for Single Sign On V1.0",
                     "pp_link": "http://www.commoncriteriaportal.org/files/ppfiles/KECS-PP-0822-2017%20Korean%20National%20PP%20for%20Single%20Sign%20On%20V1.0(eng).pdf"
                 }
diff --git a/test/test_cc_oop.py b/test/test_cc_oop.py
index c29e90e3..b5740077 100644
--- a/test/test_cc_oop.py
+++ b/test/test_cc_oop.py
@@ -7,7 +7,8 @@ import filecmp
 import shutil
 import os
 
-from sec_certs.dataset import CCDataset, DatasetJSONDecoder, DatasetJSONEncoder
+from sec_certs.dataset import CCDataset
+from sec_certs.serialization import CustomJSONEncoder, CustomJSONDecoder
 from sec_certs.certificate import CommonCriteriaCert
 
 
@@ -75,7 +76,7 @@ class TestCommonCriteriaOOP(TestCase):
         fd, path = mkstemp()
         try:
             with os.fdopen(fd, 'w') as handle:
-                json.dump(obj, handle, cls=DatasetJSONEncoder, indent=4)
+                json.dump(obj, handle, cls=CustomJSONEncoder, indent=4)
 
             return filecmp.cmp(referential_path, path)
         finally:
@@ -84,7 +85,7 @@ class TestCommonCriteriaOOP(TestCase):
     @staticmethod
     def equal_from_json(referential_path, obj):
         with open(referential_path, 'r') as handle:
-            new_obj = json.load(handle, cls=DatasetJSONDecoder)
+            new_obj = json.load(handle, cls=CustomJSONDecoder)
         return obj == new_obj
 
     def test_cert_to_json(self):