Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2025-01-30 14:58:38 +00:00
Code Issues Projects Releases Wiki Activity

command: save.file flowspec path -> None

Browse Source 
Our first user-facing command. The following commands do the obvious things:

save.file @marked /tmp/flows
save.file @focus /tmp/flows
save.file @hidden /tmp/flows
save.file "~m get" /tmp/flows
This commit is contained in:
Aldo Cortesi 2017-04-27 17:05:00 +12:00
parent b7afcb5dc2
commit 97000aa85c
8 changed files with 89 additions and 41 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
mitmproxy/addonmanager.py
View File

@ -145,7 +145,7 @@ class AddonManager:
for a in traverse([addon]):
name = _get_name(a)
if name in self.lookup:
raise exceptions.AddonError(
raise exceptions.AddonManagerError(
"An addon called '%s' already exists." % name
)
l = Loader(self.master)
@ -175,7 +175,7 @@ class AddonManager:
for a in traverse([addon]):
n = _get_name(a)
if n not in self.lookup:
raise exceptions.AddonError("No such addon: %s" % n)
raise exceptions.AddonManagerError("No such addon: %s" % n)
self.chain = [i for i in self.chain if i is not a]
del self.lookup[_get_name(a)]
with self.master.handlecontext():
@ -224,7 +224,7 @@ class AddonManager:
func = getattr(a, name, None)
if func:
if not callable(func):
raise exceptions.AddonError(
raise exceptions.AddonManagerError(
"Addon handler %s not callable" % name
)
func(*args, **kwargs)

38
mitmproxy/addons/save.py
View File

@ -1,9 +1,11 @@
import os.path
import typing
from mitmproxy import exceptions
from mitmproxy import flowfilter
from mitmproxy import io
from mitmproxy import ctx
from mitmproxy import flow
class Save:
@ -12,10 +14,18 @@ class Save:
self.filt = None
self.active_flows = set() # type: Set[flow.Flow]
def start_stream_to_path(self, path, mode, flt):
def open_file(self, path):
if path.startswith("+"):
path = path[1:]
mode = "ab"
else:
mode = "wb"
path = os.path.expanduser(path)
return open(path, mode)
def start_stream_to_path(self, path, flt):
try:
f = open(path, mode)
f = self.open_file(path)
except IOError as v:
raise exceptions.OptionsError(str(v))
self.stream = io.FilteredFlowWriter(f, flt)
@ -36,13 +46,19 @@ class Save:
if self.stream:
self.done()
if ctx.options.save_stream_file:
if ctx.options.save_stream_file.startswith("+"):
path = ctx.options.save_stream_file[1:]
mode = "ab"
else:
path = ctx.options.save_stream_file
mode = "wb"
self.start_stream_to_path(path, mode, self.filt)
self.start_stream_to_path(ctx.options.save_stream_file, self.filt)
def save(self, flows: typing.Sequence[flow.Flow], path: str) -> None:
try:
f = self.open_file(path)
except IOError as v:
raise exceptions.CommandError(v) from v
stream = io.FlowWriter(f)
for i in flows:
stream.add(i)
def load(self, l):
l.add_command("save.file", self.save)
def tcp_start(self, flow):
if self.stream:
@ -64,8 +80,8 @@ class Save:
def done(self):
if self.stream:
for flow in self.active_flows:
self.stream.add(flow)
for f in self.active_flows:
self.stream.add(f)
self.active_flows = set([])
self.stream.fo.close()
self.stream = None

40
mitmproxy/command.py
View File

@ -6,25 +6,19 @@ from mitmproxy import exceptions
from mitmproxy import flow
def typename(t: type) -> str:
def typename(t: type, ret: bool) -> str:
"""
Translates a type to an explanatory string. Ifl ret is True, we're
looking at a return type, else we're looking at a parameter type.
"""
if t in (str, int, bool):
return t.__name__
if t == typing.Sequence[flow.Flow]:
return "[flow]"
return "[flow]" if ret else "flowspec"
else: # pragma: no cover
raise NotImplementedError(t)
def parsearg(spec: str, argtype: type) -> typing.Any:
"""
Convert a string to a argument to the appropriate type.
"""
if argtype == str:
return spec
else:
raise exceptions.CommandError("Unsupported argument type: %s" % argtype)
class Command:
def __init__(self, manager, path, func) -> None:
self.path = path
@ -35,8 +29,8 @@ class Command:
self.returntype = sig.return_annotation
def signature_help(self) -> str:
params = " ".join([typename(i) for i in self.paramtypes])
ret = " -> " + typename(self.returntype) if self.returntype else ""
params = " ".join([typename(i, False) for i in self.paramtypes])
ret = " -> " + typename(self.returntype, True) if self.returntype else ""
return "%s %s%s" % (self.path, params, ret)
def call(self, args: typing.Sequence[str]):
@ -46,10 +40,12 @@ class Command:
if len(self.paramtypes) != len(args):
raise exceptions.CommandError("Usage: %s" % self.signature_help())
args = [parsearg(args[i], self.paramtypes[i]) for i in range(len(args))]
pargs = []
for i in range(len(args)):
pargs.append(parsearg(self.manager, args[i], self.paramtypes[i]))
with self.manager.master.handlecontext():
ret = self.func(*args)
ret = self.func(*pargs)
if not typecheck.check_command_return_type(ret, self.returntype):
raise exceptions.CommandError("Command returned unexpected data")
@ -81,3 +77,15 @@ class CommandManager:
if not len(parts) >= 1:
raise exceptions.CommandError("Invalid command: %s" % cmdstr)
return self.call_args(parts[0], parts[1:])
def parsearg(manager: CommandManager, spec: str, argtype: type) -> typing.Any:
"""
Convert a string to a argument to the appropriate type.
"""
if argtype == str:
return spec
elif argtype == typing.Sequence[flow.Flow]:
return manager.call_args("console.resolve", [spec])
else:
raise exceptions.CommandError("Unsupported argument type: %s" % argtype)

2
mitmproxy/exceptions.py
View File

@ -101,7 +101,7 @@ class OptionsError(MitmproxyException):
pass
class AddonError(MitmproxyException):
class AddonManagerError(MitmproxyException):
pass

17
mitmproxy/tools/console/command.py
View File

@ -1,6 +1,6 @@
import urwid
from mitmproxy import command
from mitmproxy import exceptions
from mitmproxy.tools.console import signals
@ -17,10 +17,11 @@ class CommandExecutor:
self.master = master
def __call__(self, cmd):
try:
ret = self.master.commands.call(cmd)
except command.CommandError as v:
signals.status_message.send(message=str(v))
else:
if type(ret) == str:
signals.status_message.send(message=ret)
if cmd.strip():
try:
ret = self.master.commands.call(cmd)
except exceptions.CommandError as v:
signals.status_message.send(message=str(v))
else:
if type(ret) == str:
signals.status_message.send(message=ret)

2
mitmproxy/utils/typecheck.py
View File

@ -19,6 +19,8 @@ def check_command_return_type(value: typing.Any, typeinfo: typing.Any) -> bool:
for v in value:
if not check_command_return_type(v, T):
return False
elif value is None and typeinfo is None:
return True
elif not isinstance(value, typeinfo):
return False
return True

21
test/mitmproxy/addons/test_save.py
View File

@ -7,6 +7,7 @@ from mitmproxy import io
from mitmproxy import exceptions
from mitmproxy import options
from mitmproxy.addons import save
from mitmproxy.addons import view
def test_configure(tmpdir):
@ -42,6 +43,26 @@ def test_tcp(tmpdir):
assert rd(p)
def test_save_command(tmpdir):
sa = save.Save()
with taddons.context() as tctx:
p = str(tmpdir.join("foo"))
sa.save([tflow.tflow(resp=True)], p)
assert len(rd(p)) == 1
sa.save([tflow.tflow(resp=True)], p)
assert len(rd(p)) == 1
sa.save([tflow.tflow(resp=True)], "+" + p)
assert len(rd(p)) == 2
with pytest.raises(exceptions.CommandError):
sa.save([tflow.tflow(resp=True)], str(tmpdir))
v = view.View()
tctx.master.addons.add(v)
tctx.master.addons.add(sa)
tctx.master.commands.call_args("save.file", ["@shown", p])
def test_simple(tmpdir):
sa = save.Save()
with taddons.context() as tctx:

4
test/mitmproxy/test_addonmanager.py
View File

@ -61,9 +61,9 @@ def test_lifecycle():
a = addonmanager.AddonManager(m)
a.add(TAddon("one"))
with pytest.raises(exceptions.AddonError):
with pytest.raises(exceptions.AddonManagerError):
a.add(TAddon("one"))
with pytest.raises(exceptions.AddonError):
with pytest.raises(exceptions.AddonManagerError):
a.remove(TAddon("nonexistent"))
f = tflow.tflow()