app.py

from flask import Flask, render_template, request, redirect, url_for, flash, send_file, session
from dotenv import load_dotenv
import downloader
import requests
from validator import is_valid_cve_format, cve_exists, is_url_accessible
from fetch_exploit import fetch_exploit_db_links
import os
import webbrowser

load_dotenv()

app = Flask(__name__)
app.secret_key = os.getenv('secret_key')
NIST_API_KEY = os.getenv('NIST_API_KEY')


@app.route('/', methods=['GET', 'POST'])
def index():
    if request.method == 'POST':
        cve_id = request.form['cve_id'].upper()
        if cve_id:
            if not is_valid_cve_format(cve_id):
                flash("Invalid CVE ID format", "warning")
            else:
                status = cve_exists(cve_id)
                if status == True:
                    return redirect(url_for('result', cve_id=cve_id))
                else:
                    flash(f"CVE ID status: {status}", "warning")
        else:
            flash("Please enter a CVE ID", "warning")
    return render_template('index.html')

@app.route('/result/<cve_id>')
def result(cve_id):
    # API URL
    nist_api_url = f"https://services.nvd.nist.gov/rest/json/cves/2.0?cveId={cve_id}"

    mitre_api_url = f"https://cveawg.mitre.org/api/cve/{cve_id}"

    # Add the API key to the request headers
    headers = {
        'apiKey': NIST_API_KEY
    }

    # Fetch details from NIST API with headers
    nist_response = requests.get(nist_api_url, headers=headers).json()

    mitre_response = requests.get(mitre_api_url).json()

    # Check if any results were found
    if nist_response.get("totalResults", 0) == 0:
        flash(f"No information found for CVE ID {cve_id}", "warning")
        return redirect(url_for('index'))

    # Extract the relevant information
    cve_data = nist_response['vulnerabilities'][0]['cve']
    description = cve_data['descriptions'][0]['value']
    # Handle CVSS metrics for both V31 and V21
    cvss_data = None
    score = 'N/A'
    cvss_severity = 'N/A'
    cvss_vector = 'N/A'

    if 'cvssMetricV31' in cve_data['metrics']:
        # CVSS v3.1
        cvss_data = cve_data['metrics']['cvssMetricV31'][0]['cvssData']
        score = cvss_data['baseScore']
        cvss_severity = cvss_data['baseSeverity']
        cvss_vector = cvss_data['vectorString']
    elif 'cvssMetricV30' in cve_data['metrics']:
        # CVSS v2.1
        cvss_data = cve_data['metrics']['cvssMetricV30'][0]['cvssData']
        score = cvss_data['baseScore']
        cvss_severity = cvss_data['baseSeverity']
        cvss_vector = cvss_data['vectorString']
    elif 'cvssMetricV21' in cve_data['metrics']:
        # CVSS v2.1
        cvss_data = cve_data['metrics']['cvssMetricV21'][0]['cvssData']
        score = cvss_data['baseScore']
        cvss_severity = cvss_data['baseSeverity']
        cvss_vector = cvss_data['vectorString']
    elif 'cvssMetricV2' in cve_data['metrics']:
        # CVSS v2.0
        cvss_data = cve_data['metrics']['cvssMetricV2'][0]['cvssData']
        score = cvss_data['baseScore']
        cvss_severity = cve_data['metrics']['cvssMetricV2'][0]['baseSeverity']
        cvss_vector = cvss_data['vectorString']
    
    # Fetch and filter references (excluding broken ones)
    references = [
        ref['url'] for ref in cve_data['references']
        if 'Broken Link' not in ref.get('tags', [])
        and is_url_accessible(ref['url'])
    ]

    for url in references[:3]:
        webbrowser.open(url)

    # Fetch exploit links from Exploit-DB
    exploit_links = fetch_exploit_db_links(cve_id)

    affected_vendors = mitre_response.get("containers", {}).get("cna", {}).get("affected", [])

    session['cve_data'] = {
        'cve_id': cve_id,
        'description': description,
        'score': score,
        'severity': cvss_severity,
        'vector' : cvss_vector,
        'exploit_links': exploit_links,
        'references': references,
        'vendors' : affected_vendors
    }

    return render_template('result.html', 
                           cve_id=cve_id, 
                           description=description, 
                           score=score, 
                           severity=cvss_severity,
                           vector = cvss_vector,
                           exploit_links=exploit_links,
                           references=references,
                           vendors = affected_vendors
                           )

@app.route('/download_report/<cve_id>', methods=['POST'])
def download_report(cve_id):
    file_type = request.form['file_type']
    
    # Check if CVE data exists in session
    if 'cve_data' not in session or session['cve_data']['cve_id'] != cve_id:
        flash(f"Session expired or no data found for CVE ID {cve_id}. Please search again.", "warning")
        return redirect(url_for('index'))
    
    # Retrieve CVE data from session
    cve_data = session['cve_data']

    # Generate the report using the downloader
    file_path = downloader.generate_report(cve_id, cve_data, file_type)
    
    # Send the file as a response to the client
    return send_file(file_path, as_attachment=True)

if __name__ == '__main__':
    import threading
    url = 'http://127.0.0.1:5000/'

    webbrowser.open(url)
    app.run(debug=True)