app.py

import io
import webbrowser
from flask import Flask, request, render_template, send_file, session
from reportlab.lib.pagesizes import letter
from reportlab.lib.styles import getSampleStyleSheet
from reportlab.lib import colors
from reportlab.platypus import SimpleDocTemplate, Table, TableStyle, Paragraph
from reportlab.lib.units import inch
from reportlab.lib.enums import TA_LEFT
from reportlab.lib.styles import ParagraphStyle
from docx import Document
from docx.shared import Pt
from docx.oxml.ns import nsdecls
from docx.oxml import parse_xml
from docx.shared import RGBColor

# Custom imports
from scraper import scrape_nvd, scrape_cveorg, scrape_cvedetails, scrape_exploitdb
from helpers import is_valid_cve

# App initialization
app = Flask(__name__)
app.secret_key = "your_secret_key"

@app.route("/")
def home():
    """
    RS
    Route for the home page of the web application.

    This function handles requests to the root URL ("/") and returns the
    rendered 'index.html' template, which serves as the home page of the site.
    """
    return render_template("index.html")


@app.route("/cve", methods=["POST"])
def cve_checker():
    """
    AA
    Route for the CVE checker page of the web application.

    This function handles POST requests to the "/cve" URL and renders the
    'searchresult.html' template with the scraped data from NVD, CVE.org,
    CVE Details and ExploitDB. If the CVE ID is not valid, it renders the
    same template with an error message.

    Args:
        request.form["cve_id"]: The CVE ID to be searched.

    Returns:
        Rendered 'searchresult.html' template with scraped data or an error
        message if the CVE ID is not valid.
    """

    # AA. Get the CVE ID from the form and make it uppercase
    cve_id = request.form["cve_id"].upper()

    # AA. Check if the CVE ID is valid
    cve_passed_regex = is_valid_cve(cve_id)

    # AA. If the CVE ID is valid, scrape the data from the scrapers otherwise render the same template with an error message
    if cve_passed_regex:
        nvd_data = scrape_nvd(cve_id)
        cveorg_data = scrape_cveorg(cve_id)
        cvedetails_data = scrape_cvedetails(cve_id)
        exploits = scrape_exploitdb(cve_id)

        # AA. Combine the scraped data into a single dictionary
        cve_data = {**nvd_data, **cveorg_data, **cvedetails_data}

        # RS, SA. Store the scraped data in the session
        session["cve_data"] = cve_data

        # AA. Open the CVE website in the default web browser
        websites = [
            "https://nvd.nist.gov/vuln/detail/"+cve_id,
            "https://www.cve.org/CVERecord?id="+cve_id,
            "https://www.cvedetails.com/cve/"+cve_id,
            "https://cve.mitre.org/cgi-bin/cvename.cgi?name="+cve_id,
            "https://www.tenable.com/cve/"+cve_id,
        ]

        # AA. Open the websites in the default web browser
        for url in websites:
            webbrowser.open(url)

        # AA. Render the 'searchresult.html' template with the scraped data
        return render_template(
            "searchresult.html", cve_data=cve_data, exploits=exploits
        )
    else:
        # AA. If the CVE ID is not valid, render the same template with an error message
        return render_template(
            "searchresult.html", cve_data={"error": "CVE not found or wrong CVE format"}
        )


@app.route("/download/pdf")
def download_pdf():
    """
    RS
    Download CVE report as PDF.

    This function generates a PDF report based on the stored CVE data in the
    session and sends it as a downloadable file. If no CVE data is available in
    the session, it returns an error message.

    Returns:
        A PDF file containing the CVE report, or an error message if no CVE
        data is available in the session.

    """

    # RS, SA.Get the CVE data from the session
    cve_data = session.get("cve_data")

    # If no CVE data is available in the session, return an error message
    if not cve_data:
        return "No CVE data available to generate the PDF.", 400
    
    # Generate the PDF report
    pdf_buffer = io.BytesIO()
    generate_pdf(cve_data, pdf_buffer)
    pdf_buffer.seek(0)

    # Send the PDF file as a downloadable file
    return send_file(
        pdf_buffer,
        as_attachment=True,
        download_name="cve_report.pdf",
        mimetype="application/pdf",
    )


def generate_pdf(cve_data, buffer):
    """
    RS
    Generate a PDF report from the given CVE data and write it to the given file-like object.

    Args:
        cve_data (dict): A dictionary containing the CVE data to be included in the report.
        buffer (file-like object): A file-like object to which the generated PDF report will be written.

    Returns:
        None
    """

    # Create the PDF document
    doc = SimpleDocTemplate(buffer, pagesize=letter)
    
    # Define the elements to be included in the report
    elements = []

    # Add a title
    styles = getSampleStyleSheet()
    title = Paragraph("CVE Report", styles["Title"])
    elements.append(title)

    # Define a custom paragraph style for table content
    custom_style = ParagraphStyle(
        name="Custom",
        fontSize=10,
        leading=12,
        alignment=TA_LEFT,
        wordWrap="CJK",  # Enables text wrapping
    )

    # Prepare data for the table
    table_data = [["Key", "Value"]]  # Table header

    # Iterate over the cve_data dictionary
    for key, value in cve_data.items():
        if key in ["cvss_color", "status_color"]:
            continue  # Skip these fields

        # Format the key: uppercase and replace underscores with spaces
        formatted_key = key.upper().replace("_", " ")

        # Check if value needs to be a Paragraph
        if isinstance(value, str):
            if len(value) > 50:  # Adjust threshold as needed
                value = Paragraph(value, custom_style)
            else:
                value = Paragraph(value, custom_style)
        elif isinstance(value, dict):
            # Debugging: Print out the type of value if it's a dictionary
            print(f"Debug: Value for key '{key}' is a dictionary: {value}")
            value = Paragraph(
                str(value), custom_style
            )  # Convert dict to string and then to Paragraph
        else:
            # Handle other types of values if needed
            print(f"Debug: Value for key '{key}' is of type {type(value)}")
            value = Paragraph(str(value), custom_style)
        table_data.append([formatted_key, value])

    # Create a table with style
    table = Table(
        table_data, colWidths=[2.5 * inch, 4.5 * inch]
    )  # Adjust column widths as needed
    table.setStyle(
        TableStyle(
            [
                ("BACKGROUND", (0, 0), (-1, 0), colors.grey),
                ("TEXTCOLOR", (0, 0), (-1, 0), colors.whitesmoke),
                ("ALIGN", (0, 0), (-1, -1), "CENTER"),
                ("FONTNAME", (0, 0), (-1, 0), "Helvetica-Bold"),
                ("FONTSIZE", (0, 0), (-1, 0), 14),
                ("BOTTOMPADDING", (0, 0), (-1, 0), 12),
                ("BACKGROUND", (0, 1), (-1, -1), colors.beige),
                ("GRID", (0, 0), (-1, -1), 1, colors.black),
                ("ALIGN", (0, 1), (-1, -1), "LEFT"),
            ]
        )
    )

    # Add the table to the elements
    elements.append(table)

    # Build the PDF document
    doc.build(elements)


@app.route("/download/word")
def download_word():
    """
    RS
    Download CVE report as a Word document.

    This function generates a Word document based on the stored CVE data in the
    session and sends it as a downloadable file. If no CVE data is available in
    the session, it returns an error message.

    Returns:
        A Word file containing the CVE report, or an error message if no CVE
        data is available in the session.

    """
    cve_data = session.get("cve_data")

    # Check if CVE data exists
    if not cve_data:
        return "No CVE data available to generate the Word document.", 400

    # Create an in-memory buffer
    word_buffer = io.BytesIO()

    # Generate the Word document and write it to the buffer
    generate_word(cve_data, word_buffer)

    # Move the pointer to the beginning of the buffer
    word_buffer.seek(0)

    # Send the buffer as a downloadable Word file
    return send_file(
        word_buffer,
        as_attachment=True,
        download_name="cve_report.docx",
        mimetype="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
    )


def generate_word(cve_data, buffer):
    """
    RS
    Generate a Word document from the given CVE data and write it to the given file-like object.

    Args:
        cve_data (dict): A dictionary containing the CVE data to be included in the report.
        buffer (file-like object): A file-like object to which the generated Word document will be written.

    Returns:
        None
    """

    # Create a new Word document
    doc = Document()

    # Add a title with styling
    title = doc.add_heading("CVE Report", level=1)
    title_format = title.paragraph_format
    title_format.space_after = Pt(12)

    # Style for key headings (formatted key names)
    key_style = doc.styles.add_style("KeyStyle", 1)
    key_font = key_style.font
    key_font.bold = True
    key_font.size = Pt(12)
    key_font.color.rgb = RGBColor(0, 102, 204)  # Dark blue color

    # Style for values
    value_style = doc.styles.add_style("ValueStyle", 1)
    value_font = value_style.font
    value_font.size = Pt(11)

    # Add key-value pairs to the Word document with formatting
    for key, value in cve_data.items():
        if key in ["cvss_color", "status_color"]:
            continue  # Skip these fields

        # Format the key: uppercase and replace underscores with spaces
        formatted_key = key.upper().replace("_", " ")

        # Add the key with the custom style
        key_paragraph = doc.add_paragraph(f"{formatted_key}:", style="KeyStyle")

        # Add the value with the custom style
        value_paragraph = doc.add_paragraph(str(value), style="ValueStyle")

        # Add space after each entry for better readability
        key_paragraph.paragraph_format.space_after = Pt(4)
        value_paragraph.paragraph_format.space_after = Pt(8)

        # Alternate background color for better readability
        if len(doc.paragraphs) % 2 == 0:
            shading_elm = parse_xml(r'<w:shd {} w:fill="F2F2F2"/>'.format(nsdecls("w")))
            key_paragraph._p.get_or_add_pPr().append(shading_elm)
            value_paragraph._p.get_or_add_pPr().append(shading_elm)

    # Save the document to the buffer
    doc.save(buffer)


@app.route("/download/md")
def download_md():
    """
    SA
    Download CVE report as Markdown.

    This function generates a Markdown report based on the stored CVE data in the
    session and sends it as a downloadable file. If no CVE data is available in
    the session, it returns an error message.

    Returns:
        A Markdown file containing the CVE report, or an error message if no CVE
        data is available in the session.

    """
    cve_data = session.get("cve_data")

    # Check if CVE data exists
    if not cve_data:
        return "No CVE data available to generate the Markdown report.", 400

    # Create an in-memory buffer for Markdown
    md_buffer = io.BytesIO()

    # Generate the Markdown content and write it to the buffer
    generate_md(cve_data, md_buffer)

    # Move the pointer to the beginning of the buffer
    md_buffer.seek(0)

    # Send the buffer as a downloadable Markdown file
    return send_file(
        md_buffer,
        as_attachment=True,
        download_name="cve_report.md",
        mimetype="text/markdown",
    )


def generate_md(cve_data, buffer):
    """
    SA
    Generate a Markdown report from the given CVE data and write it to the given file-like object.

    Args:
        cve_data (dict): A dictionary containing the CVE data to be included in the report.
        buffer (file-like object): A file-like object to which the generated Markdown report will be written.

    Returns:
        None
    """
    report_lines = ["# CVE Report\n\n"]

    # Add key-value pairs to the report
    for key, value in cve_data.items():
        if key in ["cvss_color", "status_color"]:
            continue

        formatted_key = key.upper().replace("_", " ")
        report_lines.append(f"## {formatted_key}\n")
        report_lines.append(f"{value}\n\n")

    # Join the list into a string and write to the buffer
    report_content = "".join(report_lines)
    buffer.write(report_content.encode("utf-8"))


@app.route("/download/html")
def download_html():
    """
    SA
    Download CVE report as HTML.

    This function generates an HTML report based on the stored CVE data in the
    session and sends it as a downloadable file. If no CVE data is available in
    the session, it returns an error message.

    Returns:
        An HTML file containing the CVE report, or an error message if no CVE
        data is available in the session.

    """
    cve_data = session.get("cve_data")

    # Check if CVE data exists
    if not cve_data:
        return "No CVE data available to generate the HTML report.", 400

    # Create an in-memory buffer for HTML
    html_buffer = io.BytesIO()

    # Generate the HTML content and write it to the buffer
    generate_html(cve_data, html_buffer)

    # Move the pointer to the beginning of the buffer
    html_buffer.seek(0)

    # Send the buffer as a downloadable HTML file
    return send_file(
        html_buffer,
        as_attachment=True,
        download_name="cve_report.html",
        mimetype="text/html",
    )


def generate_html(cve_data, buffer):
    """
    SA
    Generate an HTML report from the given CVE data and write it to the given file-like object.

    Args:
        cve_data (dict): A dictionary containing the CVE data to be included in the report.
        buffer (file-like object): A file-like object to which the generated HTML report will be written.

    Returns:
        None
    """

    # Generate the HTML content
    html_lines = ["<html><head><title>CVE Report</title></head><body>"]
    html_lines.append("<h1>CVE Report</h1>")

    # Add key-value pairs to the report
    for key, value in cve_data.items():
        if key in ["cvss_color", "status_color"]:
            continue

        formatted_key = key.upper().replace("_", " ")
        html_lines.append(f"<h2>{formatted_key}</h2>")
        html_lines.append(f"<p>{value}</p>")

    # Close the HTML body
    html_lines.append("</body></html>")

    # Join the list into a string and write to the buffer
    html_content = "".join(html_lines)
    buffer.write(html_content.encode("utf-8"))

# RS. Run the Flask app
if __name__ == "__main__":
    app.run(debug=True)