mirror of
https://github.com/Grasscutters/mitmproxy.git
synced 2024-11-23 08:11:00 +00:00
separate reading from stdin into its own addon
This commit is contained in:
parent
b3a1143338
commit
ca2827886a
@ -8,7 +8,6 @@ from mitmproxy.addons import disable_h2c
|
|||||||
from mitmproxy.addons import onboarding
|
from mitmproxy.addons import onboarding
|
||||||
from mitmproxy.addons import proxyauth
|
from mitmproxy.addons import proxyauth
|
||||||
from mitmproxy.addons import replace
|
from mitmproxy.addons import replace
|
||||||
from mitmproxy.addons import readfile
|
|
||||||
from mitmproxy.addons import script
|
from mitmproxy.addons import script
|
||||||
from mitmproxy.addons import serverplayback
|
from mitmproxy.addons import serverplayback
|
||||||
from mitmproxy.addons import setheaders
|
from mitmproxy.addons import setheaders
|
||||||
@ -38,6 +37,5 @@ def default_addons():
|
|||||||
stickycookie.StickyCookie(),
|
stickycookie.StickyCookie(),
|
||||||
streambodies.StreamBodies(),
|
streambodies.StreamBodies(),
|
||||||
streamfile.StreamFile(),
|
streamfile.StreamFile(),
|
||||||
readfile.ReadFile(),
|
|
||||||
upstream_auth.UpstreamAuth(),
|
upstream_auth.UpstreamAuth(),
|
||||||
]
|
]
|
||||||
|
@ -1,17 +1,17 @@
|
|||||||
import os.path
|
import os.path
|
||||||
|
import sys
|
||||||
import typing
|
import typing
|
||||||
|
|
||||||
import sys
|
|
||||||
|
|
||||||
from mitmproxy import ctx
|
from mitmproxy import ctx
|
||||||
from mitmproxy import io
|
|
||||||
from mitmproxy import exceptions
|
from mitmproxy import exceptions
|
||||||
|
from mitmproxy import io
|
||||||
|
|
||||||
|
|
||||||
class ReadFile:
|
class ReadFile:
|
||||||
"""
|
"""
|
||||||
An addon that handles reading from file on startup.
|
An addon that handles reading from file on startup.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def load_flows(self, fo: typing.IO[bytes]) -> int:
|
def load_flows(self, fo: typing.IO[bytes]) -> int:
|
||||||
cnt = 0
|
cnt = 0
|
||||||
freader = io.FlowReader(fo)
|
freader = io.FlowReader(fo)
|
||||||
@ -29,16 +29,13 @@ class ReadFile:
|
|||||||
return cnt
|
return cnt
|
||||||
|
|
||||||
def load_flows_from_path(self, path: str) -> int:
|
def load_flows_from_path(self, path: str) -> int:
|
||||||
if path == "-":
|
path = os.path.expanduser(path)
|
||||||
return self.load_flows(sys.stdin.buffer)
|
try:
|
||||||
else:
|
with open(path, "rb") as f:
|
||||||
path = os.path.expanduser(path)
|
return self.load_flows(f)
|
||||||
try:
|
except IOError as e:
|
||||||
with open(path, "rb") as f:
|
ctx.log.error("Cannot load flows: {}".format(e))
|
||||||
return self.load_flows(f)
|
raise exceptions.FlowReadException(str(e)) from e
|
||||||
except IOError as e:
|
|
||||||
ctx.log.error("Cannot load flows: {}".format(e))
|
|
||||||
raise exceptions.FlowReadException(str(e)) from e
|
|
||||||
|
|
||||||
def running(self):
|
def running(self):
|
||||||
if ctx.options.rfile:
|
if ctx.options.rfile:
|
||||||
@ -48,3 +45,12 @@ class ReadFile:
|
|||||||
raise exceptions.OptionsError(e) from e
|
raise exceptions.OptionsError(e) from e
|
||||||
finally:
|
finally:
|
||||||
ctx.master.addons.trigger("processing_complete")
|
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)
|
||||||
|
@ -17,8 +17,9 @@ from mitmproxy import exceptions
|
|||||||
from mitmproxy import master
|
from mitmproxy import master
|
||||||
from mitmproxy import io
|
from mitmproxy import io
|
||||||
from mitmproxy import log
|
from mitmproxy import log
|
||||||
from mitmproxy.addons import view
|
|
||||||
from mitmproxy.addons import intercept
|
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 flowlist
|
||||||
from mitmproxy.tools.console import flowview
|
from mitmproxy.tools.console import flowview
|
||||||
from mitmproxy.tools.console import grideditor
|
from mitmproxy.tools.console import grideditor
|
||||||
@ -91,7 +92,12 @@ class ConsoleMaster(master.Master):
|
|||||||
signals.sig_add_log.connect(self.sig_add_log)
|
signals.sig_add_log.connect(self.sig_add_log)
|
||||||
self.addons.add(Logger())
|
self.addons.add(Logger())
|
||||||
self.addons.add(*addons.default_addons())
|
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):
|
def sigint_handler(*args, **kwargs):
|
||||||
self.prompt_for_exit()
|
self.prompt_for_exit()
|
||||||
|
@ -1,7 +1,7 @@
|
|||||||
from mitmproxy import addons
|
from mitmproxy import addons
|
||||||
from mitmproxy import options
|
from mitmproxy import options
|
||||||
from mitmproxy import master
|
from mitmproxy import master
|
||||||
from mitmproxy.addons import dumper, termlog, termstatus, keepserving
|
from mitmproxy.addons import dumper, termlog, termstatus, keepserving, readfile
|
||||||
|
|
||||||
|
|
||||||
class ErrorCheck:
|
class ErrorCheck:
|
||||||
@ -31,5 +31,6 @@ class DumpMaster(master.Master):
|
|||||||
self.addons.add(dumper.Dumper())
|
self.addons.add(dumper.Dumper())
|
||||||
self.addons.add(
|
self.addons.add(
|
||||||
keepserving.KeepServing(),
|
keepserving.KeepServing(),
|
||||||
|
readfile.ReadFileStdin(),
|
||||||
self.errorcheck
|
self.errorcheck
|
||||||
)
|
)
|
||||||
|
@ -7,6 +7,7 @@ from mitmproxy import log
|
|||||||
from mitmproxy import master
|
from mitmproxy import master
|
||||||
from mitmproxy.addons import eventstore
|
from mitmproxy.addons import eventstore
|
||||||
from mitmproxy.addons import intercept
|
from mitmproxy.addons import intercept
|
||||||
|
from mitmproxy.addons import readfile
|
||||||
from mitmproxy.addons import termlog
|
from mitmproxy.addons import termlog
|
||||||
from mitmproxy.addons import view
|
from mitmproxy.addons import view
|
||||||
from mitmproxy.addons import termstatus
|
from mitmproxy.addons import termstatus
|
||||||
@ -32,6 +33,7 @@ class WebMaster(master.Master):
|
|||||||
self.addons.add(*addons.default_addons())
|
self.addons.add(*addons.default_addons())
|
||||||
self.addons.add(
|
self.addons.add(
|
||||||
intercept.Intercept(),
|
intercept.Intercept(),
|
||||||
|
readfile.ReadFile(),
|
||||||
self.view,
|
self.view,
|
||||||
self.events,
|
self.events,
|
||||||
)
|
)
|
||||||
|
@ -37,59 +37,68 @@ def corrupt_data():
|
|||||||
return f
|
return f
|
||||||
|
|
||||||
|
|
||||||
@mock.patch('mitmproxy.master.Master.load_flow')
|
class TestReadFile:
|
||||||
def test_configure(mck, tmpdir, data, corrupt_data):
|
@mock.patch('mitmproxy.master.Master.load_flow')
|
||||||
rf = readfile.ReadFile()
|
def test_configure(self, mck, tmpdir, data, corrupt_data):
|
||||||
with taddons.context() as tctx:
|
rf = readfile.ReadFile()
|
||||||
tf = tmpdir.join("tfile")
|
with taddons.context() as tctx:
|
||||||
|
tf = tmpdir.join("tfile")
|
||||||
|
|
||||||
tf.write(data.getvalue())
|
tf.write(data.getvalue())
|
||||||
tctx.configure(rf, rfile=str(tf))
|
tctx.configure(rf, rfile=str(tf))
|
||||||
assert not mck.called
|
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()
|
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')
|
class TestReadFileStdin:
|
||||||
@mock.patch('sys.stdin')
|
@mock.patch('mitmproxy.master.Master.load_flow')
|
||||||
def test_configure_stdin(stdin, load_flow, data, corrupt_data):
|
@mock.patch('sys.stdin')
|
||||||
rf = readfile.ReadFile()
|
def test_stdin(self, stdin, load_flow, data, corrupt_data):
|
||||||
with taddons.context() as tctx:
|
rf = readfile.ReadFileStdin()
|
||||||
stdin.buffer = data
|
with taddons.context() as tctx:
|
||||||
tctx.configure(rf, rfile="-")
|
stdin.buffer = data
|
||||||
assert not load_flow.called
|
tctx.configure(rf, rfile="-")
|
||||||
rf.running()
|
assert not load_flow.called
|
||||||
assert load_flow.called
|
|
||||||
|
|
||||||
stdin.buffer = corrupt_data
|
|
||||||
tctx.configure(rf, rfile="-")
|
|
||||||
with pytest.raises(exceptions.OptionsError):
|
|
||||||
rf.running()
|
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')
|
@mock.patch('mitmproxy.master.Master.load_flow')
|
||||||
def test_corrupt(mck, corrupt_data):
|
def test_normal(self, load_flow, tmpdir, data):
|
||||||
rf = readfile.ReadFile()
|
rf = readfile.ReadFileStdin()
|
||||||
with taddons.context() as tctx:
|
with taddons.context():
|
||||||
with pytest.raises(exceptions.FlowReadException):
|
tfile = tmpdir.join("tfile")
|
||||||
rf.load_flows(io.BytesIO(b"qibble"))
|
tfile.write(data.getvalue())
|
||||||
assert not mck.called
|
rf.load_flows_from_path(str(tfile))
|
||||||
assert len(tctx.master.logs) == 1
|
assert load_flow.called
|
||||||
|
|
||||||
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
|
|
||||||
|
Loading…
Reference in New Issue
Block a user