aboutsummaryrefslogtreecommitdiffstats
path: root/test/mitmproxy/addons/test_readfile.py
blob: 689d9779e98b5cb994dfcb4ac499f7627af0e346 (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
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


@mock.patch('mitmproxy.master.Master.load_flow')
def test_configure(mck, tmpdir, data, corrupt_data):
    rf = readfile.ReadFile()
    with taddons.context() 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()


@mock.patch('mitmproxy.master.Master.load_flow')
@mock.patch('sys.stdin')
def test_configure_stdin(stdin, load_flow, data, corrupt_data):
    rf = readfile.ReadFile()
    with taddons.context() 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_corrupt(mck, corrupt_data):
    rf = readfile.ReadFile()
    with taddons.context() as tctx:
        with pytest.raises(exceptions.FlowReadException):
            rf.load_flows(io.BytesIO(b"qibble"))
        assert not mck.called
        assert len(tctx.master.logs) == 1

        with pytest.raises(exceptions.FlowReadException):
            rf.load_flows(corrupt_data)
        assert mck.called
        assert len(tctx.master.logs) == 2


def test_nonexisting_file():
    rf = readfile.ReadFile()
    with taddons.context() as tctx:
        with pytest.raises(exceptions.FlowReadException):
            rf.load_flows_from_path("nonexistent")
        assert len(tctx.master.logs) == 1