mirror of
https://github.com/Grasscutters/mitmproxy.git
synced 2024-11-23 00:01:36 +00:00
Merge pull request #2258 from mhils/readfile
Integrate readstdin into readfile
This commit is contained in:
commit
d5ea08db62
@ -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(),
|
||||
]
|
||||
|
@ -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)
|
||||
|
@ -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")
|
@ -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()
|
||||
|
@ -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
|
||||
)
|
||||
|
@ -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,
|
||||
)
|
||||
|
@ -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
|
||||
|
@ -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
|
||||
|
@ -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
|
@ -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)
|
||||
|
Loading…
Reference in New Issue
Block a user