summaryrefslogtreecommitdiffstats
path: root/testing/mozbase/mozproxy/tests/test_mitm_addons.py
blob: ed3805ef9d2e0c5bad6d09beb867d050e33f9918 (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
#!/usr/bin/env python
import json
import os
from unittest import mock

import mozunit

here = os.path.dirname(__file__)
os.environ["MOZPROXY_DIR"] = os.path.join(here, "files")

protocol = {
    "http_protocol": {"aax-us-iad.amazon.com": "HTTP/1.1"},
    "recorded_requests": 4,
    "recorded_requests_unique": 1,
}


@mock.patch(
    "mozproxy.backends.mitm.scripts.http_protocol_extractor.HttpProtocolExtractor.get_ctx"
)
def test_http_protocol_generate_json_file(ctx_mock):
    ctx_mock.return_value.options.save_stream_file = os.path.join(
        os.environ["MOZPROXY_DIR"], "http_protocol_recording_done.mp"
    )

    from mozproxy.backends.mitm.scripts.http_protocol_extractor import (
        HttpProtocolExtractor,
    )

    test_http_protocol = HttpProtocolExtractor()
    test_http_protocol.ctx = test_http_protocol.get_ctx()

    # test data
    test_http_protocol.request_protocol = protocol["http_protocol"]
    test_http_protocol.hashes = ["Hash string"]
    test_http_protocol.request_count = protocol["recorded_requests"]

    test_http_protocol.done()

    json_path = os.path.join(
        os.environ["MOZPROXY_DIR"], "http_protocol_recording_done.json"
    )
    assert os.path.exists(json_path)
    with open(json_path) as json_file:
        output_data = json.load(json_file)

        assert output_data["recorded_requests"] == protocol["recorded_requests"]
        assert (
            output_data["recorded_requests_unique"]
            == protocol["recorded_requests_unique"]
        )
        assert output_data["http_protocol"] == protocol["http_protocol"]


@mock.patch(
    "mozproxy.backends.mitm.scripts.http_protocol_extractor.HttpProtocolExtractor.get_ctx"
)
def test_http_protocol_response(ctx_mock):
    ctx_mock.return_value.options.save_stream_file = os.path.join(
        os.environ["MOZPROXY_DIR"], "http_protocol_recording_done.mp"
    )

    from mozproxy.backends.mitm.scripts.http_protocol_extractor import (
        HttpProtocolExtractor,
    )

    test_http_protocol = HttpProtocolExtractor()
    test_http_protocol.ctx = test_http_protocol.get_ctx()

    # test data
    flow = mock.MagicMock()
    flow.type = "http"
    flow.request.url = "https://www.google.com/complete/search"
    flow.request.port = 33
    flow.response.data.http_version = b"HTTP/1.1"

    test_http_protocol.request_protocol = {}
    test_http_protocol.hashes = []
    test_http_protocol.request_count = 0

    test_http_protocol.response(flow)

    assert test_http_protocol.request_count == 1
    assert len(test_http_protocol.hashes) == 1
    assert test_http_protocol.request_protocol["www.google.com"] == "HTTP/1.1"


if __name__ == "__main__":
    mozunit.main(runwith="pytest")