FileStreamer -> StreamFile

options.wfile -> options.streamfile
This commit is contained in:
Aldo Cortesi 2016-11-03 15:00:18 +13:00
parent 43a8221989
commit 9f77c79227
8 changed files with 37 additions and 38 deletions

View File

@ -1,7 +1,7 @@
from mitmproxy.addons import anticache
from mitmproxy.addons import anticomp
from mitmproxy.addons import clientplayback
from mitmproxy.addons import filestreamer
from mitmproxy.addons import streamfile
from mitmproxy.addons import onboarding
from mitmproxy.addons import replace
from mitmproxy.addons import script
@ -20,7 +20,7 @@ def default_addons():
stickyauth.StickyAuth(),
stickycookie.StickyCookie(),
script.ScriptLoader(),
filestreamer.FileStreamer(),
streamfile.StreamFile(),
streambodies.StreamBodies(),
replace.Replace(),
setheaders.SetHeaders(),

View File

@ -5,7 +5,7 @@ from mitmproxy import flowfilter
from mitmproxy import io
class FileStreamer:
class StreamFile:
def __init__(self):
self.stream = None
self.filt = None
@ -29,11 +29,11 @@ class FileStreamer:
raise exceptions.OptionsError(
"Invalid filter specification: %s" % options.filtstr
)
if "outfile" in updated:
if "streamfile" in updated:
if self.stream:
self.done()
if options.outfile:
path, mode = options.outfile
if options.streamfile:
path, mode = options.streamfile
if mode not in ("wb", "ab"):
raise exceptions.OptionsError("Invalid mode.")
self.start_stream_to_path(path, mode, self.filt)

View File

@ -48,7 +48,7 @@ class Options(optmanager.OptManager):
stream_large_bodies: Optional[int] = None,
verbosity: int = 2,
default_contentview: str = "auto",
outfile: Optional[Tuple[str, str]] = None,
streamfile: Optional[Tuple[str, str]] = None,
server_replay_ignore_content: bool = False,
server_replay_ignore_params: Sequence[str] = (),
server_replay_ignore_payload_params: Sequence[str] = (),
@ -108,7 +108,7 @@ class Options(optmanager.OptManager):
self.stream_large_bodies = stream_large_bodies
self.verbosity = verbosity
self.default_contentview = default_contentview
self.outfile = outfile
self.streamfile = streamfile
self.server_replay_ignore_content = server_replay_ignore_content
self.server_replay_ignore_params = server_replay_ignore_params
self.server_replay_ignore_payload_params = server_replay_ignore_payload_params

View File

@ -140,8 +140,8 @@ def get_common_options(args):
raise exceptions.OptionsError(e)
setheaders.append(p)
if args.outfile and args.outfile[0] == args.rfile:
if args.outfile[1] == "wb":
if args.streamfile and args.streamfile[0] == args.rfile:
if args.streamfile[1] == "wb":
raise exceptions.OptionsError(
"Cannot use '{}' for both reading and writing flows. "
"Are you looking for --afile?".format(args.rfile)
@ -228,7 +228,7 @@ def get_common_options(args):
stickyauth=stickyauth,
stream_large_bodies=stream_large_bodies,
showhost=args.showhost,
outfile=args.outfile,
streamfile=args.streamfile,
verbosity=args.verbose,
server_replay_nopop=args.server_replay_nopop,
server_replay_ignore_content=args.server_replay_ignore_content,
@ -339,15 +339,15 @@ def basic_options(parser):
action="store_const", dest="verbose", default=2, const=3,
help="Increase log verbosity."
)
outfile = parser.add_mutually_exclusive_group()
outfile.add_argument(
streamfile = parser.add_mutually_exclusive_group()
streamfile.add_argument(
"-w", "--wfile",
action="store", dest="outfile", type=lambda f: (f, "wb"),
action="store", dest="streamfile", type=lambda f: (f, "wb"),
help="Write flows to file."
)
outfile.add_argument(
streamfile.add_argument(
"-a", "--afile",
action="store", dest="outfile", type=lambda f: (f, "ab"),
action="store", dest="streamfile", type=lambda f: (f, "ab"),
help="Append flows to file."
)
parser.add_argument(

View File

@ -393,13 +393,13 @@ class FlowListBox(urwid.ListBox):
val = not self.master.options.order_reversed
self.master.options.order_reversed = val
elif key == "W":
if self.master.options.outfile:
self.master.options.outfile = None
if self.master.options.streamfile:
self.master.options.streamfile = None
else:
signals.status_prompt_path.send(
self,
prompt="Stream flows to",
callback= lambda path: self.master.options.update(outfile=(path, "ab"))
callback= lambda path: self.master.options.update(streamfile=(path, "ab"))
)
else:
return urwid.ListBox.keypress(self, size, key)

View File

@ -223,8 +223,8 @@ class StatusBar(urwid.WidgetWrap):
r.append(("heading_key", "s"))
r.append("cripts:%s]" % len(self.master.options.scripts))
if self.master.options.outfile:
r.append("[W:%s]" % self.master.options.outfile[0])
if self.master.options.streamfile:
r.append("[W:%s]" % self.master.options.streamfile[0])
return r

View File

@ -111,10 +111,10 @@ class WebMaster(master.Master):
"error"
)
if options.outfile:
if options.streamfile:
err = self.start_stream_to_path(
options.outfile[0],
options.outfile[1]
options.streamfile[0],
options.streamfile[1]
)
if err:
print("Stream file error: {}".format(err), file=sys.stderr)

View File

@ -6,25 +6,25 @@ import os.path
from mitmproxy import io
from mitmproxy import exceptions
from mitmproxy.tools import dump
from mitmproxy.addons import filestreamer
from mitmproxy.addons import streamfile
def test_configure():
sa = filestreamer.FileStreamer()
sa = streamfile.StreamFile()
with taddons.context(options=dump.Options()) as tctx:
with tutils.tmpdir() as tdir:
p = os.path.join(tdir, "foo")
tutils.raises(
exceptions.OptionsError,
tctx.configure, sa, outfile=(tdir, "ab")
tctx.configure, sa, streamfile=(tdir, "ab")
)
tutils.raises(
"invalid filter",
tctx.configure, sa, outfile=(p, "ab"), filtstr="~~"
tctx.configure, sa, streamfile=(p, "ab"), filtstr="~~"
)
tutils.raises(
"invalid mode",
tctx.configure, sa, outfile=(p, "xx")
tctx.configure, sa, streamfile=(p, "xx")
)
@ -34,36 +34,35 @@ def rd(p):
def test_tcp():
sa = filestreamer.FileStreamer()
sa = streamfile.StreamFile()
with taddons.context() as tctx:
with tutils.tmpdir() as tdir:
p = os.path.join(tdir, "foo")
tctx.configure(sa, outfile=(p, "wb"))
tctx.configure(sa, streamfile=(p, "wb"))
tt = tflow.ttcpflow()
sa.tcp_start(tt)
sa.tcp_end(tt)
tctx.configure(sa, outfile=None)
tctx.configure(sa, streamfile=None)
assert rd(p)
def test_simple():
sa = filestreamer.FileStreamer()
sa = streamfile.StreamFile()
with taddons.context() as tctx:
with tutils.tmpdir() as tdir:
p = os.path.join(tdir, "foo")
tctx.configure(sa, outfile=(p, "wb"))
tctx.configure(sa, streamfile=(p, "wb"))
f = tflow.tflow(resp=True)
sa.request(f)
sa.response(f)
tctx.configure(sa, outfile=None)
tctx.configure(sa, streamfile=None)
assert rd(p)[0].response
tctx.configure(sa, outfile=(p, "ab"))
tctx.configure(sa, streamfile=(p, "ab"))
f = tflow.tflow()
sa.request(f)
tctx.configure(sa, outfile=None)
tctx.configure(sa, streamfile=None)
assert not rd(p)[1].response