implement tests for cpe and cve stuff

6 files changed, 179 insertions, 10 deletions

diff --git a/sec_certs/cpe.py b/sec_certs/cpe.py
index b82de340..7b06bec9 100644
--- a/sec_certs/cpe.py
+++ b/sec_certs/cpe.py
@@ -1,7 +1,7 @@
 from dataclasses import dataclass, field
 import logging
 import json
-from typing import Optional, List, Dict, Tuple, Set
+from typing import Optional, List, Dict, Tuple, Set, Union
 import itertools
 import re
 from rapidfuzz import process, fuzz
@@ -101,8 +101,8 @@ class CPEDataset:
         self.vendor_version_to_cpe = {x: list(set(y)) for x, y in self.vendor_version_to_cpe.items()}
 
     @classmethod
-    def from_json(cls, json_path: str):
-        with open(json_path, 'r') as handle:
+    def from_json(cls, json_path: Union[str, Path]):
+        with Path(json_path).open('r') as handle:
             data = json.load(handle)
         return cls({x: CPE(x, y) for x, y in data.items()})
 
diff --git a/sec_certs/cve.py b/sec_certs/cve.py
index e7bbfc0a..4b8c0912 100644
--- a/sec_certs/cve.py
+++ b/sec_certs/cve.py
@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import Dict, List,  Optional
+from typing import Dict, List,  Optional, Union
 import copy
 import datetime
 from pathlib import Path
@@ -18,9 +18,9 @@ from sec_certs.serialization import ComplexSerializableType, CustomJSONDecoder,
 logger = logging.getLogger(__name__)
 
 
-@dataclass
+@dataclass(eq=True, frozen=True)
 class CVE(ComplexSerializableType):
-    @dataclass
+    @dataclass(eq=True)
     class Impact(ComplexSerializableType):
         base_score: float
         severity: str
@@ -94,7 +94,7 @@ class CVE(ComplexSerializableType):
         return CVE(cve_id, vulnerable_cpes, impact)
 
 
-@dataclass
+@dataclass(eq=True)
 class CVEDataset(ComplexSerializableType):
     cves: Dict[str, CVE]
     cpes_to_cve_lookup: Dict[str, List[str]] = field(init=False)
@@ -174,9 +174,8 @@ class CVEDataset(ComplexSerializableType):
             json.dump(self, handle, indent=4, cls=CustomJSONEncoder, ensure_ascii=False)
 
     @classmethod
-    def from_json(cls, input_path: str):
-        input_path = Path(input_path)
-        with input_path.open('r') as handle:
+    def from_json(cls, input_path: Union[str, Path]):
+        with Path(input_path).open('r') as handle:
             dset = json.load(handle, cls=CustomJSONDecoder)
         return dset
 
diff --git a/test/data/test_cpe_cve/auxillary_datasets/cpe_dataset.json b/test/data/test_cpe_cve/auxillary_datasets/cpe_dataset.json
new file mode 100644
index 00000000..7bb22f1c
--- /dev/null
+++ b/test/data/test_cpe_cve/auxillary_datasets/cpe_dataset.json
@@ -0,0 +1,6 @@
+{
+    "cpe:2.3:a:ibm:security_access_manager_for_enterprise_single_sign-on:8.2.2:*:*:*:*:*:*:*": "IBM Security Access Manager For Enterprise Single Sign-On 8.2.2",
+    "cpe:2.3:a:ibm:security_key_lifecycle_manager:2.6.0.1:*:*:*:*:*:*:*": "IBM Security Key Lifecycle Manager 2.6.0.1",
+    "cpe:2.3:a:semperplugins:all_in_one_seo_pack:1.3.6.4:*:*:*:*:wordpress:*:*": "Semper Plugins All in One SEO Pack 1.3.6.4 for WordPress",
+    "cpe:2.3:a:tracker-software:pdf-xchange_lite_printer:6.0.320.0:*:*:*:*:*:*:*": "Tracker Software PDF-XChange Lite Printer 6.0.320.0"
+}
\ No newline at end of file
diff --git a/test/data/test_cpe_cve/auxillary_datasets/cve_dataset.json b/test/data/test_cpe_cve/auxillary_datasets/cve_dataset.json
new file mode 100644
index 00000000..0572310b
--- /dev/null
+++ b/test/data/test_cpe_cve/auxillary_datasets/cve_dataset.json
@@ -0,0 +1,33 @@
+{
+    "_type": "CVEDataset",
+    "cves": {
+        "CVE-2017-1732": {
+            "_type": "CVE",
+            "cve_id": "CVE-2017-1732",
+            "vulnerable_cpes": [
+                "cpe:2.3:a:ibm:security_access_manager_for_enterprise_single_sign-on:8.2.2:*:*:*:*:*:*:*"
+            ],
+            "impact": {
+                "_type": "Impact",
+                "base_score": 5.3,
+                "severity": "MEDIUM",
+                "explotability_score": 3.9,
+                "impact_score": 1.4
+            }
+        },
+        "CVE-2019-4513": {
+            "_type": "CVE",
+            "cve_id": "CVE-2019-4513",
+            "vulnerable_cpes": [
+                "cpe:2.3:a:ibm:security_access_manager_for_enterprise_single_sign-on:8.2.2:*:*:*:*:*:*:*"
+            ],
+            "impact": {
+                "_type": "Impact",
+                "base_score": 8.2,
+                "severity": "HIGH",
+                "explotability_score": 3.9,
+                "impact_score": 4.2
+            }
+        }
+    }
+}
\ No newline at end of file
diff --git a/test/data/test_cpe_cve/vulnerable_dataset.json b/test/data/test_cpe_cve/vulnerable_dataset.json
new file mode 100644
index 00000000..30a62017
--- /dev/null
+++ b/test/data/test_cpe_cve/vulnerable_dataset.json
@@ -0,0 +1,68 @@
+{
+    "_type": "CCDataset",
+    "state": {
+        "_type": "DatasetInternalState",
+        "meta_sources_parsed": true,
+        "pdfs_downloaded": false,
+        "pdfs_converted": false,
+        "txt_data_extracted": false,
+        "certs_analyzed": false
+    },
+    "timestamp": "2021-04-16 15:05:18.386794",
+    "sha256_digest": "not implemented",
+    "name": "cc_full_dataset",
+    "description": "sample dataset description",
+    "n_certs": 1,
+    "certs": [
+        {
+            "_type": "CommonCriteriaCert",
+            "dgst": "c01e5375331b25dc",
+            "status": "active",
+            "category": "Access Control Devices and Systems",
+            "name": "IBM Security Access Manager for Enterprise Single Sign-On Version 8.2",
+            "manufacturer": "IBM Corporation",
+            "scheme": "DE",
+            "security_level": [
+                "ALC_FLR.1",
+                "EAL3+"
+            ],
+            "not_valid_before": "2014-12-05",
+            "not_valid_after": null,
+            "report_link": "http://www.commoncriteriaportal.org/files/epfiles/0683a_pdf.pdf",
+            "st_link": "http://www.commoncriteriaportal.org/files/epfiles/0683b_pdf.pdf",
+            "src": "csv + html",
+            "cert_link": null,
+            "manufacturer_web": "http://www.ibm.com",
+            "protection_profiles": [],
+            "maintainance_updates": [],
+            "state": {
+                "_type": "InternalState",
+                "st_link_ok": true,
+                "report_link_ok": true,
+                "st_convert_ok": true,
+                "report_convert_ok": true,
+                "st_extract_ok": true,
+                "report_extract_ok": true,
+                "errors": []
+            },
+            "pdf_data": {
+                "_type": "PdfData",
+                "report_metadata": null,
+                "st_metadata": null,
+                "report_frontpage": null,
+                "st_frontpage": null,
+                "report_keywords": null,
+                "st_keywords": null
+            },
+            "heuristics": {
+                "_type": "Heuristics",
+                "extracted_versions": [
+                    "8.2"
+                ],
+                "cpe_matches": null,
+                "verified_cpe_matches": null,
+                "related_cves": null
+            }
+        }
+    ]
+}
\ No newline at end of file
diff --git a/test/test_cve_cpe_matching.py b/test/test_cve_cpe_matching.py
new file mode 100644
index 00000000..c9dd05f6
--- /dev/null
+++ b/test/test_cve_cpe_matching.py
@@ -0,0 +1,63 @@
+from unittest import TestCase
+from sec_certs.dataset import CCDataset
+from sec_certs.certificate import CommonCriteriaCert
+from sec_certs.cpe import CPEDataset, CPE
+from sec_certs.cve import CVEDataset, CVE
+from pathlib import Path
+
+
+class TestCPEandCVEMatching(TestCase):
+    def setUp(self) -> None:
+        self.test_data_dir = Path(__file__).parent / 'data' / 'test_cpe_cve'
+        self.cc_dset = CCDataset.from_json(self.test_data_dir / 'vulnerable_dataset.json')
+        self.cc_dset.compute_heuristics(update_json=False)
+
+        self.cpes = [CPE("cpe:2.3:a:ibm:security_access_manager_for_enterprise_single_sign-on:8.2.2:*:*:*:*:*:*:*", "IBM Security Access Manager For Enterprise Single Sign-On 8.2.2"),
+                CPE("cpe:2.3:a:ibm:security_key_lifecycle_manager:2.6.0.1:*:*:*:*:*:*:*", "IBM Security Key Lifecycle Manager 2.6.0.1"),
+                CPE("cpe:2.3:a:semperplugins:all_in_one_seo_pack:1.3.6.4:*:*:*:*:wordpress:*:*", "Semper Plugins All in One SEO Pack 1.3.6.4 for WordPress"),
+                CPE("cpe:2.3:a:tracker-software:pdf-xchange_lite_printer:6.0.320.0:*:*:*:*:*:*:*", "Tracker Software PDF-XChange Lite Printer 6.0.320.0")]
+        self.cpe_dset = CPEDataset({x.uri: x for x in self.cpes})
+
+        self.cves = [CVE('CVE-2017-1732', ['cpe:2.3:a:ibm:security_access_manager_for_enterprise_single_sign-on:8.2.2:*:*:*:*:*:*:*'], CVE.Impact(5.3, 'MEDIUM', 3.9, 1.4)),
+                     CVE('CVE-2019-4513', ['cpe:2.3:a:ibm:security_access_manager_for_enterprise_single_sign-on:8.2.2:*:*:*:*:*:*:*'], CVE.Impact(8.2, 'HIGH', 3.9, 4.2))]
+        self.cve_dset = CVEDataset({x.cve_id: x for x in self.cves})
+
+    def test_load_cpe_dataset(self):
+        json_cpe_dset = CPEDataset.from_json(self.test_data_dir / 'auxillary_datasets' / 'cpe_dataset.json')
+        self.assertEqual(self.cpe_dset, json_cpe_dset, 'CPE template dataset does not match CPE dataset loaded from json.')
+
+    def test_cpe_lookup_dicts(self):
+        self.assertEqual(self.cpe_dset.vendors, {'ibm', 'tracker-software', 'semperplugins'},
+                         'The set of versions in CPE dataset does not match template')
+        alt_lookup = {x: set(y) for x, y in self.cpe_dset.vendor_to_versions.items()}
+        self.assertEqual(alt_lookup, {'ibm': {'8.2.2', '2.6.0.1'}, 'semperplugins': {'1.3.6.4'}, 'tracker-software': {'6.0.320.0'}},
+                         'The CPE lookup dictionary vendor->version of CPE dataset does not match template.')
+        self.assertEqual(self.cpe_dset.vendor_version_to_cpe, {('ibm', '8.2.2'): [CPE('cpe:2.3:a:ibm:security_access_manager_for_enterprise_single_sign-on:8.2.2:*:*:*:*:*:*:*', 'IBM Security Access Manager For Enterprise Single Sign-On 8.2.2')], ('ibm', '2.6.0.1'): [CPE('cpe:2.3:a:ibm:security_key_lifecycle_manager:2.6.0.1:*:*:*:*:*:*:*', 'IBM Security Key Lifecycle Manager 2.6.0.1')], ('semperplugins', '1.3.6.4'): [CPE('cpe:2.3:a:semperplugins:all_in_one_seo_pack:1.3.6.4:*:*:*:*:wordpress:*:*', 'Semper Plugins All in One SEO Pack 1.3.6.4 for WordPress')], ('tracker-software', '6.0.320.0'): [CPE('cpe:2.3:a:tracker-software:pdf-xchange_lite_printer:6.0.320.0:*:*:*:*:*:*:*', 'Tracker Software PDF-XChange Lite Printer 6.0.320.0')]},
+                         'The CPE lookup dictionary (vendor,version)->cpe does not match the template.')
+
+    def test_cve_lookup_dicts(self):
+        alt_lookup = {x: set(y) for x,y in self.cve_dset.cpes_to_cve_lookup.items()}
+        self.assertEqual(alt_lookup, {'cpe:2.3:a:ibm:security_access_manager_for_enterprise_single_sign-on:8.2.2:*:*:*:*:*:*:*': set(['CVE-2017-1732', 'CVE-2019-4513'])},
+                         'The CVE lookup dicionary cve-> affected cpes does not match the template')
+
+    def test_load_cve_dataset(self):
+        json_cve_dset = CVEDataset.from_json(self.test_data_dir / 'auxillary_datasets' / 'cve_dataset.json')
+        self.assertEqual(self.cve_dset, json_cve_dset, 'CVE template dataset does not match CVE dataset loaded from json.')
+
+    def test_match_cpe(self):
+        self.assertTrue(self.cpes[0] in [x[1] for x in self.cc_dset['c01e5375331b25dc'].heuristics.cpe_matches], 'The CPE matching algorithm did not find the right CPE.')
+        self.assertTrue(len(self.cc_dset['c01e5375331b25dc'].heuristics.cpe_matches) == 1, 'Exactly one CPE match should be found.')
+
+    def test_find_related_cves(self):
+        self.cc_dset['c01e5375331b25dc'].heuristics.verified_cpe_matches = [self.cpes[0]]
+        self.cc_dset.compute_related_cves()
+        self.assertCountEqual([x.cve_id for x in self.cves], self.cc_dset['c01e5375331b25dc'].heuristics.related_cves, 'The computed CVEs do not match the excpected CVEs')
+
+    def test_version_extraction(self):
+        self.assertEqual(self.cc_dset['c01e5375331b25dc'].heuristics.extracted_versions, ['8.2'], 'The version extracted from the certificate does not match the template')
+        new_cert = CommonCriteriaCert('', '', 'IDOneClassIC Card : ID-One Cosmo 64 RSA v5.4 and applet IDOneClassIC v1.0 embedded on P5CT072VOP', '', '',
+                                      '', None, None, '', '', '', '', '', set(), set(), None, None, None)
+        new_cert.get_heuristics_version()
+        self.assertEqual(set(new_cert.heuristics.extracted_versions), {'5.4', '1.0'}, 'The extracted versions do not match the template.')
+
+