diff options
-rw-r--r-- | mitmproxy/tools/console/common.py | 3 | ||||
-rw-r--r-- | mitmproxy/tools/console/master.py | 18 | ||||
-rw-r--r-- | mitmproxy/tools/console/options.py | 1 | ||||
-rw-r--r-- | mitmproxy/tools/console/pathedit.py | 4 | ||||
-rw-r--r-- | mitmproxy/tools/console/statusbar.py | 39 | ||||
-rw-r--r-- | test/mitmproxy/console/test_pathedit.py | 2 |
6 files changed, 40 insertions, 27 deletions
diff --git a/mitmproxy/tools/console/common.py b/mitmproxy/tools/console/common.py index a8e7f59f..ec637cbc 100644 --- a/mitmproxy/tools/console/common.py +++ b/mitmproxy/tools/console/common.py @@ -150,9 +150,6 @@ def save_data(path, data): def ask_save_overwrite(path, data): - if not path: - return - path = os.path.expanduser(path) if os.path.exists(path): def save_overwrite(k): if k == "y": diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py index 73d7adbd..536dc03d 100644 --- a/mitmproxy/tools/console/master.py +++ b/mitmproxy/tools/console/master.py @@ -386,17 +386,11 @@ class ConsoleMaster(master.Master): ) def _write_flows(self, path, flows): - if not path: - return - path = os.path.expanduser(path) - try: - f = open(path, "wb") - fw = io.FlowWriter(f) - for i in flows: - fw.add(i) - f.close() - except IOError as v: - signals.status_message.send(message=v.strerror) + f = open(path, "wb") + fw = io.FlowWriter(f) + for i in flows: + fw.add(i) + f.close() def save_one_flow(self, path, flow): return self._write_flows(path, [flow]) @@ -405,8 +399,6 @@ class ConsoleMaster(master.Master): return self._write_flows(path, self.view) def load_flows_callback(self, path): - if not path: - return ret = self.load_flows_path(path) return ret or "Flows loaded from %s" % path diff --git a/mitmproxy/tools/console/options.py b/mitmproxy/tools/console/options.py index 8932d850..54876f87 100644 --- a/mitmproxy/tools/console/options.py +++ b/mitmproxy/tools/console/options.py @@ -171,6 +171,7 @@ class Options(urwid.WidgetWrap): def do_save(self, path): self.master.options.save(path) + return "Saved" def save(self): signals.status_prompt_path.send( diff --git a/mitmproxy/tools/console/pathedit.py b/mitmproxy/tools/console/pathedit.py index 4447070b..10ee1416 100644 --- a/mitmproxy/tools/console/pathedit.py +++ b/mitmproxy/tools/console/pathedit.py @@ -57,8 +57,8 @@ class _PathCompleter: class PathEdit(urwid.Edit, _PathCompleter): - def __init__(self, *args, **kwargs): - urwid.Edit.__init__(self, *args, **kwargs) + def __init__(self, prompt, last_path): + urwid.Edit.__init__(self, prompt, last_path) _PathCompleter.__init__(self) def keypress(self, size, key): diff --git a/mitmproxy/tools/console/statusbar.py b/mitmproxy/tools/console/statusbar.py index e3424493..dce8605f 100644 --- a/mitmproxy/tools/console/statusbar.py +++ b/mitmproxy/tools/console/statusbar.py @@ -9,6 +9,28 @@ from mitmproxy.tools.console import signals from mitmproxy.utils import human +class PromptPath: + def __init__(self, callback, args): + self.callback, self.args = callback, args + + def __call__(self, pth): + if not pth: + return + pth = os.path.expanduser(pth) + try: + return self.callback(pth, *self.args) + except IOError as v: + signals.status_message.send(message=v.strerror) + + +class PromptStub: + def __init__(self, callback, args): + self.callback, self.args = callback, args + + def __call__(self, txt): + return self.callback(txt, *self.args) + + class ActionBar(urwid.WidgetWrap): def __init__(self): @@ -21,7 +43,8 @@ class ActionBar(urwid.WidgetWrap): self.last_path = "" - self.prompting = False + self.prompting = None + self.onekey = False self.pathprompt = False @@ -42,7 +65,7 @@ class ActionBar(urwid.WidgetWrap): def sig_prompt(self, sender, prompt, text, callback, args=()): signals.focus.send(self, section="footer") self._w = urwid.Edit(self.prep_prompt(prompt), text or "") - self.prompting = (callback, args) + self.prompting = PromptStub(callback, args) def sig_path_prompt(self, sender, prompt, callback, args=()): signals.focus.send(self, section="footer") @@ -51,7 +74,7 @@ class ActionBar(urwid.WidgetWrap): os.path.dirname(self.last_path) ) self.pathprompt = True - self.prompting = (callback, args) + self.prompting = PromptPath(callback, args) def sig_prompt_onekey(self, sender, prompt, keys, callback, args=()): """ @@ -69,7 +92,7 @@ class ActionBar(urwid.WidgetWrap): prompt.append(")? ") self.onekey = set(i[1] for i in keys) self._w = urwid.Edit(prompt, "") - self.prompting = (callback, args) + self.prompting = PromptStub(callback, args) def selectable(self): return True @@ -93,10 +116,10 @@ class ActionBar(urwid.WidgetWrap): def clear(self): self._w = urwid.Text("") - self.prompting = False + self.prompting = None def prompt_done(self): - self.prompting = False + self.prompting = None self.onekey = False self.pathprompt = False signals.status_message.send(message="") @@ -105,9 +128,9 @@ class ActionBar(urwid.WidgetWrap): def prompt_execute(self, txt): if self.pathprompt: self.last_path = txt - p, args = self.prompting + p = self.prompting self.prompt_done() - msg = p(txt, *args) + msg = p(txt) if msg: signals.status_message.send(message=msg, expire=1) diff --git a/test/mitmproxy/console/test_pathedit.py b/test/mitmproxy/console/test_pathedit.py index 40d55353..b326ed6d 100644 --- a/test/mitmproxy/console/test_pathedit.py +++ b/test/mitmproxy/console/test_pathedit.py @@ -54,7 +54,7 @@ class TestPathEdit: def test_keypress(self): - pe = pathedit.PathEdit() + pe = pathedit.PathEdit("", "") with patch('urwid.widget.Edit.get_edit_text') as get_text, \ patch('urwid.widget.Edit.set_edit_text') as set_text: |