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 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,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)
|
||||
|
@ -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, 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
|
||||
)
|
||||
|
@ -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,
|
||||
)
|
||||
|
@ -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
|
||||
|
Loading…
Reference in New Issue
Block a user