mockup classes for Dataset and Certificate

6 files changed, 88 insertions, 2 deletions

diff --git a/sec_certs/analyze_certificates.py b/sec_certs/analyze_certificates.py
index cce16598..db907ce4 100644
--- a/sec_certs/analyze_certificates.py
+++ b/sec_certs/analyze_certificates.py
@@ -13,7 +13,7 @@ from graphviz import Digraph
 from tabulate import tabulate
 
 from . import sanity
-from .tags_constants import *
+from .constants import *
 
 
 plt.rcdefaults()
diff --git a/sec_certs/certificate.py b/sec_certs/certificate.py
new file mode 100644
index 00000000..3e66f278
--- /dev/null
+++ b/sec_certs/certificate.py
@@ -0,0 +1,31 @@
+from typing import Type
+import json
+
+
+# TODO: Move me to appropriate place. I'm here just to demonstrate how we're going to serialize to json
+class CertificateJSONEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, Certificate):
+            return obj.asdict()
+        return super().default(obj)
+
+
+class Certificate:
+    def __init__(self):
+        self.sha256 = None
+        self.framework = None  # CC or FIPS or PP or whatever
+        self.product_name = None
+        self.vendor_name = None
+        self.pdf_path = None
+
+    def __repr__(self):
+        return str(self.asdict())
+
+    def __str__(self):
+        return 'Not implemented'  # TODO: Introduce some meaningful representation of the certificate
+
+    def asdict(self):
+        return {'not': 'implemented'}  # TODO: Implement me, I'll be used for json serialization!
+
+    def __eq__(self, other: 'Certificate') -> bool:
+        return self.sha256 == other.sha256
diff --git a/sec_certs/tags_constants.py b/sec_certs/constants.py
index 01625265..8d2e0657 100644
--- a/sec_certs/tags_constants.py
+++ b/sec_certs/constants.py
@@ -1,3 +1,10 @@
+from enum import Enum
+
+
+class CertFramework(Enum):
+    CC = 'Common Criteria'
+    FIPS = 'FIPS'
+
 TAG_MATCH_COUNTER = 'count'
 TAG_MATCH_MATCHES = 'matches'
 
diff --git a/sec_certs/dataset.py b/sec_certs/dataset.py
new file mode 100644
index 00000000..278dfae9
--- /dev/null
+++ b/sec_certs/dataset.py
@@ -0,0 +1,47 @@
+class Dataset:
+    def __init__(self):
+        self.certs = {}
+        self.framework = None
+        self.data_dir = None
+
+        self.sha256_digest = 'not implemented'  # TODO: Implement as hash of all certs
+        self.name = 'not implemented'  # TODO: Allow for naming of the datasets
+        self.description = 'not implemented'  # Allow for descriptions of the dataset
+
+    # The idea is to iterate over dataset as: `for cert in Dataset: ... `
+    def __iter__(self):
+        for cert in self.certs.values():
+            yield cert
+
+    # The idea is to be able to access certificates by calling Dataset[cert_hash]
+    def __getitem__(self, item):
+        return self.certs.__getitem__(item.lower())
+
+    # Same as above, but setting instead of getting
+    def __setitem__(self, key, value):
+        self.certs.__setitem__(key.lower(), value)
+
+    def __len__(self):
+        return len(self.certs)
+
+    def dump_to_json(self):
+        pass
+
+    # Not sure if we're going to implement
+    def dump_to_csv(self):
+        pass
+
+    # Not sure if we want such behaviour
+    def __eq__(self, other):
+        return self.certs == other.certs
+
+    def __str__(self):
+        return 'Not implemented'  # TODO: Prepare some meaningful representation of the dataset
+
+
+class DatasetCC(Dataset):
+    pass
+
+
+class DatasetFIPS(Dataset):
+    pass
diff --git a/sec_certs/extract_certificates.py b/sec_certs/extract_certificates.py
index a857a5c4..949c7f33 100644
--- a/sec_certs/extract_certificates.py
+++ b/sec_certs/extract_certificates.py
@@ -19,7 +19,7 @@ from . import sanity
 from .analyze_certificates import is_in_dict
 from .cert_rules import rules, fips_rules
 from .files import search_files, load_cert_html_file, FILE_ERRORS_STRATEGY
-from .tags_constants import *
+from .constants import *
 
 plt.rcdefaults()
 
diff --git a/sec_certs/helpers.py b/sec_certs/helpers.py
new file mode 100644
index 00000000..8913d91f
--- /dev/null
+++ b/sec_certs/helpers.py
@@ -0,0 +1 @@
+#TBA
\ No newline at end of file