From 40cbd0572ec5ea7454145303fc8bf1d0c71d9682 Mon Sep 17 00:00:00 2001 From: GeogeFI Date: Sun, 25 Dec 2022 22:43:21 +0100 Subject: chore: Formating, fixed tests --- src/sec_certs/dataset/cve.py | 2 +- src/sec_certs/sample/cve.py | 8 ++++++-- 2 files changed, 7 insertions(+), 3 deletions(-) (limited to 'src') diff --git a/src/sec_certs/dataset/cve.py b/src/sec_certs/dataset/cve.py index 566bc2e4..9a36fa6b 100644 --- a/src/sec_certs/dataset/cve.py +++ b/src/sec_certs/dataset/cve.py @@ -84,7 +84,7 @@ class CVEDataset(JSONPathDataset, ComplexSerializableType): for cve in tqdm(self, desc="Building-up lookup dictionaries for fast CVE matching"): # See note above, we use matching_dict.get(cpe, []) instead of matching_dict[cpe] as would be expected if use_nist_mapping: - vulnerable_configurations = list( + vulnerable_configurations = set( itertools.chain.from_iterable(matching_dict.get(cpe, []) for cpe in cve.vulnerable_cpes) ) else: diff --git a/src/sec_certs/sample/cve.py b/src/sec_certs/sample/cve.py index caf2b4d7..9b2c2437 100644 --- a/src/sec_certs/sample/cve.py +++ b/src/sec_certs/sample/cve.py @@ -221,8 +221,12 @@ class CVE(PandasSerializableType, ComplexSerializableType): impact = cls.Impact.from_nist_dict(dct) cpe_and_cpe_configurations = get_vulnerable_cpes_from_nist_dict(dct) # There exist CVEs such as (CVE-2022-0177) which are rejected and do not contain any assinged CPEs - vulnerable_cpes = set(itertools.chain.from_iterable(cpe_and_cpe_configurations[0])) if cpe_and_cpe_configurations else set() - vulnerable_cpe_configurations = set(itertools.chain.from_iterable(cpe_and_cpe_configurations[1])) if cpe_and_cpe_configurations else set() + vulnerable_cpes = ( + set(itertools.chain.from_iterable(cpe_and_cpe_configurations[0])) if cpe_and_cpe_configurations else set() + ) + vulnerable_cpe_configurations = ( + set(itertools.chain.from_iterable(cpe_and_cpe_configurations[1])) if cpe_and_cpe_configurations else set() + ) published_date = dct["publishedDate"] cwe_ids = cls.parse_cwe_data(dct) -- cgit v1.3.1