1import pytest
2
3from mitmproxy.addons import blocklist
4from mitmproxy.exceptions import OptionsError
5from mitmproxy.test import taddons
6from mitmproxy.test import tflow
7
8
9@pytest.mark.parametrize("filter,err", [
10    ("/~u index.html/TOOMANY/300", "Invalid number of parameters"),
11    (":~d ~d ~d:200", "Invalid filter"),
12    ("/~u index.html/999", "Invalid HTTP status code"),
13    ("/~u index.html/abc", "Invalid HTTP status code"),
14])
15def test_parse_spec_err(filter, err):
16    with pytest.raises(ValueError, match=err):
17        blocklist.parse_spec(filter)
18
19
20class TestBlockList:
21    @pytest.mark.parametrize("filter,status_code", [
22        (":~u example.org:404", 404),
23        (":~u example.com:404", None),
24        ("/!jpg/418", None),
25        ("/!png/418", 418),
26
27    ])
28    def test_block(self, filter, status_code):
29        bl = blocklist.BlockList()
30        with taddons.context(bl) as tctx:
31            tctx.configure(bl, block_list=[filter])
32            f = tflow.tflow()
33            f.request.url = b"https://example.org/images/test.jpg"
34            bl.request(f)
35            if status_code is not None:
36                assert f.response.status_code == status_code
37                assert f.metadata['blocklisted']
38            else:
39                assert not f.response
40
41    def test_special_kill_status_closes_connection(self):
42        bl = blocklist.BlockList()
43        with taddons.context(bl) as tctx:
44            tctx.configure(bl, block_list=[':.*:444'])
45            f = tflow.tflow()
46            bl.request(f)
47            assert f.error.msg == f.error.KILLED_MESSAGE
48            assert f.response is None
49            assert f.metadata['blocklisted'] is True
50
51    def test_already_handled(self):
52        """Test that we don't interfere if another addon already killed this request."""
53        bl = blocklist.BlockList()
54        with taddons.context(bl) as tctx:
55            tctx.configure(bl, block_list=["/.*/404"])
56            f = tflow.tflow()
57            f.kill()  # done by another addon.
58            bl.request(f)
59            assert not f.response
60
61    def test_configure_err(self):
62        bl = blocklist.BlockList()
63        with taddons.context(bl) as tctx:
64            with pytest.raises(OptionsError):
65                tctx.configure(bl, block_list=["lalelu"])
66