CVE_searching_tool.py

from typing import List
from pydantic import BaseModel
import requests
from bs4 import BeautifulSoup

class CVE(BaseModel):
    cve_id: str
    description: str
    published_date: str
    cvss_score: str

def fetch_cves_by_keyword(keyword: str) -> List[dict]:
    base_url = "https://nvd.nist.gov"
    search_url = f"{base_url}/vuln/search/results"
    
    params = {
        "form_type": "Basic",
        "results_type": "overview",
        "query": keyword,
        "search_type": "all",  # last3months
        "isCpeNameSearch": "false"
    }

    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
    }

    cves = []
    page = 0

    while True:
        params["startIndex"] = page * 20  

        response = requests.get(search_url, headers=headers, params=params)

        if response.status_code != 200:
            print(f"Failed to retrieve data: {response.status_code}")
            break

        soup = BeautifulSoup(response.content, 'html.parser')

        rows = soup.find_all('tr', attrs={'data-testid': lambda x: x and x.startswith('vuln-row-')})

        if not rows:
            break 

        for row in rows:
            cve_id = row.find('a', {'data-testid': lambda x: x and 'vuln-detail-link' in x}).text.strip()
            description = row.find('p', {'data-testid': lambda x: x and 'vuln-summary' in x}).text.strip()
            published_date = row.find('span', {'data-testid': lambda x: x and 'vuln-published-on' in x}).text.strip()
            
            cvss_scores = {
                '4.0': None,
                '3.x': None,
                '2.0': None
            }
            
            cvss4_span = row.find('span', {'data-testid': lambda x: x and 'vuln-cvss4-na' not in x and 'vuln-cvss4-link' in x})
            if cvss4_span:
                cvss_scores['4.0'] = cvss4_span.text.strip()

            cvss3_span = row.find('a', {'data-testid': lambda x: x and 'vuln-cvss3-link' in x})
            if cvss3_span:
                cvss_scores['3.x'] = cvss3_span.text.strip()

            cvss2_span = row.find('span', {'data-testid': lambda x: x and 'vuln-cvss2-na' not in x and 'vuln-cvss2-link' in x})
            if cvss2_span:
                cvss_scores['2.0'] = cvss2_span.text.strip()
            
            cvss_score = "Not available"
            if cvss_scores['4.0'] and 'not available' not in cvss_scores['4.0'].lower():
                cvss_score = cvss_scores['4.0']
            elif cvss_scores['3.x'] and 'not available' not in cvss_scores['3.x'].lower():
                cvss_score = cvss_scores['3.x']
            elif cvss_scores['2.0'] and 'not available' not in cvss_scores['2.0'].lower():
                cvss_score = cvss_scores['2.0']

            cve_data = CVE(
                cve_id=cve_id,
                description=description,
                published_date=published_date,
                cvss_score=cvss_score
            )

            cves.append(cve_data)

        next_button = soup.find('a', {'data-testid': 'pagination-link-page->'})
        if not next_button or 'disabled' in next_button.attrs.get('class', []):
            break

        page += 1
        print(f"Page {page} done")

    
    return [cve.dict() for cve in cves]