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

Merge pull request #3056 from cortesi/readfile

Browse Source 
readfile fixes
This commit is contained in:
Aldo Cortesi 2018-04-15 10:03:30 +12:00 committed by GitHub
commit 4e126c0fba
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 102 additions and 57 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

55
mitmproxy/addons/readfile.py
View File

@ -1,9 +1,11 @@
import asyncio
import os.path
import sys
import typing
from mitmproxy import ctx
from mitmproxy import exceptions
from mitmproxy import flowfilter
from mitmproxy import io
@ -11,18 +13,38 @@ class ReadFile:
"""
An addon that handles reading from file on startup.
"""
def __init__(self):
self.filter = None
def load(self, loader):
loader.add_option(
"rfile", typing.Optional[str], None,
"Read flows from file."
)
loader.add_option(
"readfile_filter", typing.Optional[str], None,
"Read only matching flows."
)
def load_flows(self, fo: typing.IO[bytes]) -> int:
def configure(self, updated):
if "readfile_filter" in updated:
filt = None
if ctx.options.readfile_filter:
filt = flowfilter.parse(ctx.options.readfile_filter)
if not filt:
raise exceptions.OptionsError(
"Invalid readfile filter: %s" % ctx.options.readfile_filter
)
self.filter = filt
async def load_flows(self, fo: typing.IO[bytes]) -> int:
cnt = 0
freader = io.FlowReader(fo)
try:
for flow in freader.stream():
ctx.master.load_flow(flow)
if self.filter and not self.filter(flow):
continue
await ctx.master.load_flow(flow)
cnt += 1
except (IOError, exceptions.FlowReadException) as e:
if cnt:
@ -33,29 +55,34 @@ class ReadFile:
else:
return cnt
def load_flows_from_path(self, path: str) -> int:
async 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)
return await self.load_flows(f)
except IOError as e:
ctx.log.error("Cannot load flows: {}".format(e))
raise exceptions.FlowReadException(str(e)) from e
async def doread(self, rfile):
try:
await 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")
def running(self):
if ctx.options.rfile:
try:
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")
asyncio.get_event_loop().create_task(self.doread(ctx.options.rfile))
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)
async def load_flows_from_path(self, path: str) -> int:
if path == "-": # pragma: no cover
# Need to think about how to test this. This function is scheduled
# onto the event loop, where a sys.stdin mock has no effect.
return await self.load_flows(sys.stdin.buffer)
else:
return super().load_flows_from_path(path)
return await super().load_flows_from_path(path)

3
mitmproxy/tools/main.py
View File

@ -150,8 +150,9 @@ def mitmdump(args=None): # pragma: no cover
if args.filter_args:
v = " ".join(args.filter_args)
return dict(
view_filter=v,
save_stream_filter=v,
readfile_filter=v,
dumper_filter=v,
)
return {}

1
setup.py
View File

@ -85,6 +85,7 @@ setup(
"pydivert>=2.0.3,<2.2",
],
'dev': [
"asynctest>=0.12.0",
"flake8>=3.5, <3.6",
"Flask>=0.10.1, <0.13",
"mypy>=0.580,<0.581",

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

@ -1,7 +1,9 @@
import asyncio
import io
from unittest import mock
import pytest
import asynctest
import mitmproxy.io
from mitmproxy import exceptions
@ -38,69 +40,83 @@ def corrupt_data():
class TestReadFile:
@mock.patch('mitmproxy.master.Master.load_flow')
def test_configure(self, mck, tmpdir, data, corrupt_data):
def test_configure(self):
rf = readfile.ReadFile()
with taddons.context() as tctx:
tctx.configure(rf, readfile_filter="~q")
with pytest.raises(Exception, match="Invalid readfile filter"):
tctx.configure(rf, readfile_filter="~~")
@pytest.mark.asyncio
async def test_read(self, tmpdir, data, corrupt_data):
rf = readfile.ReadFile()
with taddons.context(rf) 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
with asynctest.patch('mitmproxy.master.Master.load_flow') as mck:
tf.write(data.getvalue())
tctx.configure(
rf,
rfile = str(tf),
readfile_filter = ".*"
)
assert not mck.awaited
rf.running()
await asyncio.sleep(0)
assert mck.awaited
tf.write(corrupt_data.getvalue())
tctx.configure(rf, rfile=str(tf))
with pytest.raises(exceptions.OptionsError):
rf.running()
rf.running()
assert await tctx.master.await_log("corrupted")
@pytest.mark.asyncio
async def test_corrupt(self, corrupt_data):
rf = readfile.ReadFile()
with taddons.context(rf) as tctx:
with mock.patch('mitmproxy.master.Master.load_flow') as mck:
with pytest.raises(exceptions.FlowReadException):
rf.load_flows(io.BytesIO(b"qibble"))
assert not mck.called
with pytest.raises(exceptions.FlowReadException):
await rf.load_flows(io.BytesIO(b"qibble"))
tctx.master.clear()
with pytest.raises(exceptions.FlowReadException):
rf.load_flows(corrupt_data)
assert await tctx.master.await_log("file corrupted")
assert mck.called
tctx.master.clear()
with pytest.raises(exceptions.FlowReadException):
await rf.load_flows(corrupt_data)
assert await tctx.master.await_log("file corrupted")
@pytest.mark.asyncio
async def test_nonexisting_file(self):
async def test_nonexistent_file(self):
rf = readfile.ReadFile()
with taddons.context(rf) as tctx:
with pytest.raises(exceptions.FlowReadException):
rf.load_flows_from_path("nonexistent")
await rf.load_flows_from_path("nonexistent")
assert await tctx.master.await_log("nonexistent")
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(rf) 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):
@asynctest.patch('sys.stdin')
@pytest.mark.asyncio
async def test_stdin(self, stdin, data, corrupt_data):
rf = readfile.ReadFileStdin()
with taddons.context(rf):
tfile = tmpdir.join("tfile")
tfile.write(data.getvalue())
rf.load_flows_from_path(str(tfile))
assert load_flow.called
with asynctest.patch('mitmproxy.master.Master.load_flow') as mck:
stdin.buffer = data
assert not mck.awaited
await rf.load_flows(stdin.buffer)
assert mck.awaited
stdin.buffer = corrupt_data
with pytest.raises(exceptions.FlowReadException):
await rf.load_flows(stdin.buffer)
@pytest.mark.asyncio
@mock.patch('mitmproxy.master.Master.load_flow')
async def test_normal(self, load_flow, tmpdir, data):
rf = readfile.ReadFileStdin()
with taddons.context(rf) as tctx:
tf = tmpdir.join("tfile")
with asynctest.patch('mitmproxy.master.Master.load_flow') as mck:
tf.write(data.getvalue())
tctx.configure(rf, rfile=str(tf))
assert not mck.awaited
rf.running()
await asyncio.sleep(0)
assert mck.awaited