random fixes

3 files changed, 94 insertions, 26 deletions

diff --git a/sec_certs/certificate/common_criteria.py b/sec_certs/certificate/common_criteria.py
index d22bea40..63a43d03 100644
--- a/sec_certs/certificate/common_criteria.py
+++ b/sec_certs/certificate/common_criteria.py
@@ -583,6 +583,6 @@ class CommonCriteriaCert(Certificate, ComplexSerializableType):
 
     def get_cve_labels(self):
         if not self.heuristics.related_cves:
-            return 'None'
+            return ['None']
         else:
-            return [x.cve_id for x in self.heuristics.related_cves]
\ No newline at end of file
+            return [x.cve_id for x in self.heuristics.related_cves]
diff --git a/sec_certs/model/cve_matching.py b/sec_certs/model/cve_matching.py
index 8d909196..f89a5510 100644
--- a/sec_certs/model/cve_matching.py
+++ b/sec_certs/model/cve_matching.py
@@ -6,7 +6,7 @@ import tqdm
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.base import BaseEstimator
 from sklearn.neighbors import NearestNeighbors
-
+import seaborn as sns
 import sec_certs.helpers as helpers
 from sec_certs.dataset.cve import CVE
 
@@ -40,7 +40,7 @@ class VulnClassifier(BaseEstimator):
         else:
             return list(map(list, zip(*[self.predict_single_cert(x, return_distances=True) for x in tqdm.tqdm(X, desc='Predicting')])))
 
-    def predict_single_cert(self, cert: str, return_distances: bool = False):
+    def predict_single_cert(self, cert: str, return_distances: bool = False, n_neighbors: int = 5):
         feature_vector = self.vectorizer_.transform([cert]).todense().A1
         max_indicies = np.argpartition(feature_vector, -self.n_tokens)[-self.n_tokens:]
         subset_feature_matrix = self.feature_matrix_[:, max_indicies]
@@ -50,18 +50,24 @@ class VulnClassifier(BaseEstimator):
         # feature_matrix = scaler.fit_transform(feature_matrix)
         # feature_vector = scaler.transform(feature_vector[max_indicies].reshape(1,-1))
 
-        clf = NearestNeighbors()
+        clf = NearestNeighbors(algorithm='brute', metric='cosine', n_neighbors=n_neighbors)
         clf.fit(subset_feature_matrix)
 
         distances, indicies = clf.kneighbors(feature_vector[max_indicies].reshape(1, -1), return_distance=True)
         filtered = list(filter(lambda x: x[0] < self.cutoff_distance, list(zip(distances.flatten(), indicies.flatten()))))
 
-        result = np.array([self.labels_[x[1]] for x in filtered]) if filtered else np.array(['None'])
+        result = np.array([self.labels_[x[1]][0] for x in filtered]) if filtered else np.array(['None'])
         if not return_distances:
             return result
         else:
             return result, [x[0] for x in filtered]
 
+    def plot_histogram_of_distances(self, x, outpath='histogram_of_distances.png'):
+        result, distances = self.predict_single_cert(x, return_distances=True, n_neighbors=len(self.labels_))
+        hist_plot = sns.histplot(data=distances, x='Cert distance to CVE', y='Distance frequency', bins=50)
+        fig = hist_plot.get_figre()
+        fig.savefig(outpath, dpli=300)
+
     def prepare_df_from_description(self, description: str):
         matrix = self.vectorizer_.transform([description]).todense()
         feature_index = matrix[0, :].nonzero()[1]
diff --git a/sec_certs/model/evaluation.py b/sec_certs/model/evaluation.py
index ac070626..353fa6e6 100644
--- a/sec_certs/model/evaluation.py
+++ b/sec_certs/model/evaluation.py
@@ -1,5 +1,6 @@
 import json
 from pathlib import Path
+import logging
 from typing import Dict, List, Set, Optional, Union
 
 import numpy as np
@@ -8,6 +9,10 @@ from sec_certs.dataset.common_criteria import CCDataset
 from sec_certs.certificate.common_criteria import CommonCriteriaCert
 from sec_certs.dataset.cve import CVEDataset, CVE
 from sec_certs.serialization import CustomJSONEncoder
+import sec_certs.helpers as helpers
+import tqdm
+
+logger = logging.getLogger(__name__)
 
 
 def binarize_labels(cve_dset: CVEDataset, instances: List[List[CVE]]) -> np.array:
@@ -27,15 +32,19 @@ def get_validation_dgsts(filepath: Union[str, Path]) -> Set[str]:
     return set(data.keys())
 
 
-def get_y_true(certs: List[CommonCriteriaCert]) -> List[List[CVE]]:
-    return [set(cert.heuristics.related_cves) if cert.heuristics.related_cves else [] for cert in certs]
+def get_y_true(certs: List[CommonCriteriaCert]) -> np.array:
+    return np.array([np.array([y.cve_id for y in cert.heuristics.related_cves]) if cert.heuristics.related_cves else np.array(['None']) for cert in certs], dtype='object')
 
 
-def compute_precision(y: List[List[CVE]], y_pred: List[List[CVE]], **kwargs):
+def compute_precision(y: np.array, y_pred: np.array, **kwargs):
     prec = []
-    for pred, true in zip(y_pred, y):
+    for true, pred in zip(y, y_pred):
         set_pred = set(pred)
+        if 'None' in set_pred:
+            set_pred.remove('None')
         set_true = set(true)
+        if 'None' in set_true:
+            set_true.remove('None')
 
         if set_pred and not set_true:
             prec.append(0)
@@ -46,25 +55,78 @@ def compute_precision(y: List[List[CVE]], y_pred: List[List[CVE]], **kwargs):
     return np.mean(prec)
 
 
-def prepare_classification_report(cert_names: List[str], y_pred: List[List[CVE]], y_true: List[List[CVE]],
-                                  keywords: Set[str],
-                                  distances: Optional[List[List[float]]],
-                                  out_filepath: Optional[Union[str, Path]] = None):
+def compute_promising_ratio(y: np.array, y_pred: np.array):
+    """
+    Computes number of matched vulnerabilities that have lower distance from a certificate than the
+    first already discovered vulnerability (any from y). If no new vulnerability with such property was identified,
+    0 is assigned instead. Assumes that vulnerabilities are ordered by their similarity to given certificate.
+    """
+    if len(y_pred) > 200:
+        logger.warning('Promising matches metric should be computed only on certificates with ground-truth-verified vulnerability.')
+
+    n_promising = []
+    for instance, ground_truth_vulns in zip(y_pred, y):
+        known_before = np.array(list(map(lambda x: x in set(ground_truth_vulns), instance)))
+        true_indices = np.where(known_before)
+        if true_indices[0].size > 0 and true_indices[0][0] != 0:
+            n_promising.append(true_indices[0][0])
+        else:
+            n_promising.append(0)
+
+    return np.mean(n_promising)
+
+
+def prepare_classification_report(cert_names, y_pred, y_true, distances, cve_dataset, keywords, classifier, out_path):
+    def get_cve_representation(cve_dataset, cve_id, keywords, classifier):
+        if cve_id == 'None':
+            return None,
+        else:
+            return {
+                'cve_id': cve_id,
+                'description': cve_dataset[cve_id].description,
+                'tokenized': helpers.tokenize(cve_dataset[cve_id].description, keywords),
+                'tfidf': classifier.prepare_df_from_description(helpers.tokenize(cve_dataset[cve_id].description, keywords))['TF-IDF'].to_dict(),
+            }
+
     correctly_classified = []
     badly_classified = []
     results = {'correctly_classified': correctly_classified, 'badly_classified': badly_classified}
 
-    for index, cert in enumerate(cert_names):
-        outcome = {'certificate name': cert, 'prediction': [x.to_brief_dict(keywords) for x in y_pred[index]],
-                   'ground_truth': [x.to_brief_dict(keywords) for x in y_true[index]]}
-        if distances:
-            outcome['distances'] = distances[index]
-
-        if set(y_true[index]).issubset(set(y_pred[index])):
-            correctly_classified.append(outcome)
+    for crt, prediction, ground_truth, dis in tqdm.tqdm(zip(cert_names, y_pred, y_true, distances), desc='Preparing classification report', total=len(cert_names)):
+        record = {'certificate_name': crt,
+                  'tokenized': helpers.tokenize(crt, keywords),
+                  'tfidf': classifier.prepare_df_from_description(crt)['TF-IDF'].to_dict(),
+                  'distances': dis,
+                  'predicted_cves': [get_cve_representation(cve_dataset, cve_id, keywords, classifier) for cve_id in prediction],
+                  'true_cves': [get_cve_representation(cve_dataset, cve_id, keywords, classifier) for cve_id in ground_truth]}
+        if set(ground_truth).issubset(set(prediction)):
+            correctly_classified.append(record)
         else:
-            badly_classified.append(outcome)
+            badly_classified.append(record)
+
+    with Path(out_path).open('w') as handle:
+        json.dump(results, handle, indent=4)
+
 
-    if out_filepath:
-        with Path(out_filepath).open('w') as handle:
-            json.dump(results, handle, indent=4, cls=CustomJSONEncoder)
+# def prepare_classification_report(cert_names: List[str], y_pred: List[List[CVE]], y_true: List[List[CVE]],
+#                                   keywords: Set[str],
+#                                   distances: Optional[List[List[float]]],
+#                                   out_filepath: Optional[Union[str, Path]] = None):
+#     correctly_classified = []
+#     badly_classified = []
+#     results = {'correctly_classified': correctly_classified, 'badly_classified': badly_classified}
+#
+#     for index, cert in enumerate(cert_names):
+#         outcome = {'certificate name': cert, 'prediction': [x.to_brief_dict(keywords) for x in y_pred[index]],
+#                    'ground_truth': [x.to_brief_dict(keywords) for x in y_true[index]]}
+#         if distances:
+#             outcome['distances'] = distances[index]
+#
+#         if set(y_true[index]).issubset(set(y_pred[index])):
+#             correctly_classified.append(outcome)
+#         else:
+#             badly_classified.append(outcome)
+#
+#     if out_filepath:
+#         with Path(out_filepath).open('w') as handle:
+#             json.dump(results, handle, indent=4, cls=CustomJSONEncoder)