KasperskyLab-TinyCheck/analysis/classes/report.py

import weasyprint
import os
import json
import hashlib

from weasyprint import HTML
from pathlib import Path
from datetime import datetime


class Report(object):

    def __init__(self, capture_directory):
        self.capture_directory = capture_directory
        self.alerts = self.read_json(os.path.join(
            capture_directory, "assets/alerts.json"))
        self.whitelist = self.read_json(os.path.join(
            capture_directory, "assets/whitelist.json"))
        self.conns = self.read_json(os.path.join(
            capture_directory, "assets/conns.json"))
        self.device = self.read_json(os.path.join(
            capture_directory, "assets/device.json"))
        self.capinfos = self.read_json(os.path.join(
            capture_directory, "assets/capinfos.json"))
        try:
            with open(os.path.join(self.capture_directory, "capture.pcap"), "rb") as f:
                self.capture_sha1 = hashlib.sha1(f.read()).hexdigest()
        except:
            self.capture_sha1 = "N/A"

    def read_json(self, json_path):
        """
            Read and convert a JSON file.
            :return: array or dict.
        """
        with open(json_path, "r") as json_file:
            return json.load(json_file)

    def generate_report(self):
        """
            Generate the full report in PDF
            :return: nothing
        """
        content = self.generate_page_header()
        content += self.generate_header()
        content += self.generate_warning()
        content += self.generate_alerts()
        content += self.generate_suspect_conns_block()
        content += self.generate_uncat_conns_block()
        content += self.generate_whitelist_block()

        htmldoc = HTML(string=content, base_url="").write_pdf()
        Path(os.path.join(self.capture_directory,
                          "report.pdf")).write_bytes(htmldoc)

    def generate_warning(self):
        """
            Generate the warning message.
            :return: str
        """
        if len(self.alerts["high"]):
            return "<div class=\"warning high\">Your device seems to be compromised as you have {} high alerts.</div>".format(self.nb_translate(len(self.alerts["high"])))
        elif len(self.alerts["moderate"]):
            return "<div class=\"warning moderate\">You have {} moderate alerts, your device might be compromised. Please look at them carefully.</div>".format(self.nb_translate(len(self.alerts["moderate"])))
        elif len(self.alerts["low"]):
            return "<div class=\"warning low\">You have only {} low alerts, don't hesitate to check them.</div>".format(self.nb_translate(len(self.alerts["low"])))
        else:
            return "<div class=\"warning low\">Everything looks fine, zero alerts. Don't hesitate to check the uncategorized communications, if any.</div>"

    def nb_translate(self, nb):
        """
            Translate a number in a string.
            :return: str
        """
        a = ["one", "two", "three", "four", "five",
             "six", "seven", "height", "nine"]
        return a[nb-1] if nb <= 9 else str(nb)

    def generate_suspect_conns_block(self):
        """
            Generate the table of the network non-whitelisted communications.
            :return: string
        """

        if not len([c for c in self.conns if c["alert_tiggered"] == True]):
            return ""

        title = "<h2>Suspect communications</h2>"
        table = """<table>
                       <thead>
                           <tr>
                               <th>Protocol</th>
                               <th>Domain</th>
                               <th>Dst IP Address</th>
                              <th>Dst port</th>
                           </tr>
                       </thead>
                <tbody>"""

        for rec in self.conns:
            if rec["alert_tiggered"] == True:
                table += "<tr>"
                table += "<td>{}</td>".format(rec["proto"].upper())
                table += "<td>{}</td>".format(rec["resolution"]
                                              if rec["resolution"] != rec["ip_dst"] else "--")
                table += "<td>{}</td>".format(rec["ip_dst"])
                table += "<td>{}</td>".format(rec["port_dst"])
                table += "</tr>"
        table += "</tbody></table>"
        return title + table

    def generate_uncat_conns_block(self):
        """
            Generate the table of the network non-whitelisted communications.
            :return: string
        """
        if not len([c for c in self.conns if c["alert_tiggered"] == False]):
            return ""

        title = "<h2>Uncategorized communications</h2>"
        table = """<table>
                       <thead>
                           <tr>
                               <th>Protocol</th>
                               <th>Domain</th>
                               <th>Dst IP Address</th>
                              <th>Dst port</th>
                           </tr>
                       </thead>
                <tbody>"""

        for rec in self.conns:
            if rec["alert_tiggered"] == False:
                table += "<tr>"
                table += "<td>{}</td>".format(rec["proto"].upper())
                table += "<td>{}</td>".format(rec["resolution"]
                                              if rec["resolution"] != rec["ip_dst"] else "--")
                table += "<td>{}</td>".format(rec["ip_dst"])
                table += "<td>{}</td>".format(rec["port_dst"])
                table += "</tr>"
        table += "</tbody></table>"
        return title + table

    def generate_whitelist_block(self):
        """
            Generate the table of the whitelisted communications.
            :return: string
        """
        if not len(self.whitelist):
            return ""

        title = "<h2>Whitelisted communications</h2>"
        table = """<table>
                    <thead>
                            <tr>
                                <th>Protocol</th>
                                <th>Domain</th>
                                <th>Dst IP Address</th>
                                <th>Dst port</th>
                            </tr>
                    </thead>
                <tbody>"""

        for rec in sorted(self.whitelist, key=lambda k: k['resolution']):
            table += "<tr>"
            table += "<td>{}</td>".format(rec["proto"].upper())
            table += "<td>{}</td>".format(rec["resolution"]
                                          if rec["resolution"] != rec["ip_dst"] else "--")
            table += "<td>{}</td>".format(rec["ip_dst"])
            table += "<td>{}</td>".format(rec["port_dst"])
            table += "</tr>"
        table += "</tbody></table>"
        return title + table

    def generate_header(self):
        """
            Generate the report header with context data.
            :return: string
        """
        header = "<div class=\"header\">"
        header += "<div class=\"logo\"></div>"
        header += "<p><br /><strong>Device name: {}</strong><br />".format(
            self.device["name"])
        header += "Device MAC address: {}<br />".format(
            self.device["mac_address"])
        header += "Report generated on {}<br />".format(
            datetime.now().strftime("%d/%m/%Y at %H:%M:%S"))
        header += "Capture duration: {}<br />".format(
            self.capinfos["Capture duration"])
        header += "Number of packets: {}<br />".format(
            self.capinfos["Number of packets"])
        header += "Capture SHA1: {}<br />".format(self.capture_sha1)
        header += "</p>"
        header += "</div>"
        return header

    def generate_alerts(self):
        """
            Generate the alerts.
            :return: string
        """
        alerts = "<ul class=\"alerts\">"
        for alert in self.alerts["high"]:
            alerts += "<li class =\"alert\">"
            alerts += "<span class=\"high-label\">High</span>"
            alerts += "<span class=\"alert-id\">{}</span>".format(alert["id"])
            alerts += "<div class = \"alert-body\">"
            alerts += "<span class=\"title\">{}</span>".format(alert["title"])
            alerts += "<p class=\"description\">{}</p>".format(
                alert["description"])
            alerts += "</div>"
            alerts += "</li>"

        for alert in self.alerts["moderate"]:
            alerts += "<li class =\"alert\">"
            alerts += "<span class=\"moderate-label\">moderate</span>"
            alerts += "<span class=\"alert-id\">{}</span>".format(alert["id"])
            alerts += "<div class = \"alert-body\">"
            alerts += "<span class=\"title\">{}</span>".format(alert["title"])
            alerts += "<p class=\"description\">{}</p>".format(
                alert["description"])
            alerts += "</div>"
            alerts += "</li>"
        for alert in self.alerts["low"]:
            alerts += "<li class =\"alert\">"
            alerts += "<span class=\"low-label\">low</span>"
            alerts += "<span class=\"alert-id\">{}</span>".format(alert["id"])
            alerts += "<div class = \"alert-body\">"
            alerts += "<span class=\"title\">{}</span>".format(alert["title"])
            alerts += "<p class=\"description\">{}</p>".format(
                alert["description"])
            alerts += "</div>"
            alerts += "</li>"

        alerts += "</ul>"
        return alerts

    def generate_page_footer(self):
        """
            Generate the html footer.
            :return: string
        """
        return "</body></html>"

    def generate_page_header(self):
        """
            Generate the html header.
            :return: string
        """
        return """<html
                    <head>
                        <style>
                            * {
                                font-family: Arial, Helvetica, sans-serif;
                            }

                            h2 {
                                padding-top: 30px;
                                font-weight: 400;
                                font-size: 18px;
                            }

                            td {
                                width: auto;
                                padding: 10px;
                            }

                            table {
                                background: #FFF;
                                border: 2px solid #FAFAFA;
                                border-radius: 5px;
                                border-collapse: separate;
                                border-spacing: 0px;
                                width: 100%;
                                font-size: 12px;
                            }

                            p {
                                font-size: 13px;
                            }

                            thead tr th {
                                border-bottom: 1px solid #CCC;
                                border-collapse: separate;
                                border-spacing: 5px 5px;
                                background-color: #FFF;
                                padding: 10px;
                                text-align: left;
                            }

                            tbody tr#first td {
                                border-top: 3px solid #4d4d4d;
                                border-collapse: separate;
                                border-spacing: 5px 5px;
                            }

                            tr:nth-of-type(odd) {
                                background-color: #fafafa;
                            }

                            .logo {
                                background-image: url("data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAUAAAABkCAYAAAD32uk+AAAMYmlDQ1BJQ0MgUHJvZmlsZQAASImVVwdYU8kWnltSSWiBCEgJvYkiNYCUEFoEAekgKiEJJJQYE4KKHV1WwbWgIoplRVdFFF1dAVkLInYXxd4XCyor6+IqNlTehAR09ZXvne+bO3/OnPlPycy9MwDodPBlsjxUF4B8aYE8LjyYlZKaxiI9AiRgABjwacQXKGSc2NgoAGWw/6e8vgYQVX/ZRcX17fh/FX2hSCEAAEmHOFOoEORD3AwAXiyQyQsAIIZAvfW0ApkKiyE2kMMAIZ6lwtlqvFyFM9V464BNQhwX4kYAyDQ+X54NgHYr1LMKBdmQR/sRxK5SoUQKgI4BxAECMV8IcQLEI/Lzp6jwPIgdoL0M4h0QszO/4Mz+B3/mED+fnz2E1XkNCDlEopDl8Wf8n6X535Kfpxz0YQcbTSyPiFPlD2t4I3dKpArTIO6WZkbHqGoN8VuJUF13AFCqWBmRqLZHTQUKLqwfYELsKuSHREJsCnGYNC86SqPPzJKE8SCGqwWdLingJWjmLhIpQuM1nOvlU+JiBnGWnMvRzK3jywf8quxblbmJHA3/DbGIN8j/qkickAwxFQCMWihJioZYG2IDRW58pNoGsyoSc6MHbeTKOFX8NhCzRdLwYDU/lp4lD4vT2MvyFYP5YiViCS9agysLxAkR6vpgOwX8gfiNIK4XSTmJgzwiRUrUYC5CUUioOnesTSRN1OSL3ZMVBMdp5vbI8mI19jhZlBeu0ltBbKIojNfMxccUwMWp5sejZAWxCeo48Ywc/thYdTx4IYgCXBACWEAJWyaYAnKApK27oRv+Uo+EAT6Qg2wgAi4azeCM5IERKXzGgyLwJ0QioBiaFzwwKgKFUP9xSKt+uoCsgdHCgRm54DHE+SAS5MHfyoFZ0iFvSeAR1Ei+8S6AsebBphr7VseBmiiNRjnIy9IZtCSGEkOIEcQwoiNuggfgfngUfAbB5oazcZ/BaD/bEx4T2gkPCFcJHYSbkyXF8q9iGQc6IH+YJuPMLzPG7SCnJx6M+0N2yIwzcRPggntAPxw8EHr2hFquJm5V7qx/k+dQBl/UXGNHcaWglGGUIIrD1zO1nbQ9h1hUFf2yPupYM4eqyh0a+do/94s6C2Ef+bUltgjbj53CjmFnsENYA2BhR7FG7Dx2WIWH1tCjgTU06C1uIJ5cyCP5xh9f41NVSYVrrWuX6wfNGCgQTS9QbTDuFNkMuSRbXMDiwK+AiMWTCkaOYLm5urkCoPqmqF9TL5kD3wqEefazboElAP4z+vv7D33WRV4EYP9huM1vfdbZd8LXwVkATq8RKOWFah2uehDg20AH7ihjYA6sgQPMyA14AT8QBELBWBADEkAqmATrLIbrWQ6mgVlgPigBZWA5WA3WgU1gC9gBdoN9oAEcAsfASXAOXARXwW24fjrBM9ADXoM+BEFICB1hIMaIBWKLOCNuCBsJQEKRKCQOSUUykGxEiiiRWcgCpAwpR9Yhm5Ea5GfkIHIMOYO0IzeR+0gX8jfyHsVQGmqAmqF26CiUjXLQSDQBnYhmo1PRInQhuhStRKvRXWg9egw9h15FO9BnaC8GMC2MiVliLhgb42IxWBqWhcmxOVgpVoFVY3VYE/ynL2MdWDf2DifiDJyFu8A1HIEn4gJ8Kj4HX4Kvw3fg9Xgrfhm/j/fgnwh0ginBmeBL4BFSCNmEaYQSQgVhG+EA4QTcTZ2E10QikUm0J3rD3ZhKzCHOJC4hbiDuITYT24kPib0kEsmY5EzyJ8WQ+KQCUglpLWkX6SjpEqmT9JasRbYgu5HDyGlkKbmYXEHeST5CvkR+Qu6j6FJsKb6UGIqQMoOyjLKV0kS5QOmk9FH1qPZUf2oCNYc6n1pJraOeoN6hvtTS0rLS8tEaryXRmqdVqbVX67TWfa13NH2aE41LS6cpaUtp22nNtJu0l3Q63Y4eRE+jF9CX0mvox+n36G+1GdojtXnaQu252lXa9dqXtJ/rUHRsdTg6k3SKdCp09utc0OnWpeja6XJ1+bpzdKt0D+pe1+3VY+iN1ovRy9dbordT74zeU32Svp1+qL5Qf6H+Fv3j+g8ZGMOawWUIGAsYWxknGJ0GRAN7A55BjkGZwW6DNoMeQ31DD8Mkw+mGVYaHDTuYGNOOyWPmMZcx9zGvMd8PMxvGGSYatnhY3bBLw94YDTcKMhIZlRrtMbpq9N6YZRxqnGu8wrjB+K4JbuJkMt5kmslGkxMm3cMNhvsNFwwvHb5v+C1T1NTJNM50pukW0/OmvWbmZuFmMrO1ZsfNus2Z5kHmOearzI+Yd1kwLAIsJBarLI5a/MEyZHFYeaxKViurx9LUMsJSabnZss2yz8reKtGq2GqP1V1rqjXbOst6lXWLdY+Nhc04m1k2tTa3bCm2bFux7RrbU7Zv7Oztku2+t2uwe2pvZM+zL7Kvtb/jQHcIdJjqUO1wxZHoyHbMddzgeNEJdfJ0EjtVOV1wRp29nCXOG5zbRxBG+IyQjqgecd2F5sJxKXSpdbk/kjkyamTxyIaRz0fZjEobtWLUqVGfXD1d81y3ut4erT967Oji0U2j/3ZzchO4Vbldcae7h7nPdW90f+Hh7CHy2Ohxw5PhOc7ze88Wz49e3l5yrzqvLm8b7wzv9d7X2QbsWPYS9mkfgk+wz1yfQz7vfL18C3z3+f7l5+KX67fT7+kY+zGiMVvHPPS38uf7b/bvCGAFZAT8GNARaBnID6wOfBBkHSQM2hb0hOPIyeHs4jwPdg2WBx8IfsP15c7mNodgIeEhpSFtofqhiaHrQu+FWYVlh9WG9YR7hs8Mb44gRERGrIi4zjPjCXg1vJ6x3mNnj22NpEXGR66LfBDlFCWPahqHjhs7buW4O9G20dLohhgQw4tZGXM31j52auyv44njY8dXjX8cNzpuVtypeEb85Pid8a8TghOWJdxOdEhUJrYk6SSlJ9UkvUkOSS5P7kgZlTI75VyqSaoktTGNlJaUti2td0LohNUTOtM900vSr020nzh94plJJpPyJh2erDOZP3l/BiEjOWNnxgd+DL+a35vJy1yf2SPgCtYIngmDhKuEXSJ/UbnoSZZ/VnnW02z/7JXZXeJAcYW4W8KVrJO8yInI2ZTzJjcmd3tuf15y3p58cn5G/kGpvjRX2jrFfMr0Ke0yZ1mJrGOq79TVU3vkkfJtCkQxUdFYYAAP7+eVDsrvlPcLAwqrCt9OS5q2f7redOn08zOcZiye8aQorOinmfhMwcyWWZaz5s+6P5sze/McZE7mnJa51nMXzu2cFz5vx3zq/Nz5vxW7FpcXv1qQvKBpodnCeQsffhf+XW2Jdom85Pr3ft9vWoQvkixqW+y+eO3iT6XC0rNlrmUVZR+WCJac/WH0D5U/9C/NWtq2zGvZxuXE5dLl11YErthRrldeVP5w5biV9atYq0pXvVo9efWZCo+KTWuoa5RrOiqjKhvX2qxdvvbDOvG6q1XBVXvWm65fvP7NBuGGSxuDNtZtMttUtun9j5Ifb2wO31xfbVddsYW4pXDL461JW0/9xP6pZpvJtrJtH7dLt3fsiNvRWuNdU7PTdOeyWrRWWdu1K33Xxd0huxvrXOo272HuKdsL9ir3/vFzxs/X9kXua9nP3l/3i+0v6w8wDpTWI/Uz6nsaxA0djamN7QfHHmxp8ms68OvIX7cfsjxUddjw8LIj1CMLj/QfLTra2yxr7j6Wfexhy+SW28dTjl9pHd/adiLyxOmTYSePn+KcOnra//ShM75nDp5ln20453Wu/rzn+QO/ef52oM2rrf6C94XGiz4Xm9rHtB+5FHjp2OWQyyev8K6cuxp9tf1a4rUb19Ovd9wQ3nh6M+/mi1uFt/puz7tDuFN6V/duxT3Te9W/O/6+p8Or4/D9kPvnH8Q/uP1Q8PDZI8WjD50LH9MfVzyxeFLz1O3poa6wrot/TPij85nsWV93yZ96f65/7vD8l7+C/jrfk9LT+UL+ov
                                width: 200px;
                                height: 60px;
                                background-size: cover;
                                position: absolute;
                                right: 0px;
                            }

                            .warning {
                                padding: 10px;
                                text-align: center;
                                border-radius: 5px;
                                color: #FFF;
                                margin-top: 40px;
                                margin-bottom: 40px;
                                font-weight:900;
                            }

                            .high {
                                background-color: #F44336;
                            }

                            .moderate {
                                background-color: #ff7e33eb;
                            }

                            .low {
                                background-color: #4fce0eb8;
                            }

                            ul {
                                list-style: none;
                                margin: 0;
                                padding: 0;
                            }

                            .alert {
                                margin-top: 15px;
                            }

                            .alert-body {
                                background-color: #FFF;
                                list-style: none;
                                padding: 10px;
                                border-radius: 5px;
                                border: 1px solid #EEE;
                                margin-top: 3px;
                            }

                            .alert-body>.title {
                                display: block;
                                padding: 5px 5px 5px 10px;
                                font-size: 13px;
                            }

                            .high-label {
                                background-color: #F44336;
                                padding: 5px;
                                text-transform: uppercase;
                                font-size: 10px;
                                font-weight: bold;
                                border-radius: 3px 0px 0px 0px;
                                margin: 0px;
                                color: #FFF;
                                margin-left: 10px;
                            }

                            .moderate-label {
                                background-color: #ff7e33eb;
                                padding: 5px;
                                text-transform: uppercase;
                                font-size: 10px;
                                font-weight: bold;
                                border-radius: 3px 0px 0px 0px;
                                margin: 0px;
                                color: #FFF;
                                margin-left: 10px;
                            }

                            .low-label {
                                background-color: #4fce0eb8;
                                padding: 5px;
                                text-transform: uppercase;
                                font-size: 10px;
                                font-weight: bold;
                                border-radius: 3px 0px 0px 0px;
                                margin: 0px;
                                color: #FFF;
                                margin-left: 10px;
                            }

                            .description {
                                margin: 0;
                                padding: 10px;
                                color:#333;
                                font-size:12px;
                            }

                            ul {
                                list-style: none;
                                margin: 0;
                                padding: 0;
                            }

                            .alert-id {
                                background-color: #636363;
                                padding: 5px;
                                text-transform: uppercase;
                                font-size: 10px;
                                font-weight: bold;
                                border-radius: 0px 3px 0px 0px;
                                margin: 0px;
                                color: #FFF;
                                margin-right: 10px;
                            }
                 
                            .header>p {
                                font-size:12px;
                            }
                            @page {
                                @top-center { 
                                    content: "REPORT_HEADER - Page " counter(page) " of " counter(pages) ".";
                                    font-size:12px;
                                    color:#CCC;
                                }
                                @bottom-center { 
                                    content: "This report has been autogenerated by a Tinycheck device. For any question, bug report or feedback, please contact tinycheck@kaspersky.com.";
                                    font-size:12px;  
                                    color:#CCC;
                                }
                            }   
                        </style>
                    </head>
                    <body>""".replace("REPORT_HEADER", "Report for the capture {}".format(self.capture_sha1))