aboutsummaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
Diffstat
-rw-r--r--mitmproxy/addonmanager.py6
-rw-r--r--mitmproxy/addons/save.py38
-rw-r--r--mitmproxy/command.py40
-rw-r--r--mitmproxy/exceptions.py2
-rw-r--r--mitmproxy/tools/console/command.py17
-rw-r--r--mitmproxy/utils/typecheck.py2
-rw-r--r--test/mitmproxy/addons/test_save.py21
-rw-r--r--test/mitmproxy/test_addonmanager.py4
8 files changed, 89 insertions, 41 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
diff --git a/test/mitmproxy/addons/test_save.py b/test/mitmproxy/addons/test_save.py
index fcd2feab..85c2a398 100644
--- a/test/mitmproxy/addons/test_save.py
+++ b/test/mitmproxy/addons/test_save.py
@@ -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:
diff --git a/test/mitmproxy/test_addonmanager.py b/test/mitmproxy/test_addonmanager.py
index 7b461580..fc1db941 100644
--- a/test/mitmproxy/test_addonmanager.py
+++ b/test/mitmproxy/test_addonmanager.py
@@ -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()
generated by cgit v1.2.3 (git 2.25.1) at 2024-06-18 11:36:46 +0000