CLI for retrieving CVE data from NIST with caching.

cve.py

import sqlite3
import sys
import requests

# The public rate limit (without an API key) is 5 requests in a rolling 30 second window.
# The rate limit with an API key is 50 requests in a rolling 30 second window
# https://nvd.nist.gov/developers/start-here
nist_url = 'https://services.nvd.nist.gov/rest/json/cves/2.0'
db_file = 'cve.db'


class CVE:
    def __init__(self, cve):
        self.id = None
        self._conn = None
        self._init_db()
        self._init_from_db(cve)
        if self.id is None:
            self._init_from_nist(cve)
            self._save_to_db()

    def __str__(self):
        return f"""CVE                    | {self.id}
Published              | {self.published}
Base Severity          | {self.base_severity}
Base Score             | {self.base_score}
Confidentiality Impact | {self.confidentiality_impact}
Integrity Impact       | {self.integrity_impact}
Availability Impact    | {self.availability_impact}
Attack Vector          | {self.attack_vector}
Attack Complexity      | {self.attack_complexity}
Impact Score           | {self.impact_score}
Exploitability Score   | {self.exploitability_score}
Weaknesses             | {', '.join(self.weaknesses)}

{self.description}"""

    def _init_from_nist(self, cve):
        data = requests.get(nist_url, params={'cveId': cve}).json()['vulnerabilities'][0]['cve']
        self.id = data['id']
        self.published = data['published']
        for description in data['descriptions']:
            if description['lang'] == 'en':
                self.description = description['value']
                break
        self.weaknesses = []
        for weakness in data['weaknesses']:
            for description in weakness['description']:
                if description['lang'] == 'en':
                    self.weaknesses.append(description['value'])
                    break
        if 'cvssMetricV31' in data['metrics']:
            cvss = data['metrics']['cvssMetricV31'][0]
            self.attack_vector = cvss['cvssData']['attackVector']
            self.attack_complexity = cvss['cvssData']['attackComplexity']
            self.privileges_required = cvss['cvssData']['privilegesRequired']
            self.user_interaction = cvss['cvssData']['userInteraction']
            self.scope = cvss['cvssData']['scope']
            self.confidentiality_impact = cvss['cvssData']['confidentialityImpact']
            self.integrity_impact = cvss['cvssData']['integrityImpact']
            self.availability_impact = cvss['cvssData']['availabilityImpact']
            self.base_score = cvss['cvssData']['baseScore']
            self.base_severity = cvss['cvssData']['baseSeverity']
            self.exploitability_score = cvss['exploitabilityScore']
            self.impact_score = cvss['impactScore']
        elif 'cvsMetricV30' in data['metrics']:
            cvss = data['metrics']['cvssMetricV31'][0]
            self.attack_vector = cvss['cvssData']['attackVector']
            self.attack_complexity = cvss['cvssData']['attackComplexity']
            self.privileges_required = cvss['cvssData']['privilegesRequired']
            self.user_interaction = cvss['cvssData']['userInteraction']
            self.scope = cvss['cvssData']['scope']
            self.confidentiality_impact = cvss['cvssData']['confidentialityImpact']
            self.integrity_impact = cvss['cvssData']['integrityImpact']
            self.availability_impact = cvss['cvssData']['availabilityImpact']
            self.base_score = cvss['cvssData']['baseScore']
            self.base_severity = cvss['cvssData']['baseSeverity']
            self.exploitability_score = cvss['exploitabilityScore']
            self.impact_score = cvss['impactScore']
        else:
            cvss = data['metrics']['cvssMetricV2'][0]
            self.attack_vector = cvss['cvssData']['accessVector']
            self.attack_complexity = cvss['cvssData']['accessComplexity']
            self.privileges_required = "N/A"
            self.user_interaction = "N/A"
            self.scope = "N/A"
            self.confidentiality_impact = cvss['cvssData']['confidentialityImpact']
            self.integrity_impact = cvss['cvssData']['integrityImpact']
            self.availability_impact = cvss['cvssData']['availabilityImpact']
            self.base_score = cvss['cvssData']['baseScore']
            self.base_severity = cvss['baseSeverity']
            self.exploitability_score = cvss['exploitabilityScore']
            self.impact_score = cvss['impactScore']

    def _init_from_db(self, cve):
        sql_statement = """
        SELECT * FROM cve WHERE id = ?;
        """
        data = self._conn.execute(sql_statement, (cve,)).fetchone()
        if data is not None:
            self.id = data[0]
            self.published = data[1]
            self.description = data[2]
            self.weaknesses = data[3].split(',')
            self.attack_vector = data[4]
            self.attack_complexity = data[5]
            self.privileges_required = data[6]
            self.user_interaction = data[7]
            self.scope = data[8]
            self.confidentiality_impact = data[9]
            self.integrity_impact = data[10]
            self.availability_impact = data[11]
            self.base_score = data[12]
            self.base_severity = data[13]
            self.exploitability_score = data[14]
            self.impact_score = data[15]

    def _init_db(self):
        if self._conn is None:
            conn = sqlite3.connect(db_file)
            sql_statement = """
            CREATE TABLE IF NOT EXISTS cve (id, published, description, weaknesses, attack_vector, attack_complexity,
                                            privileges_required, user_interaction, scope, confidentiality_impact,
                                            integrity_impact, availability_impact, base_score, base_severity, 
                                            exploitability_score, impact_score);
            """
            conn.execute(sql_statement)
            self._conn = conn

    def _save_to_db(self):
        sql_statement = """
        INSERT INTO cve (id, published, description, weaknesses, attack_vector, attack_complexity, privileges_required,
                         user_interaction, scope, confidentiality_impact, integrity_impact, availability_impact, 
                         base_score, base_severity, exploitability_score, impact_score) 
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);
        """
        self._conn.execute(sql_statement, (self.id, self.published, self.description, ','.join(self.weaknesses),
                                           self.attack_vector, self.attack_complexity, self.privileges_required,
                                           self.user_interaction, self.scope, self.confidentiality_impact,
                                           self.integrity_impact, self.availability_impact, self.base_score,
                                           self.base_severity, self.exploitability_score, self.impact_score))
        self._conn.commit()


if __name__ == '__main__':
    print(CVE(sys.argv[1]))