aboutsummaryrefslogtreecommitdiffstats
path: root/test/mitmproxy/addons/test_readfile.py
blob: a0826a2659ecf148a26caf0d1d3dcc81b46b42f3 (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
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
import io
from unittest import mock

import pytest

import mitmproxy.io
from mitmproxy import exceptions
from mitmproxy.addons import readfile
from mitmproxy.test import taddons
from mitmproxy.test import tflow


@pytest.fixture
def data():
    f = io.BytesIO()

    w = mitmproxy.io.FlowWriter(f)
    flows = [
        tflow.tflow(resp=True),
        tflow.tflow(err=True),
        tflow.ttcpflow(),
        tflow.ttcpflow(err=True)
    ]
    for flow in flows:
        w.add(flow)

    f.seek(0)
    return f


@pytest.fixture
def corrupt_data():
    f = data()
    f.seek(0, io.SEEK_END)
    f.write(b"qibble")
    f.seek(0)
    return f


class TestReadFile:
    @mock.patch('mitmproxy.master.Master.load_flow')
    def test_configure(self, mck, tmpdir, data, corrupt_data):
        rf = readfile.ReadFile()
        with taddons.context(rf) as tctx:
            tf = tmpdir.join("tfile")

            tf.write(data.getvalue())
            tctx.configure(rf, rfile=str(tf))
            assert not mck.called
            rf.running()
            assert mck.called

            tf.write(corrupt_data.getvalue())
            tctx.configure(rf, rfile=str(tf))
            with pytest.raises(exceptions.OptionsError):
                rf.running()

    @pytest.mark.asyncio
    async def test_corrupt(self, corrupt_data):
        rf = readfile.ReadFile()
        with taddons.context(rf) as tctx:
            with mock.patch('mitmproxy.master.Master.load_flow') as mck:
                with pytest.raises(exceptions.FlowReadException):
                    rf.load_flows(io.BytesIO(b"qibble"))
                assert not mck.called

                tctx.master.clear()
                with pytest.raises(exceptions.FlowReadException):
                    rf.load_flows(corrupt_data)
                assert await tctx.master.await_log("file corrupted")
                assert mck.called

    @pytest.mark.asyncio
    async def test_nonexisting_file(self):
        rf = readfile.ReadFile()
        with taddons.context(rf) as tctx:
            with pytest.raises(exceptions.FlowReadException):
                rf.load_flows_from_path("nonexistent")
            assert await tctx.master.await_log("nonexistent")


class TestReadFileStdin:
    @mock.patch('mitmproxy.master.Master.load_flow')
    @mock.patch('sys.stdin')
    def test_stdin(self, stdin, load_flow, data, corrupt_data):
        rf = readfile.ReadFileStdin()
        with taddons.context(rf) as tctx:
            stdin.buffer = data
            tctx.configure(rf, rfile="-")
            assert not load_flow.called
            rf.running()
            assert load_flow.called

            stdin.buffer = corrupt_data
            tctx.configure(rf, rfile="-")
            with pytest.raises(exceptions.OptionsError):
                rf.running()

    @mock.patch('mitmproxy.master.Master.load_flow')
    def test_normal(self, load_flow, tmpdir, data):
        rf = readfile.ReadFileStdin()
        with taddons.context(rf):
            tfile = tmpdir.join("tfile")
            tfile.write(data.getvalue())
            rf.load_flows_from_path(str(tfile))
            assert load_flow.called