mirror of
https://github.com/Grasscutters/mitmproxy.git
synced 2024-11-23 00:01:36 +00:00
integrate readstdin into readfile addon
This commit is contained in:
Maximilian Hils
parent
2b500f234f
commit
b3a1143338
@ -1,4 +1,7 @@
|
||||
import os.path
|
||||
import typing
|
||||
|
||||
import sys
|
||||
|
||||
from mitmproxy import ctx
|
||||
from mitmproxy import io
|
||||
@ -9,30 +12,39 @@ 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:
|
||||
if path == "-":
|
||||
return self.load_flows(sys.stdin.buffer)
|
||||
else:
|
||||
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")
|
||||
|
||||
@ -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")
|
||||
@ -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
|
||||
|
||||
|
||||
class ErrorCheck:
|
||||
@ -30,7 +30,6 @@ class DumpMaster(master.Master):
|
||||
if with_dumper:
|
||||
self.addons.add(dumper.Dumper())
|
||||
self.addons.add(
|
||||
readstdin.ReadStdin(),
|
||||
keepserving.KeepServing(),
|
||||
self.errorcheck
|
||||
)
|
||||
|
||||
@ -1,62 +1,95 @@
|
||||
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
|
||||
|
||||
import mitmproxy.io
|
||||
from mitmproxy import exceptions
|
||||
from mitmproxy.addons import readfile
|
||||
from mitmproxy.test import taddons
|
||||
from mitmproxy.test import tflow
|
||||
|
||||
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")
|
||||
|
||||
@pytest.fixture
|
||||
def data():
|
||||
f = io.BytesIO()
|
||||
|
||||
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)
|
||||
|
||||
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
|
||||
|
||||
|
||||
@mock.patch('mitmproxy.master.Master.load_flow')
|
||||
def test_configure(mck, tmpdir):
|
||||
|
||||
def test_configure(mck, tmpdir, data, corrupt_data):
|
||||
rf = readfile.ReadFile()
|
||||
with taddons.context() as tctx:
|
||||
tf = str(tmpdir.join("tfile"))
|
||||
write_data(tf)
|
||||
tf = tmpdir.join("tfile")
|
||||
|
||||
tf.write(data.getvalue())
|
||||
tctx.configure(rf, rfile=str(tf))
|
||||
assert not mck.called
|
||||
rf.running()
|
||||
assert mck.called
|
||||
|
||||
write_data(tf, corrupt=True)
|
||||
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_corruption(mck, tmpdir):
|
||||
@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):
|
||||
rf = readfile.ReadFile()
|
||||
with taddons.context() as tctx:
|
||||
with pytest.raises(exceptions.FlowReadException):
|
||||
rf.load_flows_file("nonexistent")
|
||||
rf.load_flows(io.BytesIO(b"qibble"))
|
||||
assert not mck.called
|
||||
assert len(tctx.master.logs) == 1
|
||||
|
||||
tfc = str(tmpdir.join("tfile"))
|
||||
write_data(tfc, corrupt=True)
|
||||
|
||||
with pytest.raises(exceptions.FlowReadException):
|
||||
rf.load_flows_file(tfc)
|
||||
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
|
||||
|
||||
@ -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
|
||||
Loading…
Reference in New Issue
Block a user