Merge pull request #624 from drahosj/master

Added flow marking functionality in the console

4 files changed, 63 insertions, 5 deletions

diff --git a/libmproxy/console/__init__.py b/libmproxy/console/__init__.py
index 052ac7dd..3d20947b 100644
--- a/libmproxy/console/__init__.py
+++ b/libmproxy/console/__init__.py
@@ -48,6 +48,7 @@ class ConsoleState(flow.State):
             self.set_focus(0)
         elif self.follow_focus:
             self.set_focus(len(self.view) - 1)
+        self.set_flow_marked(f, False)
         return f
 
     def update_flow(self, f):
@@ -100,9 +101,29 @@ class ConsoleState(flow.State):
         return ret
 
     def clear(self):
-        self.focus = None
+        marked_flows = []
+        for f in self.flows:
+            if self.flow_marked(f):
+                marked_flows.append(f)
+                
         super(ConsoleState, self).clear()
-
+        
+        for f in marked_flows:
+            self.add_flow(f)
+            self.set_flow_marked(f, True)
+        
+        if len(self.flows.views) == 0:
+            self.focus = None
+        else:
+            self.focus = 0
+        self.set_focus(self.focus)
+        
+    def flow_marked(self, flow):
+        return self.get_flow_setting(flow, "marked", False)
+    
+    def set_flow_marked(self, flow, marked):
+        self.add_flow_setting(flow, "marked", marked)
+        
 
 class Options(object):
     attributes = [
@@ -591,6 +612,13 @@ class ConsoleMaster(flow.FlowMaster):
 
     def save_flows(self, path):
         return self._write_flows(path, self.state.view)
+    
+    def save_marked_flows(self, path):
+        marked_flows = []
+        for f in self.state.view:
+            if self.state.flow_marked(f):
+                marked_flows.append(f)
+        return self._write_flows(path, marked_flows)
 
     def load_flows_callback(self, path):
         if not path:
diff --git a/libmproxy/console/common.py b/libmproxy/console/common.py
index e5bebf7f..90bccfe7 100644
--- a/libmproxy/console/common.py
+++ b/libmproxy/console/common.py
@@ -115,9 +115,11 @@ def fcol(s, attr):
 if urwid.util.detected_encoding:
     SYMBOL_REPLAY = u"\u21ba"
     SYMBOL_RETURN = u"\u2190"
+    SYMBOL_MARK = u"\u25cf"
 else:
     SYMBOL_REPLAY = u"[r]"
     SYMBOL_RETURN = u"<-"
+    SYMBOL_MARK = "[m]"
 
 
 def raw_format_flow(f, focus, extended, padding):
@@ -133,6 +135,10 @@ def raw_format_flow(f, focus, extended, padding):
         )
     else:
         req.append(fcol(">>" if focus else "  ", "focus"))
+        
+    if f["marked"]:
+        req.append(fcol(SYMBOL_MARK, "mark"))
+
     if f["req_is_replay"]:
         req.append(fcol(SYMBOL_REPLAY, "replay"))
     req.append(fcol(f["req_method"], "method"))
@@ -372,7 +378,8 @@ def ask_save_body(part, master, state, flow):
 flowcache = utils.LRUCache(800)
 
 
-def format_flow(f, focus, extended=False, hostheader=False, padding=2):
+def format_flow(f, focus, extended=False, hostheader=False, padding=2,
+        marked=False):
     d = dict(
         intercepted = f.intercepted,
         acked = f.reply.acked,
@@ -384,6 +391,8 @@ def format_flow(f, focus, extended=False, hostheader=False, padding=2):
 
         err_msg = f.error.msg if f.error else None,
         resp_code = f.response.code if f.response else None,
+        
+        marked = marked,
     )
     if f.response:
         if f.response.content:
diff --git a/libmproxy/console/flowlist.py b/libmproxy/console/flowlist.py
index 39245984..bb23df75 100644
--- a/libmproxy/console/flowlist.py
+++ b/libmproxy/console/flowlist.py
@@ -17,9 +17,11 @@ def _mkhelp():
         ("F", "toggle follow flow list"),
         ("l", "set limit filter pattern"),
         ("L", "load saved flows"),
+        ("m", "toggle flow mark"),
         ("n", "create a new request"),
         ("P", "copy flow to clipboard"),
         ("r", "replay request"),
+        ("U", "unmark all marked flows"),
         ("V", "revert changes to request"),
         ("w", "save flows "),
         ("W", "stream flows to file"),
@@ -108,7 +110,8 @@ class ConnectionItem(urwid.WidgetWrap):
         return common.format_flow(
             self.flow,
             self.f,
-            hostheader = self.master.showhost
+            hostheader = self.master.showhost,
+            marked=self.state.flow_marked(self.flow)
         )
 
     def selectable(self):
@@ -120,6 +123,11 @@ class ConnectionItem(urwid.WidgetWrap):
                 prompt = "Save all flows to",
                 callback = self.master.save_flows
             )
+        elif k == "m":
+            signals.status_prompt_path.send(
+                prompt = "Save marked flows to",
+                callback = self.master.save_marked_flows
+            )
         else:
             signals.status_prompt_path.send(
                 prompt = "Save this flow to",
@@ -177,6 +185,12 @@ class ConnectionItem(urwid.WidgetWrap):
         elif key == "D":
             f = self.master.duplicate_flow(self.flow)
             self.master.view_flow(f)
+        elif key == "m":
+            if self.state.flow_marked(self.flow):
+                self.state.set_flow_marked(self.flow, False)
+            else:
+                self.state.set_flow_marked(self.flow, True)
+            signals.flowlist_change.send(self)
         elif key == "r":
             r = self.master.replay_request(self.flow)
             if r:
@@ -202,6 +216,10 @@ class ConnectionItem(urwid.WidgetWrap):
                     ),
                     callback = self.stop_server_playback_prompt,
                 )
+        elif key == "U":
+            for f in self.state.flows:
+                self.state.set_flow_marked(f, False)
+            signals.flowlist_change.send(self)       
         elif key == "V":
             if not self.flow.modified():
                 signals.status_message.send(message="Flow not modified.")
@@ -216,6 +234,7 @@ class ConnectionItem(urwid.WidgetWrap):
                 keys = (
                     ("all flows", "a"),
                     ("this flow", "t"),
+                    ("marked flows", "m"),
                 ),
                 callback = self.save_flows_prompt,
             )
diff --git a/libmproxy/console/palettes.py b/libmproxy/console/palettes.py
index ea3d1b62..d897a0a2 100644
--- a/libmproxy/console/palettes.py
+++ b/libmproxy/console/palettes.py
@@ -24,7 +24,7 @@ class Palette:
         'method', 'focus',
         'code_200', 'code_300', 'code_400', 'code_500', 'code_other',
         'error',
-        'header', 'highlight', 'intercept', 'replay',
+        'header', 'highlight', 'intercept', 'replay', 'mark',
 
         # Hex view
         'offset',
@@ -104,6 +104,7 @@ class LowDark(Palette):
         highlight = ('white,bold', 'default'),
         intercept = ('brown', 'default'),
         replay = ('light green', 'default'),
+        mark = ('light red', 'default'),
 
         # Hex view
         offset = ('dark cyan', 'default'),
@@ -167,6 +168,7 @@ class LowLight(Palette):
         highlight = ('black,bold', 'default'),
         intercept = ('brown', 'default'),
         replay = ('dark green', 'default'),
+        mark = ('dark red', 'default'),
 
         # Hex view
         offset = ('dark blue', 'default'),