Report.py

import NVD
import Exploit

from docx import Document
from docx.enum.text import WD_PARAGRAPH_ALIGNMENT
from docx.shared import Pt

from reportlab.lib.pagesizes import letter
from reportlab.platypus import SimpleDocTemplate, Paragraph
from reportlab.lib.styles import getSampleStyleSheet


# This function creates report in word format
def create_docx(cve_id):
    summary, cvss_score, vector, reference = NVD.get_info(cve_id)
    exploit_number, exploit_header, exploit_script = Exploit.check_exploit_db(cve_id)

    doc = Document()

    # Title
    title = doc.add_heading('CVE Information', level=1)
    title.alignment = WD_PARAGRAPH_ALIGNMENT.CENTER

    # CVE Information
    cve_paragraph = doc.add_paragraph()
    cve_paragraph.add_run("CVE ID: ").bold = True
    cve_paragraph.add_run(f"{cve_id}\n\n")

    cve_paragraph.add_run("Summary: ").bold = True
    cve_paragraph.add_run(f"{summary}\n\n")

    cve_paragraph.add_run("Base Score: ").bold = True
    cve_paragraph.add_run(f"{cvss_score}\n\n")

    cve_paragraph.add_run("Vector: ").bold = True
    cve_paragraph.add_run(f"{vector}\n\n")

    cve_paragraph.add_run("Reference: ").bold = True
    cve_paragraph.add_run(f"{reference}\n\n")

    # Add spacing
    cve_paragraph.add_run("\n" * 2)

    # Exploit Information
    exploit_paragraph = doc.add_paragraph()
    exploit_paragraph.add_run("Number of exploit scripts found: ").bold = True
    exploit_paragraph.add_run(f"{exploit_number}\n\n")
    exploit_paragraph.add_run(f"\nExploit Header: \n\n")
    exploit_paragraph.add_run(f"{exploit_header}\n\n")

    exploit_script_lines = exploit_script.split('\n')
    for line in exploit_script_lines:
        exploit_paragraph.add_run(line).font.size = Pt(10)
        exploit_paragraph.add_run('\n')

    # Save the document
    docx_filename = f"report_{cve_id}.docx"
    doc.save(docx_filename)

    return docx_filename


# This function creates report in pdf format
def create_pdf(cve_id):
    summary, cvss_score, vector, reference = NVD.get_info(cve_id)
    exploit_number, exploit_header, exploit_script = Exploit.check_exploit_db(cve_id)

    pdf_file_path = f'report_{cve_id}.pdf'
    doc = SimpleDocTemplate(pdf_file_path, pagesize=letter)

    story = []

    styles = getSampleStyleSheet()
    style_heading = styles['Heading1']
    style_normal = styles['Normal']

    story.append(Paragraph('CVE Information', style_heading))
    story.append(Paragraph(f'<b>CVE ID:</b> {cve_id}\n', style_normal))
    story.append(Paragraph(f'<b>Summary:</b> {summary}\n', style_normal))
    story.append(Paragraph(f'<b>Base Score:</b> {cvss_score}\n', style_normal))
    story.append(Paragraph(f'<b>Vector:</b> {vector}\n', style_normal))
    story.append(Paragraph(f'<b>Reference:</b> <a color="blue" href="{reference}">{reference}</a>\n', style_normal))

    story.append(Paragraph('<br/><br/>', style_normal))

    story.append(Paragraph('Exploit Information', style_heading))
    story.append(Paragraph(f'<b>Number of exploit scripts found:</b> {exploit_number}', style_normal))
    story.append(Paragraph(f'<b>Exploit Header:</b> {exploit_header}', style_normal))
    story.append(Paragraph(f'<b>Exploit Script:</b>', style_normal))

    exploit_script_lines = exploit_script.split('\n')
    for line in exploit_script_lines:
        story.append(Paragraph(f'<font face="Courier">{line}</font>', style_normal))

    doc.build(story)

    return pdf_file_path


# This function creates report in md format
def create_md(cve_id):
    summary, cvss_score, vector, reference = NVD.get_info(cve_id)
    exploit_number, exploit_header, exploit_script = Exploit.check_exploit_db(cve_id)

    md_content = ['## CVE Information', f'**CVE ID:** {cve_id}', f'**Summary:** {summary}',
                  f'**Base Score:** {cvss_score}', f'**Vector:** {vector}',
                  f'**Reference:** [{reference}]({reference})\n', '## Exploit Information',
                  f'**Number of exploit scripts found:** {exploit_number}', f'**Exploit Header:** {exploit_header}',
                  '**Exploit Script:**']

    exploit_script_lines = exploit_script.split('\n')
    md_content.extend(['```', *exploit_script_lines, '```'])

    md_filename = f"report_{cve_id}.md"
    with open(md_filename, 'w', encoding='utf-8') as md_file:
        md_file.write('\n'.join(md_content))

    return md_filename