From 091db9cdf206be16ca9492405402378f7aead136 Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Fri, 28 Jan 2011 12:55:02 +1300
Subject: [PATCH] Prompt user on quit.

---
 libmproxy/console.py | 22 +++++++++++++++++-----
 mitmproxy            |  5 +----
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/libmproxy/console.py b/libmproxy/console.py
index 757a74970..907d684fd 100644
--- a/libmproxy/console.py
+++ b/libmproxy/console.py
@@ -341,7 +341,7 @@ class ConnectionView(WWrap):
             if not conn.set_url(url):
                 return "Invalid URL."
         elif part == "m" and self.viewing == self.REQ:
-            self.master.prompt_onekey("Method ", self.methods, self.edit_method)
+            self.master.prompt_onekey("Method", self.methods, self.edit_method)
             key = None
         self.master.refresh_connection(self.flow)
 
@@ -372,7 +372,7 @@ class ConnectionView(WWrap):
         elif key == "e":
             if self.viewing == self.REQ:
                 self.master.prompt_onekey(
-                    "Edit request ",
+                    "Edit request",
                     (
                         ("header", "h"),
                         ("body", "b"),
@@ -383,7 +383,7 @@ class ConnectionView(WWrap):
                 )
             else:
                 self.master.prompt_onekey(
-                    "Edit response ",
+                    "Edit response",
                     (
                         ("header", "h"),
                         ("body", "b"),
@@ -946,7 +946,7 @@ class ConsoleMaster(controller.Master):
             Keys are a set of (word, key) tuples. The appropriate key in the
             word is highlighted. 
         """
-        prompt = [prompt, "("]
+        prompt = [prompt, " ("]
         mkup = []
         for i, e in enumerate(keys):
             parts = e[0].split(e[1], 1)
@@ -1076,7 +1076,15 @@ class ConsoleMaster(controller.Master):
                             if self.nested:
                                 self.view_connlist()
                             else:
-                                raise Stop
+                                self.prompt_onekey(
+                                    "Quit",
+                                    (
+                                        ("yes", "y"),
+                                        ("no", "n"),
+                                    ),
+                                    self.quit,
+                                )
+                                k = None
                         elif k == "S":
                             self.path_prompt("Save flows: ", self.save_flows)
                             k = None
@@ -1091,6 +1099,10 @@ class ConsoleMaster(controller.Master):
         except (Stop, KeyboardInterrupt):
             pass
 
+    def quit(self, a):
+        if a != "n":
+            raise Stop
+
     def shutdown(self):
         for i in self.state.flow_list:
             i.kill()
diff --git a/mitmproxy b/mitmproxy
index 4153ea44b..2fc61a80d 100755
--- a/mitmproxy
+++ b/mitmproxy
@@ -80,10 +80,7 @@ if __name__ == '__main__':
     m = console.ConsoleMaster(server, options)
 
     for i in args:
-        msg = m.load_flows(i)
-        if msg:
-            print >> sys.stderr, msg
-            sys.exit(1)
+        m.load_flows(i)
     m.run()