addon options: readfile, save, script

This commit is contained in:
Aldo Cortesi 2018-02-24 15:28:43 +13:00
parent 144b559b46
commit 12633adeb9
5 changed files with 28 additions and 27 deletions

View file

@ -11,6 +11,11 @@ class ReadFile:
"""
An addon that handles reading from file on startup.
"""
def load(self, loader):
loader.add_option(
"rfile", typing.Optional[str], None,
"Read flows from file."
)
def load_flows(self, fo: typing.IO[bytes]) -> int:
cnt = 0

View file

@ -16,6 +16,16 @@ class Save:
self.filt = None
self.active_flows = set() # type: Set[flow.Flow]
def load(self, loader):
loader.add_option(
"save_stream_file", typing.Optional[str], None,
"Stream flows to file as they arrive. Prefix path with + to append."
)
loader.add_option(
"save_stream_filter", typing.Optional[str], None,
"Filter which flows are written to file."
)
def open_file(self, path):
if path.startswith("+"):
path = path[1:]

View file

@ -98,6 +98,14 @@ class ScriptLoader:
self.is_running = False
self.addons = []
def load(self, loader):
loader.add_option(
"scripts", typing.Sequence[str], [],
"""
Execute a script.
"""
)
def running(self):
self.is_running = True

View file

@ -67,10 +67,6 @@ class Options(optmanager.OptManager):
view_filter = None # type: Optional[str]
# FIXME: Options that should be uncomplicated to migrate to addons
rfile = None # type: Optional[str]
save_stream_file = None # type: Optional[str]
save_stream_filter = None # type: Optional[str]
scripts = None # type: Sequence[str]
setheaders = None # type: Sequence[str]
stickyauth = None # type: Optional[str]
stickycookie = None # type: Optional[str]
@ -90,16 +86,6 @@ class Options(optmanager.OptManager):
"server", bool, True,
"Start a proxy server. Enabled by default."
)
self.add_option(
"rfile", Optional[str], None,
"Read flows from file."
)
self.add_option(
"scripts", Sequence[str], [],
"""
Execute a script.
"""
)
self.add_option(
"showhost", bool, False,
"Use the Host header to construct URLs for display."
@ -144,14 +130,6 @@ class Options(optmanager.OptManager):
"The default content view mode.",
choices = [i.name.lower() for i in contentviews.views]
)
self.add_option(
"save_stream_file", Optional[str], None,
"Stream flows to file as they arrive. Prefix path with + to append."
)
self.add_option(
"save_stream_filter", Optional[str], None,
"Filter which flows are written to file."
)
# Proxy options
self.add_option(

View file

@ -41,7 +41,7 @@ class TestReadFile:
@mock.patch('mitmproxy.master.Master.load_flow')
def test_configure(self, mck, tmpdir, data, corrupt_data):
rf = readfile.ReadFile()
with taddons.context() as tctx:
with taddons.context(rf) as tctx:
tf = tmpdir.join("tfile")
tf.write(data.getvalue())
@ -58,7 +58,7 @@ class TestReadFile:
@mock.patch('mitmproxy.master.Master.load_flow')
def test_corrupt(self, mck, corrupt_data):
rf = readfile.ReadFile()
with taddons.context() as tctx:
with taddons.context(rf) as tctx:
with pytest.raises(exceptions.FlowReadException):
rf.load_flows(io.BytesIO(b"qibble"))
assert not mck.called
@ -71,7 +71,7 @@ class TestReadFile:
def test_nonexisting_file(self):
rf = readfile.ReadFile()
with taddons.context() as tctx:
with taddons.context(rf) as tctx:
with pytest.raises(exceptions.FlowReadException):
rf.load_flows_from_path("nonexistent")
assert len(tctx.master.logs) == 1
@ -82,7 +82,7 @@ class TestReadFileStdin:
@mock.patch('sys.stdin')
def test_stdin(self, stdin, load_flow, data, corrupt_data):
rf = readfile.ReadFileStdin()
with taddons.context() as tctx:
with taddons.context(rf) as tctx:
stdin.buffer = data
tctx.configure(rf, rfile="-")
assert not load_flow.called
@ -97,7 +97,7 @@ class TestReadFileStdin:
@mock.patch('mitmproxy.master.Master.load_flow')
def test_normal(self, load_flow, tmpdir, data):
rf = readfile.ReadFileStdin()
with taddons.context():
with taddons.context(rf):
tfile = tmpdir.join("tfile")
tfile.write(data.getvalue())
rf.load_flows_from_path(str(tfile))