From 2a6698bf5a2ebe576ae0bbcacdee69d6eed10be9 Mon Sep 17 00:00:00 2001
From: Jake Drahos <drahos@iastate.edu>
Date: Thu, 11 Jun 2015 20:27:33 -0500
Subject: Moved marking from flow to console

No longer taints the flow primitive
---
 libmproxy/console/__init__.py    | 20 +++++++++++++++++++-
 libmproxy/console/common.py      |  5 +++--
 libmproxy/console/flowlist.py    |  8 ++++++--
 libmproxy/flow.py                | 15 +--------------
 libmproxy/protocol/primitives.py |  7 -------
 5 files changed, 29 insertions(+), 26 deletions(-)

diff --git a/libmproxy/console/__init__.py b/libmproxy/console/__init__.py
index 20580a28..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,12 +101,29 @@ class ConsoleState(flow.State):
         return ret
 
     def clear(self):
+        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 = [
@@ -598,7 +616,7 @@ class ConsoleMaster(flow.FlowMaster):
     def save_marked_flows(self, path):
         marked_flows = []
         for f in self.state.view:
-            if f.marked:
+            if self.state.flow_marked(f):
                 marked_flows.append(f)
         return self._write_flows(path, marked_flows)
 
diff --git a/libmproxy/console/common.py b/libmproxy/console/common.py
index cbf39e6e..90bccfe7 100644
--- a/libmproxy/console/common.py
+++ b/libmproxy/console/common.py
@@ -378,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,
@@ -391,7 +392,7 @@ 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 = f.marked,
+        marked = marked,
     )
     if f.response:
         if f.response.content:
diff --git a/libmproxy/console/flowlist.py b/libmproxy/console/flowlist.py
index 72d507c8..87e7c77a 100644
--- a/libmproxy/console/flowlist.py
+++ b/libmproxy/console/flowlist.py
@@ -109,7 +109,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):
@@ -184,7 +185,10 @@ class ConnectionItem(urwid.WidgetWrap):
             f = self.master.duplicate_flow(self.flow)
             self.master.view_flow(f)
         elif key == "m":
-            self.flow.toggle_mark()
+            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)
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 0dffc8d3..1a052f51 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -494,19 +494,8 @@ class FlowStore(FlowList):
         self._recalculate_views()
 
     def _clear(self):
-        marked_flows = []
-        for f in self._list:
-            if f.marked:
-                marked_flows.append(f)
-        if len(marked_flows) == len(self._list):
-            marked_flows = []
-                
         self._list = []
         self._set = set()
-        
-        for f in marked_flows:
-            self._add(f)
-            
         self._recalculate_views()
 
     def _recalculate_views(self):
@@ -815,9 +804,7 @@ class FlowMaster(controller.Master):
         return super(FlowMaster, self).tick(q, timeout)
 
     def duplicate_flow(self, f):
-        new_flow = f.copy()
-        new_flow.marked = False
-        return self.load_flow(new_flow)
+        return self.load_flow(f.copy())
 
     def create_request(self, method, scheme, host, port, path):
         """
diff --git a/libmproxy/protocol/primitives.py b/libmproxy/protocol/primitives.py
index 1d2389c3..2f8ea3e0 100644
--- a/libmproxy/protocol/primitives.py
+++ b/libmproxy/protocol/primitives.py
@@ -77,7 +77,6 @@ class Flow(stateobject.StateObject):
         """@type: bool"""
         self._backup = None
         self.reply = None
-        self.marked = False
 
     _stateobject_attributes = dict(
         id=str,
@@ -166,12 +165,6 @@ class Flow(stateobject.StateObject):
         self.intercepted = False
         self.reply()
         master.handle_accept_intercept(self)
-        
-    def toggle_mark(self):
-        if self.marked:
-            self.marked = False
-        else:
-            self.marked = True
 
 
 class ProtocolHandler(object):
-- 
cgit v1.2.3