Implements alpha version of cpe fuzzy matching

4 files changed, 47 insertions, 4 deletions

diff --git a/requirements.txt b/requirements.txt
index 3e282942..4a4277ba 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -24,3 +24,4 @@ tabula-py==2.2.0
 tabulate==0.8.7
 tqdm==4.51.0
 urllib3==1.25.11
+rapidfuzz==1.1.1
diff --git a/sec_certs/certificate.py b/sec_certs/certificate.py
index 31675bf3..c07a4f1b 100644
--- a/sec_certs/certificate.py
+++ b/sec_certs/certificate.py
@@ -872,7 +872,8 @@ class CommonCriteriaCert(Certificate, ComplexSerializableType):
                  protection_profiles: set,
                  maintainance_updates: set,
                  state: Optional[InternalState],
-                 pdf_data: Optional[PdfData]):
+                 pdf_data: Optional[PdfData],
+                 cpe_matching: Optional[List[Tuple[str]]]):
         super().__init__()
 
         self.status = status
@@ -899,6 +900,10 @@ class CommonCriteriaCert(Certificate, ComplexSerializableType):
             pdf_data = self.PdfData()
         self.pdf_data = pdf_data
 
+        if cpe_matching is None:
+            cpe_matching = []
+        self.cpe_matching = cpe_matching
+
     @property
     def dgst(self) -> str:
         """
@@ -1052,7 +1057,7 @@ class CommonCriteriaCert(Certificate, ComplexSerializableType):
             maintainance_div) if maintainance_div else set()
 
         return cls(status, category, name, manufacturer, scheme, security_level, not_valid_before, not_valid_after, report_link,
-                   st_link, 'html', cert_link, manufacturer_web, protection_profiles, maintainances, None, None)
+                   st_link, 'html', cert_link, manufacturer_web, protection_profiles, maintainances, None, None, None)
 
     def set_local_paths(self,
                         report_pdf_dir: Optional[Union[str, Path]],
diff --git a/sec_certs/dataset.py b/sec_certs/dataset.py
index e662acfb..8197ab60 100644
--- a/sec_certs/dataset.py
+++ b/sec_certs/dataset.py
@@ -16,6 +16,9 @@ from graphviz import Digraph
 import requests
 import pandas as pd
 from bs4 import BeautifulSoup, Tag
+from rapidfuzz import process, fuzz
+import xml.etree.ElementTree as ET
+
 
 import sec_certs.helpers as helpers
 import sec_certs.constants as constants
@@ -362,7 +365,7 @@ class CCDataset(Dataset, ComplexSerializableType):
 
         certs = {x.dgst: CommonCriteriaCert(cert_status, x.category, x.cert_name, x.manufacturer, x.scheme, x.security_level,
                                             x.not_valid_before, x.not_valid_after, x.report_link, x.st_link, 'csv',
-                                            None, None, profiles.get(x.dgst, None), updates.get(x.dgst, None), None, None) for
+                                            None, None, profiles.get(x.dgst, None), updates.get(x.dgst, None), None, None, None) for
                  x in
                  df_base.itertuples()}
         return certs
@@ -629,6 +632,32 @@ class CCDataset(Dataset, ComplexSerializableType):
         if update_json is True:
             self.to_json(self.json_path)
 
+    # TODO: Probably breaks a logic of ceritifcate managing itself. Needs design refactoring
+    def fuzzy_match_cpe(self, cpe_path: Path, update_json: bool = False):
+        def get_cpe_titles(cpe_path: Path):
+            root = ET.parse(str(cpe_path)).getroot()
+            return [child.text for child in root.findall('{http://cpe.mitre.org/dictionary/2.0}cpe-item/{http://cpe.mitre.org/dictionary/2.0}title')]
+
+        digests = [x for x in self.certs.keys()]
+        cpe_titles = get_cpe_titles(cpe_path)
+
+        def chunk_list(a: List, n: int):
+            k, m = divmod(len(a), n)
+            return (a[i * k + min(i, m):(i + 1) * k + min(i + 1, m)] for i in range(n))
+
+        chunks = chunk_list(digests, constants.N_THREADS)
+        chunks_dicts = [{x: self[x].name for x in y} for y in chunks]
+
+        results = cert_processing.process_parallel(helpers.match_certs, list(zip(chunks_dicts, [cpe_titles for _ in range(constants.N_THREADS)])), constants.N_THREADS, use_threading=False, unpack=True)
+
+        for chunk in results:
+            for digest, matches in chunk.items():
+                self[digest].cpe_matching = matches
+
+        if update_json is True:
+            self.to_json(self.json_path)
+
+
 
 class FIPSDataset(Dataset, ComplexSerializableType):
     FIPS_BASE_URL: ClassVar[str] = 'https://csrc.nist.gov'
diff --git a/sec_certs/helpers.py b/sec_certs/helpers.py
index 7a3bc78a..01729891 100644
--- a/sec_certs/helpers.py
+++ b/sec_certs/helpers.py
@@ -15,6 +15,7 @@ import pandas as pd
 import subprocess
 import sec_certs.constants as constants
 from enum import Enum
+from rapidfuzz import process, fuzz
 
 from PyPDF2 import PdfFileReader
 
@@ -486,4 +487,11 @@ def extract_keywords(filepath: Path) -> Tuple[int, Optional[Dict[str, str]]]:
         error_msg = f'Failed to parse keywords from: {filepath}; {e}'
         logger.error(error_msg)
         return error_msg, None
-    return constants.RETURNCODE_OK, result
\ No newline at end of file
+    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
\ No newline at end of file