aboutsummaryrefslogtreecommitdiffstats
path: root/test/examples/webscanner_helper/test_urldict.py
blob: 7bd4fb017f97c33a9d826edcd7c06b3b4acdb1fc (plain)
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
from mitmproxy.test import tflow, tutils
from examples.complex.webscanner_helper.urldict import URLDict

url = "http://10.10.10.10"
new_content_body = "New Body"
new_content_title = "New Title"
content = f'{{"body": "{new_content_body}", "title": "{new_content_title}"}}'
url_error = "i~nvalid"
input_file_content = f'{{"{url}": {content}}}'
input_file_content_error = f'{{"{url_error}": {content}}}'


class TestUrlDict:

    def test_urldict_empty(self):
        urldict = URLDict()
        dump = urldict.dumps()
        assert dump == '{}'

    def test_urldict_loads(self):
        urldict = URLDict.loads(input_file_content)
        dump = urldict.dumps()
        assert dump == input_file_content

    def test_urldict_set_error(self, tmpdir):
        tmpfile = tmpdir.join("tmpfile")
        with open(tmpfile, "w") as tfile:
            tfile.write(input_file_content_error)
        with open(tmpfile, "r") as tfile:
            try:
                URLDict.load(tfile)
            except ValueError:
                assert True
            else:
                assert False

    def test_urldict_get(self, tmpdir):
        tmpfile = tmpdir.join("tmpfile")
        with open(tmpfile, "w") as tfile:
            tfile.write(input_file_content)
        with open(tmpfile, "r") as tfile:
            urldict = URLDict.load(tfile)

        f = tflow.tflow(resp=tutils.tresp())
        f.request.url = url
        selection = urldict[f]
        assert "body" in selection[0]
        assert new_content_body in selection[0]["body"]
        assert "title" in selection[0]
        assert new_content_title in selection[0]["title"]

        selection_get = urldict.get(f)
        assert "body" in selection_get[0]
        assert new_content_body in selection_get[0]["body"]
        assert "title" in selection_get[0]
        assert new_content_title in selection_get[0]["title"]

        try:
            urldict["body"]
        except KeyError:
            assert True
        else:
            assert False

        assert urldict.get("body", default="default") == "default"

    def test_urldict_dumps(self, tmpdir):
        tmpfile = tmpdir.join("tmpfile")
        with open(tmpfile, "w") as tfile:
            tfile.write(input_file_content)
        with open(tmpfile, "r") as tfile:
            urldict = URLDict.load(tfile)

        dump = urldict.dumps()
        assert dump == input_file_content

    def test_urldict_dump(self, tmpdir):
        tmpfile = tmpdir.join("tmpfile")
        outfile = tmpdir.join("outfile")
        with open(tmpfile, "w") as tfile:
            tfile.write(input_file_content)
        with open(tmpfile, "r") as tfile:
            urldict = URLDict.load(tfile)
        with open(outfile, "w") as ofile:
            urldict.dump(ofile)

        with open(outfile, "r") as ofile:
            output = ofile.read()
        assert output == input_file_content