diff options
author | Aldo Cortesi <aldo@nullcube.com> | 2017-04-27 17:05:00 +1200 |
---|---|---|
committer | Aldo Cortesi <aldo@nullcube.com> | 2017-04-27 17:05:00 +1200 |
commit | 97000aa85c6d4cd52d97bda66b0c05ffed93e9da (patch) | |
tree | b5ed172d016a6341e98d2eb765200ba39446b2c2 /mitmproxy | |
parent | b7afcb5dc2007ce451161e513d3e55d4cc91379e (diff) | |
download | mitmproxy-97000aa85c6d4cd52d97bda66b0c05ffed93e9da.tar.gz mitmproxy-97000aa85c6d4cd52d97bda66b0c05ffed93e9da.tar.bz2 mitmproxy-97000aa85c6d4cd52d97bda66b0c05ffed93e9da.zip |
command: save.file flowspec path -> None
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
Diffstat (limited to 'mitmproxy')
-rw-r--r-- | mitmproxy/addonmanager.py | 6 | ||||
-rw-r--r-- | mitmproxy/addons/save.py | 38 | ||||
-rw-r--r-- | mitmproxy/command.py | 40 | ||||
-rw-r--r-- | mitmproxy/exceptions.py | 2 | ||||
-rw-r--r-- | mitmproxy/tools/console/command.py | 17 | ||||
-rw-r--r-- | mitmproxy/utils/typecheck.py | 2 |
6 files changed, 66 insertions, 39 deletions
diff --git a/mitmproxy/addonmanager.py b/mitmproxy/addonmanager.py index 5cfe79f5..ea23b6ff 100644 --- a/mitmproxy/addonmanager.py +++ b/mitmproxy/addonmanager.py @@ -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) diff --git a/mitmproxy/addons/save.py b/mitmproxy/addons/save.py index 9cf65750..37dc6021 100644 --- a/mitmproxy/addons/save.py +++ b/mitmproxy/addons/save.py @@ -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 diff --git a/mitmproxy/command.py b/mitmproxy/command.py index 3d24675b..acf938d5 100644 --- a/mitmproxy/command.py +++ b/mitmproxy/command.py @@ -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) diff --git a/mitmproxy/exceptions.py b/mitmproxy/exceptions.py index 04525d1f..71517480 100644 --- a/mitmproxy/exceptions.py +++ b/mitmproxy/exceptions.py @@ -101,7 +101,7 @@ class OptionsError(MitmproxyException): pass -class AddonError(MitmproxyException): +class AddonManagerError(MitmproxyException): pass diff --git a/mitmproxy/tools/console/command.py b/mitmproxy/tools/console/command.py index e6eabf3a..4cb4fe6d 100644 --- a/mitmproxy/tools/console/command.py +++ b/mitmproxy/tools/console/command.py @@ -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) diff --git a/mitmproxy/utils/typecheck.py b/mitmproxy/utils/typecheck.py index 33dd70b0..20791e17 100644 --- a/mitmproxy/utils/typecheck.py +++ b/mitmproxy/utils/typecheck.py @@ -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 |