diff options
| author | GeogeFI | 2022-12-25 21:35:23 +0100 |
|---|---|---|
| committer | GeogeFI | 2022-12-25 21:37:11 +0100 |
| commit | 9087cf1e35b9ad1bbe6a003bbfecdb96b9f7bb0e (patch) | |
| tree | b07df7921448d55ab7ab7900eaba709e511f82e3 /src | |
| parent | b5b0e9d1798338dca205b001e7b38f02c8929a89 (diff) | |
| download | sec-certs-9087cf1e35b9ad1bbe6a003bbfecdb96b9f7bb0e.tar.gz sec-certs-9087cf1e35b9ad1bbe6a003bbfecdb96b9f7bb0e.tar.zst sec-certs-9087cf1e35b9ad1bbe6a003bbfecdb96b9f7bb0e.zip | |
refactor: Refactoring from code review
Diffstat (limited to 'src')
| -rw-r--r-- | src/sec_certs/dataset/cve.py | 9 | ||||
| -rw-r--r-- | src/sec_certs/sample/cve.py | 12 |
2 files changed, 9 insertions, 12 deletions
diff --git a/src/sec_certs/dataset/cve.py b/src/sec_certs/dataset/cve.py index c370f3ce..6960c6fc 100644 --- a/src/sec_certs/dataset/cve.py +++ b/src/sec_certs/dataset/cve.py @@ -171,11 +171,10 @@ class CVEDataset(JSONPathDataset, ComplexSerializableType): First are matched the classic CPEs to CVEs with lookup dict and then are matched the 'AND' type CPEs containing platform. """ - cves = self._get_cves_from_exactly_matched_cpes(cpe_matches) - cves_matched_by_configurations = self._get_cves_from_cpe_configurations(cpe_matches) - cves.update(cves_matched_by_configurations) - - return cves + return { + *self._get_cves_from_exactly_matched_cpes(cpe_matches), + *self._get_cves_from_cpe_configurations(cpe_matches), + } def filter_related_cpes(self, relevant_cpes: set[CPE]): """ diff --git a/src/sec_certs/sample/cve.py b/src/sec_certs/sample/cve.py index ab18b310..d3ea1590 100644 --- a/src/sec_certs/sample/cve.py +++ b/src/sec_certs/sample/cve.py @@ -187,7 +187,7 @@ class CVE(PandasSerializableType, ComplexSerializableType): return [CPEConfiguration(platform.uri, vulnerable_cpe_uris) for platform in platforms] - def get_vulnerable_cpes_from_nist_dict(dct: dict) -> tuple[list[CPE], list[CPEConfiguration]]: + def get_vulnerable_cpes_from_nist_dict(dct: dict) -> list[list]: def get_vulnerable_cpes_and_cpe_configurations( node: dict, cpes: list[CPE], cpe_configurations: list[CPEConfiguration] ) -> tuple[list[CPE], list[CPEConfiguration]]: @@ -214,16 +214,14 @@ class CVE(PandasSerializableType, ComplexSerializableType): cpes_and_cpe_configurations = [ get_vulnerable_cpes_and_cpe_configurations(x, [], []) for x in dct["configurations"]["nodes"] ] - vulnerable_cpes = list(itertools.chain.from_iterable(map(lambda x: x[0], cpes_and_cpe_configurations))) - vulnerable_cpe_configurations = list( - itertools.chain.from_iterable(map(lambda x: x[1], cpes_and_cpe_configurations)) - ) - return vulnerable_cpes, vulnerable_cpe_configurations + return [list(t) for t in zip(*cpes_and_cpe_configurations)] cve_id = dct["cve"]["CVE_data_meta"]["ID"] impact = cls.Impact.from_nist_dict(dct) - vulnerable_cpes, vulnerable_cpe_configurations = get_vulnerable_cpes_from_nist_dict(dct) + cpe_and_cpe_configurations = get_vulnerable_cpes_from_nist_dict(dct) + vulnerable_cpes = list(itertools.chain.from_iterable(cpe_and_cpe_configurations[0])) + vulnerable_cpe_configurations = list(itertools.chain.from_iterable(cpe_and_cpe_configurations[1])) published_date = dct["publishedDate"] cwe_ids = cls.parse_cwe_data(dct) |