From ca8eb1614019d867283d57aee228fe5b966ec910 Mon Sep 17 00:00:00 2001
From: ausmaster <cooljazzdrummer@cox.net>
Date: Sun, 23 Jun 2024 03:15:01 -0700
Subject: [PATCH] Implement CPE to CVE functionality.

---
 README.md |  4 +++-
 query.py  | 49 +++++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index dcdc323..9d37f7c 100644
--- a/README.md
+++ b/README.md
@@ -3,10 +3,12 @@ NIST currently restricts NVD API usage to 50 requests with an API key and 5 requ
 
 For commercial usage, this is an issue when there are multiple users requesting NVD data. 
 
-This solves this by storing a replica of NVD's CVE and CPE information in a MongoDB database.
+This solves this by storing a replica of NVD's CVE and CPE information in a localized MongoDB database.
 
 Requires Python 3.8+
 ## Features
+### Implemented
+* Given a CPE name (ex: cpe:2.3:a:nsa:ghidra:9.2:\*:\*:\*:\*:\*:\*:\*) or CPE ID, find all CVEs.
 ### In Progress
 * Machine Learning to obtain CPEs (and CVEs) from plain text entries.
 ## Usage
diff --git a/query.py b/query.py
index bc13804..ed30f4e 100644
--- a/query.py
+++ b/query.py
@@ -5,10 +5,11 @@
 
 from operator import itemgetter
 from sys import stdout
-from typing import Callable, Literal, Any
+from typing import Callable, Literal
 
 import nltk
 from nltk.tokenize import word_tokenize
+from pymongo.cursor import Cursor
 from rapidfuzz.fuzz import WRatio
 
 # pylint: disable=E0401,E0611
@@ -42,7 +43,16 @@ def q_cpe_id(self, cpe_id: str) -> CPESchema | None:
         """
         return self.client.cpes.find_one({"_id": cpe_id})
 
-    def q_cpe_matches(self, cpe_id: str) -> Any:
+    def q_cpe_name(self, cpe_name: str) -> CPESchema | None:
+        """
+        Returns CPE information given CPE Name.
+
+        :param cpe_name: The CPE Name
+        :return: Dict of CPE.
+        """
+        return self.client.cpes.find_one({"cpe_name": cpe_name})
+
+    def q_cpe_matches(self, cpe_id: str) -> Cursor:
         """
         Returns CPE matches information given CPE ID.
 
@@ -51,6 +61,41 @@ def q_cpe_matches(self, cpe_id: str) -> Any:
         """
         return self.client.cpematches.find({"matches": cpe_id})
 
+    def cpe_to_cves(self, cpe_id: str) -> Cursor[CVESchema] | None:
+        """
+        Returns all CVEs given a CPE ID.
+
+        :param cpe_id: The CPE ID
+        :return: List of all CVEs
+        """
+        matches = [match["_id"] for match in self.q_cpe_matches(cpe_id)]
+        return self.client.cves.find({
+            "configurations.nodes.cpeMatch": {
+                "$elemMatch": {
+                    "matchCriteriaId": {"$in": matches}
+                }
+            }
+        })
+
+    def cpe_name_to_cves(self, cpe_name: str) -> Cursor[CVESchema] | None:
+        """
+        Returns all CVEs given a CPE Name.
+
+        :param cpe_name: The CPE Name
+        :return: List of all CVEs
+        """
+        cpe = self.q_cpe_name(cpe_name)
+        if not cpe:
+            return None
+        matches = [match["_id"] for match in self.q_cpe_matches(cpe["_id"])]
+        return self.client.cves.find({
+            "configurations.nodes.cpeMatch": {
+                "$elemMatch": {
+                    "matchCriteriaId": {"$in": matches}
+                }
+            }
+        })
+
     def ml_find_cpe(
             self,
             cpe_search_str: str,