html_report_builder.py

import json
import os
from typing import Any

from jinja2 import Template, FileSystemLoader, Environment

from contrib.descriptions import VulnDescriptionProvider
from contrib.report_builders import JsonReportBuilder


class ReportInfo:
    def __init__(self, severity='Low', is_empty=True):
        self.level = severity
        self.num_level = 0

        self.empty = is_empty

        if severity == 'Medium':
            self.num_level = 1

        if severity == 'High':
            self.num_level = 2

    def toJSON(self):
        return json.dumps(self, default=lambda o: o.__dict__,
                          sort_keys=True, indent=4)

    def __lt__(self, other):
        return self.num_level < other.num_level


class JinjaHtmlReportBuilder(JsonReportBuilder):
    def __init__(self, description_provider: VulnDescriptionProvider):
        super().__init__(description_provider)
        self.template_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'templates')
        self.template_name = 'jinja2_report.html'
        self._read_template()  # type: Template

    def build(self) -> Any:
        json_data = self._buffer

        max_severity = ReportInfo()
        for vv in json_data['vulnerable'].values():
            for v in vv['vulnerabilities']:
                s = ReportInfo(v['severity_str'], False)

                if s > max_severity:
                    max_severity = s

        with open(f'{json_data["output_path"]}.info', 'w+') as severity:
            severity.write(max_severity.toJSON())

        return self._template.render(data=json_data)

    def _read_template(self):
        template_loader = FileSystemLoader(searchpath=self.template_path)
        template_env = Environment(loader=template_loader, autoescape=True)
        self._template = template_env.get_template(self.template_name)