mirror of
https://github.com/Grasscutters/mitmproxy.git
synced 2024-11-27 10:26:23 +00:00
166 lines
5.6 KiB
Python
166 lines
5.6 KiB
Python
|
from typing import TextIO, Callable
|
||
|
|
from unittest import mock
|
||
|
|
from unittest.mock import MagicMock
|
||
|
|
|
||
|
|
from mitmproxy.test import tflow
|
||
|
|
from mitmproxy.test import tutils
|
||
|
|
|
||
|
|
from examples.complex.webscanner_helper.mapping import MappingAddon, MappingAddonConfig
|
||
|
|
|
||
|
|
|
||
|
|
class TestConfig:
|
||
|
|
|
||
|
|
def test_config(self):
|
||
|
|
assert MappingAddonConfig.HTML_PARSER == "html.parser"
|
||
|
|
|
||
|
|
|
||
|
|
url = "http://10.10.10.10"
|
||
|
|
new_content = "My Text"
|
||
|
|
mapping_content = f'{{"{url}": {{"body": "{new_content}"}}}}'
|
||
|
|
|
||
|
|
|
||
|
|
class TestMappingAddon:
|
||
|
|
|
||
|
|
def test_init(self, tmpdir):
|
||
|
|
tmpfile = tmpdir.join("tmpfile")
|
||
|
|
with open(tmpfile, "w") as tfile:
|
||
|
|
tfile.write(mapping_content)
|
||
|
|
mapping = MappingAddon(tmpfile)
|
||
|
|
assert "My Text" in str(mapping.mapping_templates._dump())
|
||
|
|
|
||
|
|
def test_load(self, tmpdir):
|
||
|
|
tmpfile = tmpdir.join("tmpfile")
|
||
|
|
with open(tmpfile, "w") as tfile:
|
||
|
|
tfile.write(mapping_content)
|
||
|
|
mapping = MappingAddon(tmpfile)
|
||
|
|
loader = MagicMock()
|
||
|
|
|
||
|
|
mapping.load(loader)
|
||
|
|
assert 'mapping_file' in str(loader.add_option.call_args_list)
|
||
|
|
assert 'map_persistent' in str(loader.add_option.call_args_list)
|
||
|
|
|
||
|
|
def test_configure(self, tmpdir):
|
||
|
|
tmpfile = tmpdir.join("tmpfile")
|
||
|
|
with open(tmpfile, "w") as tfile:
|
||
|
|
tfile.write(mapping_content)
|
||
|
|
mapping = MappingAddon(tmpfile)
|
||
|
|
new_filename = "My new filename"
|
||
|
|
updated = {str(mapping.OPT_MAPPING_FILE): new_filename, str(mapping.OPT_MAP_PERSISTENT): True}
|
||
|
|
|
||
|
|
open_mock = mock.mock_open(read_data="{}")
|
||
|
|
with mock.patch("builtins.open", open_mock):
|
||
|
|
mapping.configure(updated)
|
||
|
|
assert new_filename in str(open_mock.mock_calls)
|
||
|
|
assert mapping.filename == new_filename
|
||
|
|
assert mapping.persistent
|
||
|
|
|
||
|
|
def test_response_filtered(self, tmpdir):
|
||
|
|
tmpfile = tmpdir.join("tmpfile")
|
||
|
|
with open(tmpfile, "w") as tfile:
|
||
|
|
tfile.write(mapping_content)
|
||
|
|
mapping = MappingAddon(tmpfile)
|
||
|
|
f = tflow.tflow(resp=tutils.tresp())
|
||
|
|
test_content = b"Test"
|
||
|
|
f.response.content = test_content
|
||
|
|
|
||
|
|
mapping.response(f)
|
||
|
|
assert f.response.content == test_content
|
||
|
|
|
||
|
|
def test_response(self, tmpdir):
|
||
|
|
tmpfile = tmpdir.join("tmpfile")
|
||
|
|
with open(tmpfile, "w") as tfile:
|
||
|
|
tfile.write(mapping_content)
|
||
|
|
mapping = MappingAddon(tmpfile)
|
||
|
|
f = tflow.tflow(resp=tutils.tresp())
|
||
|
|
test_content = b"<body> Test </body>"
|
||
|
|
f.response.content = test_content
|
||
|
|
f.request.url = url
|
||
|
|
|
||
|
|
mapping.response(f)
|
||
|
|
assert f.response.content.decode("utf-8") == new_content
|
||
|
|
|
||
|
|
def test_response_content_type(self, tmpdir):
|
||
|
|
tmpfile = tmpdir.join("tmpfile")
|
||
|
|
with open(tmpfile, "w") as tfile:
|
||
|
|
tfile.write(mapping_content)
|
||
|
|
mapping = MappingAddon(tmpfile)
|
||
|
|
f = tflow.tflow(resp=tutils.tresp())
|
||
|
|
test_content = b"<body> Test </body>"
|
||
|
|
f.response.content = test_content
|
||
|
|
f.request.url = url
|
||
|
|
f.response.headers.add("content-type", "content-type")
|
||
|
|
|
||
|
|
mapping.response(f)
|
||
|
|
assert f.response.content == test_content
|
||
|
|
|
||
|
|
def test_response_not_existing(self, tmpdir):
|
||
|
|
tmpfile = tmpdir.join("tmpfile")
|
||
|
|
with open(tmpfile, "w") as tfile:
|
||
|
|
tfile.write(mapping_content)
|
||
|
|
mapping = MappingAddon(tmpfile)
|
||
|
|
f = tflow.tflow(resp=tutils.tresp())
|
||
|
|
test_content = b"<title> Test </title>"
|
||
|
|
f.response.content = test_content
|
||
|
|
f.request.url = url
|
||
|
|
mapping.response(f)
|
||
|
|
assert f.response.content == test_content
|
||
|
|
|
||
|
|
def test_persistance_false(self, tmpdir):
|
||
|
|
tmpfile = tmpdir.join("tmpfile")
|
||
|
|
with open(tmpfile, "w") as tfile:
|
||
|
|
tfile.write(mapping_content)
|
||
|
|
mapping = MappingAddon(tmpfile)
|
||
|
|
|
||
|
|
open_mock = mock.mock_open(read_data="{}")
|
||
|
|
with mock.patch("builtins.open", open_mock):
|
||
|
|
mapping.done()
|
||
|
|
assert len(open_mock.mock_calls) == 0
|
||
|
|
|
||
|
|
def test_persistance_true(self, tmpdir):
|
||
|
|
tmpfile = tmpdir.join("tmpfile")
|
||
|
|
with open(tmpfile, "w") as tfile:
|
||
|
|
tfile.write(mapping_content)
|
||
|
|
mapping = MappingAddon(tmpfile, persistent=True)
|
||
|
|
|
||
|
|
open_mock = mock.mock_open(read_data="{}")
|
||
|
|
with mock.patch("builtins.open", open_mock):
|
||
|
|
mapping.done()
|
||
|
|
with open(tmpfile, "r") as tfile:
|
||
|
|
results = tfile.read()
|
||
|
|
assert len(open_mock.mock_calls) != 0
|
||
|
|
assert results == mapping_content
|
||
|
|
|
||
|
|
def test_persistance_true_add_content(self, tmpdir):
|
||
|
|
tmpfile = tmpdir.join("tmpfile")
|
||
|
|
with open(tmpfile, "w") as tfile:
|
||
|
|
tfile.write(mapping_content)
|
||
|
|
mapping = MappingAddon(tmpfile, persistent=True)
|
||
|
|
|
||
|
|
f = tflow.tflow(resp=tutils.tresp())
|
||
|
|
test_content = b"<title> Test </title>"
|
||
|
|
f.response.content = test_content
|
||
|
|
f.request.url = url
|
||
|
|
|
||
|
|
mapping.response(f)
|
||
|
|
mapping.done()
|
||
|
|
with open(tmpfile, "r") as tfile:
|
||
|
|
results = tfile.read()
|
||
|
|
assert mapping_content in results
|
||
|
|
|
||
|
|
def mock_dump(self, f: TextIO, value_dumper: Callable):
|
||
|
|
assert value_dumper(None) == "None"
|
||
|
|
try:
|
||
|
|
value_dumper("Test")
|
||
|
|
except RuntimeError:
|
||
|
|
assert True
|
||
|
|
else:
|
||
|
|
assert False
|
||
|
|
|
||
|
|
def test_dump(selfself, tmpdir):
|
||
|
|
tmpfile = tmpdir.join("tmpfile")
|
||
|
|
with open(tmpfile, "w") as tfile:
|
||
|
|
tfile.write("{}")
|
||
|
|
mapping = MappingAddon(tmpfile, persistent=True)
|
||
|
|
with mock.patch('examples.complex.webscanner_helper.urldict.URLDict.dump', selfself.mock_dump):
|
||
|
|
mapping.done()
|