From e947eccc48d189a362a3597d04c9cb65d5d27157 Mon Sep 17 00:00:00 2001
From: Pierre-Olivier Mercier <nemunaire@nemunai.re>
Date: Thu, 3 Aug 2017 21:28:56 +0200
Subject: [PATCH] cve: improve read of partial and inexistant CVE

---
 modules/cve.py | 66 +++++++++++++++++++++++++++++++++++---------------
 1 file changed, 46 insertions(+), 20 deletions(-)

diff --git a/modules/cve.py b/modules/cve.py
index c470e29..6cdb339 100644
--- a/modules/cve.py
+++ b/modules/cve.py
@@ -5,6 +5,7 @@
 from bs4 import BeautifulSoup
 from urllib.parse import quote
 
+from nemubot.exception import IMException
 from nemubot.hooks import hook
 from nemubot.tools.web import getURLContent, striphtml
 
@@ -15,31 +16,44 @@ BASEURL_NIST = 'https://nvd.nist.gov/vuln/detail/'
 
 # MODULE CORE #########################################################
 
+VULN_DATAS = {
+    "alert-title": "vuln-warning-status-name",
+    "alert-content": "vuln-warning-banner-content",
+
+    "description": "vuln-description",
+    "published": "vuln-published-on",
+    "last_modified": "vuln-last-modified-on",
+    "source": "vuln-source",
+
+    "base_score": "vuln-cvssv3-base-score-link",
+    "severity": "vuln-cvssv3-base-score-severity",
+    "impact_score": "vuln-cvssv3-impact-score",
+    "exploitability_score": "vuln-cvssv3-exploitability-score",
+
+    "av": "vuln-cvssv3-av",
+    "ac": "vuln-cvssv3-ac",
+    "pr": "vuln-cvssv3-pr",
+    "ui": "vuln-cvssv3-ui",
+    "s": "vuln-cvssv3-s",
+    "c": "vuln-cvssv3-c",
+    "i": "vuln-cvssv3-i",
+    "a": "vuln-cvssv3-a",
+}
+
+
 def get_cve(cve_id):
     search_url = BASEURL_NIST + quote(cve_id.upper())
 
     soup = BeautifulSoup(getURLContent(search_url))
 
-    return {
-        "description": soup.body.find(attrs={"data-testid":"vuln-description"}).text.strip(),
-        "published": soup.body.find(attrs={"data-testid":"vuln-published-on"}).text.strip(),
-        "last_modified": soup.body.find(attrs={"data-testid":"vuln-last-modified-on"}).text.strip(),
-        "source": soup.body.find(attrs={"data-testid":"vuln-source"}).text.strip(),
+    vuln = {}
 
-        "base_score": float(soup.body.find(attrs={"data-testid":"vuln-cvssv3-base-score-link"}).text.strip()),
-        "severity": soup.body.find(attrs={"data-testid":"vuln-cvssv3-base-score-severity"}).text.strip(),
-        "impact_score": float(soup.body.find(attrs={"data-testid":"vuln-cvssv3-impact-score"}).text.strip()),
-        "exploitability_score": float(soup.body.find(attrs={"data-testid":"vuln-cvssv3-exploitability-score"}).text.strip()),
+    for vd in VULN_DATAS:
+        r = soup.body.find(attrs={"data-testid": VULN_DATAS[vd]})
+        if r:
+            vuln[vd] = r.text.strip()
 
-        "av": soup.body.find(attrs={"data-testid":"vuln-cvssv3-av"}).text.strip(),
-        "ac": soup.body.find(attrs={"data-testid":"vuln-cvssv3-ac"}).text.strip(),
-        "pr": soup.body.find(attrs={"data-testid":"vuln-cvssv3-pr"}).text.strip(),
-        "ui": soup.body.find(attrs={"data-testid":"vuln-cvssv3-ui"}).text.strip(),
-        "s": soup.body.find(attrs={"data-testid":"vuln-cvssv3-s"}).text.strip(),
-        "c": soup.body.find(attrs={"data-testid":"vuln-cvssv3-c"}).text.strip(),
-        "i": soup.body.find(attrs={"data-testid":"vuln-cvssv3-i"}).text.strip(),
-        "a": soup.body.find(attrs={"data-testid":"vuln-cvssv3-a"}).text.strip(),
-    }
+    return vuln
 
 
 def display_metrics(av, ac, pr, ui, s, c, i, a, **kwargs):
@@ -68,7 +82,19 @@ def get_cve_desc(msg):
             cve_id = 'cve-' + cve_id
 
         cve = get_cve(cve_id)
-        metrics = display_metrics(**cve)
-        res.append_message("{cveid}: Base score: \x02{base_score} {severity}\x0F (impact: \x02{impact_score}\x0F, exploitability: \x02{exploitability_score}\x0F; {metrics}), from \x02{source}\x0F, last modified on \x02{last_modified}\x0F. {description}".format(cveid=cve_id, metrics=metrics, **cve))
+        if not cve:
+            raise IMException("CVE %s doesn't exists." % cve_id)
+
+        if "alert-title" in cve or "alert-content" in cve:
+            alert = "\x02%s:\x0F %s " % (cve["alert-title"] if "alert-title" in cve else "",
+                                         cve["alert-content"] if "alert-content" in cve else "")
+        else:
+            alert = ""
+
+        if "base_score" not in cve and "description" in cve:
+            res.append_message("{alert}From \x02{source}\x0F, last modified on \x02{last_modified}\x0F. {description}".format(alert=alert, **cve), title=cve_id)
+        else:
+            metrics = display_metrics(**cve)
+            res.append_message("{alert}Base score: \x02{base_score} {severity}\x0F (impact: \x02{impact_score}\x0F, exploitability: \x02{exploitability_score}\x0F; {metrics}), from \x02{source}\x0F, last modified on \x02{last_modified}\x0F. {description}".format(alert=alert, metrics=metrics, **cve), title=cve_id)
 
     return res