Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2024-11-23 00:01:36 +00:00
Code Issues Projects Releases Wiki Activity

separate reading from stdin into its own addon

Browse Source
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mitmproxy/addons/__init__.py
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(),
]

32
mitmproxy/addons/readfile.py
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,16 +29,13 @@ 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:
return self.load_flows(f)
except IOError as e:
ctx.log.error("Cannot load flows: {}".format(e))
raise exceptions.FlowReadException(str(e)) from e
path = os.path.expanduser(path)
try:
with open(path, "rb") as f:
return self.load_flows(f)
except IOError as e:
ctx.log.error("Cannot load flows: {}".format(e))
raise exceptions.FlowReadException(str(e)) from e
def running(self):
if ctx.options.rfile:
@ -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)

10
mitmproxy/tools/console/master.py
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()

3
mitmproxy/tools/dump.py
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
)

2
mitmproxy/tools/web/master.py
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,
)

107
test/mitmproxy/addons/test_readfile.py
View File

@ -37,59 +37,68 @@ def corrupt_data():
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")
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")
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):
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')
def test_corrupt(self, 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(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
@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):
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_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
@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