separate reading from stdin into its own addon

This commit is contained in:
Maximilian Hils 2017-04-26 11:12:14 +02:00
parent b3a1143338
commit ca2827886a
6 changed files with 89 additions and 67 deletions

View File

@ -8,7 +8,6 @@ from mitmproxy.addons import disable_h2c
from mitmproxy.addons import onboarding
from mitmproxy.addons import proxyauth
from mitmproxy.addons import replace
from mitmproxy.addons import readfile
from mitmproxy.addons import script
from mitmproxy.addons import serverplayback
from mitmproxy.addons import setheaders
@ -38,6 +37,5 @@ def default_addons():
stickycookie.StickyCookie(),
streambodies.StreamBodies(),
streamfile.StreamFile(),
readfile.ReadFile(),
upstream_auth.UpstreamAuth(),
]

View File

@ -1,17 +1,17 @@
import os.path
import sys
import typing
import sys
from mitmproxy import ctx
from mitmproxy import io
from mitmproxy import exceptions
from mitmproxy import io
class ReadFile:
"""
An addon that handles reading from file on startup.
"""
def load_flows(self, fo: typing.IO[bytes]) -> int:
cnt = 0
freader = io.FlowReader(fo)
@ -29,9 +29,6 @@ class ReadFile:
return cnt
def load_flows_from_path(self, path: str) -> int:
if path == "-":
return self.load_flows(sys.stdin.buffer)
else:
path = os.path.expanduser(path)
try:
with open(path, "rb") as f:
@ -48,3 +45,12 @@ class ReadFile:
raise exceptions.OptionsError(e) from e
finally:
ctx.master.addons.trigger("processing_complete")
class ReadFileStdin(ReadFile):
"""Support the special case of "-" for reading from stdin"""
def load_flows_from_path(self, path: str) -> int:
if path == "-":
return self.load_flows(sys.stdin.buffer)
else:
return super().load_flows_from_path(path)

View File

@ -17,8 +17,9 @@ from mitmproxy import exceptions
from mitmproxy import master
from mitmproxy import io
from mitmproxy import log
from mitmproxy.addons import view
from mitmproxy.addons import intercept
from mitmproxy.addons import readfile
from mitmproxy.addons import view
from mitmproxy.tools.console import flowlist
from mitmproxy.tools.console import flowview
from mitmproxy.tools.console import grideditor
@ -91,7 +92,12 @@ class ConsoleMaster(master.Master):
signals.sig_add_log.connect(self.sig_add_log)
self.addons.add(Logger())
self.addons.add(*addons.default_addons())
self.addons.add(intercept.Intercept(), self.view, UnsupportedLog())
self.addons.add(
intercept.Intercept(),
self.view,
UnsupportedLog(),
readfile.ReadFile(),
)
def sigint_handler(*args, **kwargs):
self.prompt_for_exit()

View File

@ -1,7 +1,7 @@
from mitmproxy import addons
from mitmproxy import options
from mitmproxy import master
from mitmproxy.addons import dumper, termlog, termstatus, keepserving
from mitmproxy.addons import dumper, termlog, termstatus, keepserving, readfile
class ErrorCheck:
@ -31,5 +31,6 @@ class DumpMaster(master.Master):
self.addons.add(dumper.Dumper())
self.addons.add(
keepserving.KeepServing(),
readfile.ReadFileStdin(),
self.errorcheck
)

View File

@ -7,6 +7,7 @@ from mitmproxy import log
from mitmproxy import master
from mitmproxy.addons import eventstore
from mitmproxy.addons import intercept
from mitmproxy.addons import readfile
from mitmproxy.addons import termlog
from mitmproxy.addons import view
from mitmproxy.addons import termstatus
@ -32,6 +33,7 @@ class WebMaster(master.Master):
self.addons.add(*addons.default_addons())
self.addons.add(
intercept.Intercept(),
readfile.ReadFile(),
self.view,
self.events,
)

View File

@ -37,8 +37,9 @@ def corrupt_data():
return f
@mock.patch('mitmproxy.master.Master.load_flow')
def test_configure(mck, tmpdir, data, corrupt_data):
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:
tf = tmpdir.join("tfile")
@ -54,26 +55,8 @@ def test_configure(mck, tmpdir, data, corrupt_data):
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):
@mock.patch('mitmproxy.master.Master.load_flow')
def test_corrupt(self, mck, corrupt_data):
rf = readfile.ReadFile()
with taddons.context() as tctx:
with pytest.raises(exceptions.FlowReadException):
@ -86,10 +69,36 @@ def test_corrupt(mck, corrupt_data):
assert mck.called
assert len(tctx.master.logs) == 2
def test_nonexisting_file():
def test_nonexisting_file(self):
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
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() 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():
tfile = tmpdir.join("tfile")
tfile.write(data.getvalue())
rf.load_flows_from_path(str(tfile))
assert load_flow.called