Merge pull request #2258 from mhils/readfile

Integrate readstdin into readfile
This commit is contained in:
Maximilian Hils 2017-04-26 13:40:51 +02:00 committed by GitHub
commit d5ea08db62
10 changed files with 146 additions and 151 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,38 +1,56 @@
import os.path
import sys
import typing
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_file(self, path: str) -> int:
path = os.path.expanduser(path)
def load_flows(self, fo: typing.IO[bytes]) -> int:
cnt = 0
freader = io.FlowReader(fo)
try:
with open(path, "rb") as f:
freader = io.FlowReader(f)
for i in freader.stream():
cnt += 1
ctx.master.load_flow(i)
return cnt
except (IOError, exceptions.FlowReadException) as v:
for flow in freader.stream():
ctx.master.load_flow(flow)
cnt += 1
except (IOError, exceptions.FlowReadException) as e:
if cnt:
ctx.log.warn(
"Flow file corrupted - loaded %i flows." % cnt,
)
ctx.log.warn("Flow file corrupted - loaded %i flows." % cnt)
else:
ctx.log.error("Flow file corrupted.")
raise exceptions.FlowReadException(v)
raise exceptions.FlowReadException(str(e)) from e
else:
return cnt
def load_flows_from_path(self, path: str) -> int:
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:
try:
self.load_flows_file(ctx.options.rfile)
except exceptions.FlowReadException as v:
raise exceptions.OptionsError(v)
self.load_flows_from_path(ctx.options.rfile)
except exceptions.FlowReadException as e:
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

@ -1,26 +0,0 @@
from mitmproxy import ctx
from mitmproxy import io
from mitmproxy import exceptions
import sys
class ReadStdin:
"""
An addon that reads from stdin if we're not attached to (someting like)
a tty.
"""
def running(self, stdin = sys.stdin):
if not stdin.isatty():
ctx.log.info("Reading from stdin")
try:
stdin.buffer.read(0)
except Exception as e:
ctx.log.warn("Cannot read from stdin: {}".format(e))
return
freader = io.FlowReader(stdin.buffer)
try:
for i in freader.stream():
ctx.master.load_flow(i)
except exceptions.FlowReadException as e:
ctx.log.error("Error reading from stdin: %s" % e)
ctx.master.addons.trigger("processing_complete")

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, readstdin, keepserving
from mitmproxy.addons import dumper, termlog, termstatus, keepserving, readfile
class ErrorCheck:
@ -30,7 +30,7 @@ class DumpMaster(master.Master):
if with_dumper:
self.addons.add(dumper.Dumper())
self.addons.add(
readstdin.ReadStdin(),
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

@ -68,5 +68,7 @@ def check_type(name: str, value: typing.Any, typeinfo: typing.Any) -> None:
return
else:
raise e
elif typename.startswith("typing.Any"):
return
elif not isinstance(value, typeinfo):
raise e

View File

@ -1,62 +1,104 @@
from mitmproxy.addons import readfile
from mitmproxy.test import taddons
from mitmproxy.test import tflow
from mitmproxy import io
from mitmproxy import exceptions
import io
from unittest import mock
import pytest
def write_data(path, corrupt=False):
with open(path, "wb") as tf:
w = io.FlowWriter(tf)
for i in range(3):
f = tflow.tflow(resp=True)
w.add(f)
for i in range(3):
f = tflow.tflow(err=True)
w.add(f)
f = tflow.ttcpflow()
w.add(f)
f = tflow.ttcpflow(err=True)
w.add(f)
if corrupt:
tf.write(b"flibble")
import mitmproxy.io
from mitmproxy import exceptions
from mitmproxy.addons import readfile
from mitmproxy.test import taddons
from mitmproxy.test import tflow
@mock.patch('mitmproxy.master.Master.load_flow')
def test_configure(mck, tmpdir):
@pytest.fixture
def data():
f = io.BytesIO()
rf = readfile.ReadFile()
with taddons.context() as tctx:
tf = str(tmpdir.join("tfile"))
write_data(tf)
tctx.configure(rf, rfile=str(tf))
assert not mck.called
rf.running()
assert mck.called
w = mitmproxy.io.FlowWriter(f)
flows = [
tflow.tflow(resp=True),
tflow.tflow(err=True),
tflow.ttcpflow(),
tflow.ttcpflow(err=True)
]
for flow in flows:
w.add(flow)
write_data(tf, corrupt=True)
tctx.configure(rf, rfile=str(tf))
with pytest.raises(exceptions.OptionsError):
f.seek(0)
return f
@pytest.fixture
def corrupt_data():
f = data()
f.seek(0, io.SEEK_END)
f.write(b"qibble")
f.seek(0)
return f
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):
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')
def test_corruption(mck, tmpdir):
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
rf = readfile.ReadFile()
with taddons.context() as tctx:
with pytest.raises(exceptions.FlowReadException):
rf.load_flows_file("nonexistent")
assert not mck.called
assert len(tctx.master.logs) == 1
stdin.buffer = corrupt_data
tctx.configure(rf, rfile="-")
with pytest.raises(exceptions.OptionsError):
rf.running()
tfc = str(tmpdir.join("tfile"))
write_data(tfc, corrupt=True)
with pytest.raises(exceptions.FlowReadException):
rf.load_flows_file(tfc)
assert mck.called
assert len(tctx.master.logs) == 2
@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

View File

@ -1,53 +0,0 @@
import io
from mitmproxy.addons import readstdin
from mitmproxy.test import taddons
from mitmproxy.test import tflow
import mitmproxy.io
from unittest import mock
def gen_data(corrupt=False):
tf = io.BytesIO()
w = mitmproxy.io.FlowWriter(tf)
for i in range(3):
f = tflow.tflow(resp=True)
w.add(f)
for i in range(3):
f = tflow.tflow(err=True)
w.add(f)
f = tflow.ttcpflow()
w.add(f)
f = tflow.ttcpflow(err=True)
w.add(f)
if corrupt:
tf.write(b"flibble")
tf.seek(0)
return tf
class mStdin:
def __init__(self, d):
self.buffer = d
def isatty(self):
return False
@mock.patch('mitmproxy.master.Master.load_flow')
def test_read(m, tmpdir):
rf = readstdin.ReadStdin()
with taddons.context() as tctx:
assert not m.called
rf.running(stdin=mStdin(gen_data()))
assert m.called
rf.running(stdin=mStdin(None))
assert tctx.master.logs
tctx.master.clear()
m.reset_mock()
assert not m.called
rf.running(stdin=mStdin(gen_data(corrupt=True)))
assert m.called
assert tctx.master.logs

View File

@ -79,3 +79,9 @@ def test_check_io():
typecheck.check_type("foo", io.StringIO(), typing.IO[str])
with pytest.raises(TypeError):
typecheck.check_type("foo", "foo", typing.IO[str])
def test_check_any():
typecheck.check_type("foo", 42, typing.Any)
typecheck.check_type("foo", object(), typing.Any)
typecheck.check_type("foo", None, typing.Any)