implements function for manual cpe match verification

4 files changed, 42 insertions, 31 deletions

diff --git a/sec_certs/certificate.py b/sec_certs/certificate.py
index 7053e057..7057d80b 100644
--- a/sec_certs/certificate.py
+++ b/sec_certs/certificate.py
@@ -1,6 +1,6 @@
 import re
 from datetime import datetime, date
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 import logging
 from pathlib import Path
 import os
@@ -874,19 +874,19 @@ class CommonCriteriaCert(Certificate, ComplexSerializableType):
     @dataclass(init=False)
     class Heuristics(ComplexSerializableType):
         extracted_versions: List[str]
-        cpe_candidate_vendors: List[str]
+        cpe_candidate_vendors: Optional[List[str]] = field(init=False)
         cpe_matches: Optional[List[Tuple[float, CPE]]]
         verified_cpe_matches: Optional[List[CPE]]
         related_cves: Optional[List[str]]
 
-        def __init__(self, extracted_versions: Optional[List[str]] = None,
+        def __init__(self,
+                     extracted_versions: Optional[List[str]] = None,
                      cpe_matches: Optional[List[str]] = None,
-                     cpe_candidate_vendors: Optional[List[str]] = None,
                      verified_cpe_matches: Optional[List[str]] = None,
                      related_cves: Optional[List[CVE]] = None):
             self.extracted_versions = extracted_versions
             self.cpe_matches = cpe_matches
-            self.cpe_candidate_vendors = cpe_candidate_vendors
+            self.cpe_candidate_vendors = None
             self.verified_cpe_matches = verified_cpe_matches
             self.related_cves = related_cves
 
diff --git a/sec_certs/dataset.py b/sec_certs/dataset.py
index fb5b05cc..96c74f0c 100644
--- a/sec_certs/dataset.py
+++ b/sec_certs/dataset.py
@@ -6,6 +6,7 @@ from typing import Dict, List, ClassVar, Collection, Union, Set, Tuple, Optional
 from itertools import groupby
 from dataclasses import dataclass
 import copy
+import time
 
 import json
 from abc import ABC, abstractmethod
@@ -685,36 +686,47 @@ class CCDataset(Dataset, ComplexSerializableType):
         compute_candidate_cpe_vendors(cpe_dset)
         compute_cpe_matches(cpe_dset)
 
-        if not cve_dataset_path:
-            cve_dataset = CVEDataset.from_web()
-        else:
-            cve_dataset = CVEDataset.from_json(cve_dataset_path)
-
         # TODO: Invoke me back
+        # if not cve_dataset_path:
+        #     cve_dataset = CVEDataset.from_web()
+        # else:
+        #     cve_dataset = CVEDataset.from_json(cve_dataset_path)
         # compute_related_cves(cve_dataset)
 
         if update_json is True:
             self.to_json(self.json_path)
 
     def manually_verify_cpe_matches(self, update_json=True):
-        certs_to_verify: List[CommonCriteriaCert] = [x for x in self if x.heuristics.cpe_matches and not x.heuristics.verified_cpe_matches]
-        logger.info('Manually verifying CPE matches')
-        n_certs_to_verify = len(certs_to_verify)
-        for i, x in enumerate(certs_to_verify):
-            print(f'[{i}/{n_certs_to_verify}]Vendor: {x.manufacturer}, name: {x.name}')
-            for index, c in enumerate(x.heuristics.cpe_matches):
-                print(f'\t- {[index]}: {c[1]}')
-            print(f'\t- [X]: No fitting match')
-            inpt = input('Select fitting CPE matches (split with comma if choosing more):')
-            inpts = [x for x in inpt.split(',')]
+        def verify_certs(certificates_to_verify: List[CommonCriteriaCert]):
+            n_certs_to_verify = len(certificates_to_verify)
+            for i, x in enumerate(certificates_to_verify):
+                print(f'[{i}/{n_certs_to_verify}] Vendor: {x.manufacturer}, Name: {x.name}')
+                for index, c in enumerate(x.heuristics.cpe_matches):
+                    print(f'\t- {[index]}: {c[1]}')
+                print(f'\t- [X]: No fitting match')
+                inpt = input('Select fitting CPE matches (split with comma if choosing more):')
+                inpts = [x for x in inpt.strip().split(',')]
+
+                if 'X' not in inpts:
+                    try:
+                        inpts = [int(x) for x in inpts]
+                        if min(inpts) < 0 or max(inpts) > len(x.heuristics.cpe_matches) - 1:
+                            raise ValueError(f'Incorrect number chosen, choose in range 0-{len(x.heuristics.cpe_matches) - 1}')
+                    except ValueError:
+                        logger.error('Bad input from user, repeating instance')
+                        verify_certs([x])
+                    else:
+                        matches = [x.heuristics.cpe_matches[y][1] for y in inpts]
+                        self[x.dgst].heuristics.verified_cpe_matches = matches
 
-            if 'X' not in inpts:
-                inpts = [int(x) for x in inpts]
-                matches = [x.heuristics.cpe_matches[y][1] for y in inpts]
-                self[x.dgst].heuristics.verified_cpe_matches = matches
-            if not i % 10:
-                print(f'Saving progress.')
-                self.to_json()
+                        if i != 0 and not i % 10:
+                            print(f'Saving progress.')
+                            self.to_json()
+
+        certs_to_verify: List[CommonCriteriaCert] = [x for x in self if (x.heuristics.cpe_matches and not x.heuristics.verified_cpe_matches)]
+        logger.info('Manually verifying CPE matches')
+        time.sleep(0.05)  # easier than flushing the logger
+        verify_certs(certs_to_verify)
 
         if update_json is True:
             self.to_json()
@@ -781,7 +793,6 @@ class FIPSDataset(Dataset, ComplexSerializableType):
         return output
 
 
-
     def download_all_pdfs(self):
         sp_paths, sp_urls = [], []
         self.policies_dir.mkdir(exist_ok=True)
diff --git a/test/data/test_cc_oop/fictional_cert.json b/test/data/test_cc_oop/fictional_cert.json
index d30fe35b..3856548d 100644
--- a/test/data/test_cc_oop/fictional_cert.json
+++ b/test/data/test_cc_oop/fictional_cert.json
@@ -55,7 +55,7 @@
         "_type": "Heuristics",
         "extracted_versions": null,
         "cpe_matches": null,
-        "verified_cpe_match": null,
+        "verified_cpe_matches": null,
         "related_cves": null
     }
 }
\ No newline at end of file
diff --git a/test/data/test_cc_oop/toy_dataset.json b/test/data/test_cc_oop/toy_dataset.json
index cf1c7de9..62ea04de 100644
--- a/test/data/test_cc_oop/toy_dataset.json
+++ b/test/data/test_cc_oop/toy_dataset.json
@@ -58,7 +58,7 @@
                 "_type": "Heuristics",
                 "extracted_versions": null,
                 "cpe_matches": null,
-                "verified_cpe_match": null,
+                "verified_cpe_matches": null,
                 "related_cves": null
             }
         },
@@ -109,7 +109,7 @@
                 "_type": "Heuristics",
                 "extracted_versions": null,
                 "cpe_matches": null,
-                "verified_cpe_match": null,
+                "verified_cpe_matches": null,
                 "related_cves": null
             }
         }