#!/usr/bin/env python
# -*- coding=utf-8 -*-
import time
from enum import unique, Enum


@unique
class RISK(Enum):
    INFO = 0
    LOW = 1
    MEDIUM = 2
    HIGH = 3


class Reference(object):
    def __init__(self, location="", detail=""):
        self.location = location
        self.detail = detail

    def to_dict(self):
        return vars(self)


class Vulnerability(object):
    def __init__(self, name, i18n_name, category, description, solution, risk_level):
        self.name = name  # 漏洞唯一标识
        self.i18n_name = i18n_name  # 漏洞类型(展示用, 可被翻译)
        self.description = description  # 漏洞介绍(可被翻译)
        self.solution = solution  # 解决方案(可被翻译)
        if not isinstance(risk_level, type(RISK.INFO)):
            raise Exception("Invalid Vulnerability: {}".format(name))
        self.risk_level = risk_level  # 危险等级 0-4
        self.category = category  # 类型

    def to_dict(self):
        dic = vars(self)
        dic['risk_level'] = self.risk_level.value
        return dic


class VulnReport(object):
    def __init__(self, vuln=None, name=None, i18n_name=None, category=None, description=None, solution=None,
                 risk_level=None,
                 plugin_name=None):
        if not vuln:
            vuln = Vulnerability(name=name, i18n_name=i18n_name, category=None, description=description,
                                 solution=solution,
                                 risk_level=risk_level)
        self.vulnerability = vuln
        self.update_time = int(time.time())  # 扫描时间
        self.plugin_name = plugin_name
        self.reference = []  # 关于此Vul的所有发生位置

    def to_dict(self):
        return dict(
            vulnerability=self.vulnerability.to_dict(),
            update_time=self.update_time,
            plugin_name=self.plugin_name,
            reference=map(lambda e: e.to_dict(), self.reference)
        )


if __name__ == '__main__':
    pass
