Merge pull request #150 from crocs-muni/feat/slots

Add __slots__ to CPE and CVE classes.

9 files changed, 65 insertions, 10 deletions

diff --git a/.gitignore b/.gitignore
index f6054f09..3176de71 100644
--- a/.gitignore
+++ b/.gitignore
@@ -54,6 +54,8 @@ coverage.xml
 *.cover
 .hypothesis/
 .pytest_cache/
+prof/
+.pymon
 
 # Translations
 *.mo
diff --git a/dev_requirements.txt b/dev_requirements.txt
index edb5f93a..f2cceafd 100644
--- a/dev_requirements.txt
+++ b/dev_requirements.txt
@@ -2,6 +2,10 @@ mypy
 types-PyYAML
 types-python-dateutil
 types-requests
+pytest==6.2.5
+pytest-cov==3.0.0
+pytest-monitor==1.6.3
+pytest-profiling==1.7.0
 black
 isort==5.10.1
 flake8==4.0.1
diff --git a/pyproject.toml b/pyproject.toml
index 031110f1..53b1f9bd 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -26,4 +26,10 @@ skip=["certsvenv", "build"]
 [tool.mypy]
 plugins = ["numpy.typing.mypy_plugin"]
 ignore_missing_imports = true
-exclude="build/"
\ No newline at end of file
+exclude="build/"
+
+[tool.pytest.ini_options]
+markers = [
+    "slow: marks tests as slow (deselect with '-m \"not slow\"')"
+]
+addopts = "--cov sec_certs"
diff --git a/sec_certs/dataset/cpe.py b/sec_certs/dataset/cpe.py
index 6727f930..5057587e 100644
--- a/sec_certs/dataset/cpe.py
+++ b/sec_certs/dataset/cpe.py
@@ -11,7 +11,7 @@ import pandas as pd
 
 import sec_certs.helpers as helpers
 from sec_certs.dataset.cve import CVEDataset
-from sec_certs.sample.cpe import CPE
+from sec_certs.sample.cpe import CPE, cached_cpe
 from sec_certs.serialization.json import ComplexSerializableType, serialize
 
 logger = logging.getLogger(__name__)
@@ -105,7 +105,7 @@ class CPEDataset(ComplexSerializableType):
                 )
             cpe_uri = found_cpe_uri.attrib["name"]
 
-            dct[cpe_uri] = CPE(cpe_uri, title)
+            dct[cpe_uri] = cached_cpe(cpe_uri, title)
         return cls(False, Path(json_path), dct)
 
     @classmethod
diff --git a/sec_certs/dataset/cve.py b/sec_certs/dataset/cve.py
index 6e662ac6..4dc1fadc 100644
--- a/sec_certs/dataset/cve.py
+++ b/sec_certs/dataset/cve.py
@@ -16,7 +16,7 @@ import sec_certs.constants as constants
 import sec_certs.helpers as helpers
 from sec_certs.config.configuration import config
 from sec_certs.parallel_processing import process_parallel
-from sec_certs.sample.cpe import CPE
+from sec_certs.sample.cpe import CPE, cached_cpe
 from sec_certs.sample.cve import CVE
 from sec_certs.serialization.json import ComplexSerializableType, CustomJSONDecoder, CustomJSONEncoder
 
@@ -189,10 +189,10 @@ class CVEDataset(ComplexSerializableType):
             elif "versionEndExcluding" in field:
                 end_version = ("excluding", field["versionEndExcluding"])
 
-            return CPE(field["cpe23Uri"], start_version=start_version, end_version=end_version)
+            return cached_cpe(field["cpe23Uri"], start_version=start_version, end_version=end_version)
 
         def parse_values_cpe(field: Dict) -> List[CPE]:
-            return [CPE(x["cpe23Uri"]) for x in field["cpe_name"]]
+            return [cached_cpe(x["cpe23Uri"]) for x in field["cpe_name"]]
 
         logger.debug("Attempting to get NIST mapping file.")
         if not input_filepath or not input_filepath.is_file():
diff --git a/sec_certs/sample/cpe.py b/sec_certs/sample/cpe.py
index aeab6dc9..df08c105 100644
--- a/sec_certs/sample/cpe.py
+++ b/sec_certs/sample/cpe.py
@@ -1,4 +1,5 @@
 from dataclasses import dataclass
+from functools import lru_cache
 from typing import ClassVar, Dict, List, Optional, Tuple
 
 from sec_certs.serialization.json import ComplexSerializableType
@@ -15,6 +16,8 @@ class CPE(PandasSerializableType, ComplexSerializableType):
     start_version: Optional[Tuple[str, str]]
     end_version: Optional[Tuple[str, str]]
 
+    __slots__ = ["uri", "title", "version", "vendor", "item_name", "start_version", "end_version"]
+
     pandas_columns: ClassVar[List[str]] = [
         "uri",
         "vendor",
@@ -32,6 +35,7 @@ class CPE(PandasSerializableType, ComplexSerializableType):
         start_version: Optional[Tuple[str, str]] = None,
         end_version: Optional[Tuple[str, str]] = None,
     ):
+        super().__init__()
         self.uri = uri
         self.title = title
         self.start_version = start_version
@@ -85,3 +89,8 @@ class CPE(PandasSerializableType, ComplexSerializableType):
             and self.start_version == other.start_version
             and self.end_version == other.end_version
         )
+
+
+@lru_cache(maxsize=4096)
+def cached_cpe(*args, **kwargs):
+    return CPE(*args, **kwargs)
diff --git a/sec_certs/sample/cve.py b/sec_certs/sample/cve.py
index 2bc28c4a..080a4e51 100644
--- a/sec_certs/sample/cve.py
+++ b/sec_certs/sample/cve.py
@@ -5,7 +5,7 @@ from typing import ClassVar, Dict, List, Optional, Tuple
 
 from dateutil.parser import isoparse
 
-from sec_certs.sample.cpe import CPE
+from sec_certs.sample.cpe import CPE, cached_cpe
 from sec_certs.serialization.json import ComplexSerializableType
 from sec_certs.serialization.pandas import PandasSerializableType
 
@@ -16,9 +16,11 @@ class CVE(PandasSerializableType, ComplexSerializableType):
     class Impact(ComplexSerializableType):
         base_score: float
         severity: str
-        explotability_score: float
+        exploitability_score: float
         impact_score: float
 
+        __slots__ = ["base_score", "severity", "exploitability_score", "impact_score"]
+
         @classmethod
         def from_nist_dict(cls, dct: Dict):
             """
@@ -46,6 +48,8 @@ class CVE(PandasSerializableType, ComplexSerializableType):
     impact: Impact
     published_date: Optional[datetime.datetime]
 
+    __slots__ = ["cve_id", "vulnerable_cpes", "impact", "published_date"]
+
     pandas_columns: ClassVar[List[str]] = [
         "cve_id",
         "vulnerable_cpes",
@@ -57,6 +61,7 @@ class CVE(PandasSerializableType, ComplexSerializableType):
     ]
 
     def __init__(self, cve_id: str, vulnerable_cpes: List[CPE], impact: Impact, published_date: str):
+        super().__init__()
         self.cve_id = cve_id
         self.vulnerable_cpes = vulnerable_cpes
         self.impact = impact
@@ -87,7 +92,7 @@ class CVE(PandasSerializableType, ComplexSerializableType):
             self.vulnerable_cpes,
             self.impact.base_score,
             self.impact.severity,
-            self.impact.explotability_score,
+            self.impact.exploitability_score,
             self.impact.impact_score,
             self.published_date,
         )
@@ -125,7 +130,7 @@ class CVE(PandasSerializableType, ComplexSerializableType):
                             else:
                                 version_end = None
 
-                            cpe_uris.append(CPE(cpe_uri, start_version=version_start, end_version=version_end))
+                            cpe_uris.append(cached_cpe(cpe_uri, start_version=version_start, end_version=version_end))
 
                 return cpe_uris
 
diff --git a/tests/test_cpe.py b/tests/test_cpe.py
new file mode 100644
index 00000000..e2f1fdb9
--- /dev/null
+++ b/tests/test_cpe.py
@@ -0,0 +1,16 @@
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import pytest
+
+from sec_certs.dataset.cpe import CPEDataset
+
+
+class TestCPE:
+    @pytest.mark.slow
+    @pytest.mark.monitor_test
+    def test_from_web(self):
+        with TemporaryDirectory() as tmpdir:
+            dset = CPEDataset.from_web(Path(tmpdir) / "cpe.json")
+        assert dset is not None
+        assert "cpe:2.3:o:infineon:trusted_platform_firmware:6.40:*:*:*:*:*:*:*" in dset.cpes
diff --git a/tests/test_cve.py b/tests/test_cve.py
new file mode 100644
index 00000000..6c84ace3
--- /dev/null
+++ b/tests/test_cve.py
@@ -0,0 +1,13 @@
+import pytest
+
+from sec_certs.dataset.cve import CVEDataset
+
+
+class TestCVE:
+    @pytest.mark.slow
+    @pytest.mark.monitor_test
+    def test_from_web(self):
+        dset = CVEDataset.from_web()
+        assert dset is not None
+        assert "CVE-2019-15809" in dset.cves
+        assert "CVE-2017-15361" in dset.cves