Merge pull request #32 from vulnerability-lookup/codex/add-optional-filter-for-cpe-part-parameter

Add optional CPE part filtering to search, unique, and CLI

Author: adulau

README.md

@@ -17,6 +17,12 @@ Alternatively, you can call the web server (after running `server.py`). For exam
 
 ```bash
 curl -s -X POST http://localhost:8000/search -d '{"query": ["tomcat"]}' | jq .
+
+# Optional CPE part filter (a=application, h=hardware, o=operating system)
+curl -s -X POST http://localhost:8000/search -d '{"query": ["cisco", "router"], "part": "h"}' | jq .
+
+# GET is also supported
+curl -s "http://localhost:8000/search?q=cisco+router&part=h" | jq .
 ```
 
 ### Installation
@@ -93,16 +99,17 @@ curl -s -X POST https://cpe-guesser.cve-search.org/unique -d "{\"query\": [\"out
 ### Command line - `lookup.py`
 
 ```text
-usage: lookup.py [-h] [--unique] WORD [WORD ...]
+usage: lookup.py [-h] [--unique] [--part {a,h,o}] WORD [WORD ...]
 
 Find potential CPE names from a list of keyword(s) and return a JSON of the results
 
 positional arguments:
   WORD        One or more keyword(s) to lookup
 
 options:
-  -h, --help  show this help message and exit
-  --unique    Return the best CPE matching the keywords given
+  -h, --help        show this help message and exit
+  --unique          Return the best CPE matching the keywords given
+  --part {a,h,o}    Optional CPE part filter: a (application), h (hardware), o (operating system)
 ```
 
 ```bash

bin/lookup.py

@@ -27,10 +27,15 @@
         help="Return the best CPE matching the keywords given",
         default=False,
     )
+    parser.add_argument(
+        "--part",
+        choices=sorted(CPEGuesser.VALID_CPE_PARTS),
+        help="Optional CPE part filter: a (application), h (hardware), o (operating system)",
+    )
     args = parser.parse_args()
 
     cpeGuesser = CPEGuesser()
-    r = cpeGuesser.guessCpe(args.word)
+    r = cpeGuesser.guessCpe(args.word, part=args.part)
     if not args.unique:
         print(json.dumps(r))
     else:

bin/server.py

@@ -19,6 +19,26 @@
 
 
 class Search:
+    def _is_valid_part(self, part):
+        return part in CPEGuesser.VALID_CPE_PARTS
+
+    def _split_query(self, query):
+        return [keyword for keyword in query.split() if keyword]
+
+    def _search(self, query, part, resp):
+        if not query:
+            resp.status = falcon.HTTP_400
+            resp.media = "Missing query array or incorrect JSON format"
+            return
+
+        if part is not None and not self._is_valid_part(part):
+            resp.status = falcon.HTTP_400
+            resp.media = "Invalid part parameter. Allowed values are: a, h, o"
+            return
+
+        cpeGuesser = CPEGuesser()
+        resp.media = cpeGuesser.guessCpe(query, part=part)
+
     def on_post(self, req, resp):
         data_post = req.bounded_stream.read()
         js = data_post.decode("utf-8")
@@ -36,8 +56,20 @@ def on_post(self, req, resp):
             resp.media = "Missing query array or incorrect JSON format"
             return
 
-        cpeGuesser = CPEGuesser()
-        resp.media = cpeGuesser.guessCpe(q["query"])
+        part = q.get("part")
+        if isinstance(part, str):
+            part = part.lower()
+
+        self._search(q["query"], part, resp)
+
+    def on_get(self, req, resp):
+        query = req.get_param("q")
+        part = req.get_param("part")
+
+        if isinstance(part, str):
+            part = part.lower()
+
+        self._search(self._split_query(query or ""), part, resp)
 
 
 class Unique:
@@ -58,9 +90,20 @@ def on_post(self, req, resp):
             resp.media = "Missing query array or incorrect JSON format"
             return
 
+        part = q.get("part")
+        if isinstance(part, str):
+            part = part.lower()
+        elif part is not None:
+            part = None
+
+        if part is not None and part not in CPEGuesser.VALID_CPE_PARTS:
+            resp.status = falcon.HTTP_400
+            resp.media = "Invalid part parameter. Allowed values are: a, h, o"
+            return
+
         cpeGuesser = CPEGuesser()
         try:
-            r = cpeGuesser.guessCpe(q["query"])[:1][0][1]
+            r = cpeGuesser.guessCpe(q["query"], part=part)[:1][0][1]
         except:
             r = []
         resp.media = r

lib/cpeguesser.py

@@ -12,6 +12,8 @@
 
 
 class CPEGuesser:
+    VALID_CPE_PARTS = {"a", "h", "o"}
+
     def __init__(self, rdb=None):
         self.rdb = rdb or valkey.Valkey(
             host=valkey_host,
@@ -28,7 +30,20 @@ def _rank_score(self, cpe):
         score = self.rdb.zscore("rank:cpe", cpe)
         return score or 0
 
-    def guessCpe(self, words):
+    def _is_matching_part(self, cpe, part):
+        if part is None:
+            return True
+        fields = cpe.split(":")
+        if len(fields) < 3:
+            return False
+        return fields[2] == part
+
+    def guessCpe(self, words, part=None):
+        if part is not None:
+            part = part.lower()
+            if part not in self.VALID_CPE_PARTS:
+                return []
+
         k = []
         for keyword in words:
             k.append(f"w:{keyword.lower()}")
@@ -44,6 +59,8 @@ def guessCpe(self, words):
         lowered_words = [word.lower() for word in words]
 
         for cpe in result:
+            if not self._is_matching_part(cpe, part):
+                continue
             search_score = sum(self._word_score(word, cpe) for word in lowered_words)
             rank_score = self._rank_score(cpe)
             total_score = search_score + rank_score

tests/test_cpeguesser.py

@@ -83,6 +83,31 @@ def test_guess_cpe_uses_rank_score_when_word_scores_are_missing(self):
             ],
         )
 
+    def test_guess_cpe_filters_results_by_part(self):
+        rdb = FakeRDB()
+        application = "cpe:2.3:a:cisco:router"
+        hardware = "cpe:2.3:h:cisco:router"
+
+        for key in ("w:cisco", "w:router"):
+            rdb.sadd(key, application)
+            rdb.sadd(key, hardware)
+
+        rdb.zadd("rank:cpe", {application: 5, hardware: 5})
+
+        guesser = CPEGuesser(rdb=rdb)
+
+        self.assertEqual(guesser.guessCpe(["cisco", "router"], part="h"), [(5, hardware)])
+        self.assertEqual(
+            guesser.guessCpe(["cisco", "router"], part="a"), [(5, application)]
+        )
+
+    def test_guess_cpe_rejects_unknown_part(self):
+        rdb = FakeRDB()
+        rdb.sadd("w:cisco", "cpe:2.3:a:cisco:router")
+        guesser = CPEGuesser(rdb=rdb)
+
+        self.assertEqual(guesser.guessCpe(["cisco"], part="x"), [])
+
 
 if __name__ == "__main__":
     unittest.main()