tlsassistant/tests/jinja_test.py at master · stfbk/tlsassistant · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
# imports
from jinja2 import FileSystemLoader, Environment
import datetime
import random
import argparse

# parse inputs
parser = argparse.ArgumentParser()
parser.add_argument(
    "-f",
    "--file",
    help="file to render",
    choices=["hosts_report.html", "modules_report.html"],
    required=True,
)
# list of modules
parser.add_argument(
    "-m",
    "--modules",
    help="list of modules to import",
    nargs="+",
    type=str,
    required=True,
)
# list of hosts
parser.add_argument(
    "-s", "--hosts", help="list of hosts to import", nargs="+", type=str, required=True
)

args = parser.parse_args()

env = Environment(loader=FileSystemLoader("../configs/out_template/"))

template = env.get_template(args.file)

from json import load

m_dict = {}


class CaseInsensitiveDict(dict):
    @classmethod
    def _k(cls, key):
        return key.lower()

    def __init__(self, *args, **kwargs):
        super(CaseInsensitiveDict, self).__init__(*args, **kwargs)
        self._convert_keys()

    def __getitem__(self, key):
        return super(CaseInsensitiveDict, self).__getitem__(self.__class__._k(key))

    def __setitem__(self, key, value):
        super(CaseInsensitiveDict, self).__setitem__(self.__class__._k(key), value)

    def __delitem__(self, key):
        return super(CaseInsensitiveDict, self).__delitem__(self.__class__._k(key))

    def __contains__(self, key):
        return super(CaseInsensitiveDict, self).__contains__(self.__class__._k(key))

    def has_key(self, key):
        return super(CaseInsensitiveDict, self).has_key(self.__class__._k(key))

    def pop(self, key, *args, **kwargs):
        return super(CaseInsensitiveDict, self).pop(
            self.__class__._k(key), *args, **kwargs
        )

    def get(self, key, *args, **kwargs):
        return super(CaseInsensitiveDict, self).get(
            self.__class__._k(key), *args, **kwargs
        )

    def setdefault(self, key, *args, **kwargs):
        return super(CaseInsensitiveDict, self).setdefault(
            self.__class__._k(key), *args, **kwargs
        )

    def update(self, E={}, **F):
        super(CaseInsensitiveDict, self).update(self.__class__(E))
        super(CaseInsensitiveDict, self).update(self.__class__(**F))

    def _convert_keys(self):
        for k in list(self.keys()):
            v = super(CaseInsensitiveDict, self).pop(k)
            self.__setitem__(k, v)


for module in args.modules:
    with open("../configs/mitigations/" + module.upper() + ".json") as f:
        data = CaseInsensitiveDict(load(f))
        m_dict[module] = CaseInsensitiveDict(data["Entry"])
print(m_dict)

results = {}
if args.file == "hosts_report.html":
    for host in args.hosts:
        results[host] = {}
        for module in args.modules:
            # choose randomly if the module is in the list
            if random.randint(0, 1) == 1:
                results[host][module] = m_dict[module]
        # if no modules, delete the host
        if not results[host]:
            del results[host]
else:
    for module in args.modules:
        results[module] = m_dict[module]
        results[module]["hosts"] = []
        for host in args.hosts:
            # choose randomly if the host is in the list
            if random.randint(0, 1) == 1:
                results[module]["hosts"].append(host)
        # if no host, delete the module from the results
        if not results[module]["hosts"]:
            del results[module]

print(results)

with open("../outtest.html", "w") as f:
    f.write(
        template.render(
            date=datetime.datetime.now(),
            version="1.2.2 closed beta",
            modules=args.modules,
            results=results,
        )
    )