import requests
import json
import csv
import datetime
import os
import time

def get_nist_vulnerabilities(days=3, limit=5000):  # Added limit parameter
    """Retrieves the latest vulnerabilities from the NIST NVD API v2.0 with a limit."""

    end_date = datetime.date.today()
    start_date = end_date - datetime.timedelta(days=days)

    all_vulnerabilities = []
    startIndex = 0
    resultsPerPage = 2000

    while True:
        url = f"https://services.nvd.nist.gov/rest/json/cves/2.0?startIndex={startIndex}&resultsPerPage={resultsPerPage}"

        try:
            response = requests.get(url)
            response.raise_for_status()
            data = response.json()
            vulnerabilities = data.get("vulnerabilities", [])

            if not vulnerabilities:
                break

            all_vulnerabilities.extend(vulnerabilities)
            startIndex += resultsPerPage
            time.sleep(0.5)  # Reduced sleep time slightly

            if len(all_vulnerabilities) >= limit: #Enforce limit
                print(f"Reached CVE limit of {limit}. Stopping retrieval.")
                break

        except requests.exceptions.RequestException as e:
            print(f"Error fetching data from NIST API: {e}")
            return []
        except json.JSONDecodeError as e:
            print(f"Error decoding JSON response: {e}")
            return []

    filtered_vulnerabilities = [
        cve for cve in all_vulnerabilities
        if datetime.datetime.fromisoformat(cve.get('cve', {}).get('published', '').replace('Z', '+00:00')).date() >= start_date
    ]
    return filtered_vulnerabilities



def write_to_csv(vulnerabilities, filename):
    # ... (This function is exactly the same as in my previous responses)
    if not vulnerabilities:
        print("No vulnerabilities found for the specified period.")
        return

    fieldnames = ["Product", "Vendor", "CVE ID", "Date Added", "Date Updated", "Description", "Severity"]

    try:
        with open(filename, 'w', newline='', encoding='utf-8') as csvfile:
            writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
            writer.writeheader()
            for cve in vulnerabilities:
                cve_data = cve.get('cve', {})
                
                products = []
                vendors = []

                if 'cpeMatch' in cve_data.get('configurations', {}):
                  for match in cve_data['configurations']['cpeMatch']:
                    cpe_string = match.get('cpe23Uri', '')
                    parts = cpe_string.split(':')
                    if len(parts) >= 5:
                        vendors.append(parts[3])
                        products.append(parts[4])
                    else:
                        vendors.append("N/A")
                        products.append("N/A")

                product_str = ", ".join(products)
                vendor_str = ", ".join(vendors)

                description = cve_data.get('descriptions', [{}])[0].get('value', 'N/A')
                published_date = cve_data.get('published', 'N/A')
                last_modified_date = cve_data.get('lastModified', 'N/A')
                cve_id = cve_data.get('id', 'N/A')

                severity = 'N/A'
                metrics = cve_data.get('metrics', {})
                if 'cvssMetricV31' in metrics:
                    severity = metrics['cvssMetricV31'][0].get('cvssData', {}).get('baseSeverity', 'N/A')
                elif 'cvssMetricV30' in metrics:
                    severity = metrics['cvssMetricV30'][0].get('cvssData', {}).get('baseSeverity', 'N/A')
                elif 'cvssMetricV2' in metrics:
                    severity = metrics['cvssMetricV2'][0].get('cvssData', {}).get('baseSeverity', 'N/A')

                writer.writerow({
                    "Product": product_str,
                    "Vendor": vendor_str,
                    "CVE ID": cve_id,
                    "Date Added": published_date,
                    "Date Updated": last_modified_date,
                    "Description": description,
                    "Severity": severity
                })
        print(f"Vulnerabilities written to {filename}")

    except Exception as e:
        print(f"Error writing to CSV: {e}")



if __name__ == "__main__":
    vulnerabilities = get_nist_vulnerabilities()

    if vulnerabilities:
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"nist_vulnerabilities_{timestamp}.csv"
        write_to_csv(vulnerabilities, filename)