Some refactoring of the console inteface.

Reduce some state duplication, by removing currentflow variable.

Fixes #141

5 files changed, 30 insertions, 34 deletions

diff --git a/libmproxy/console/__init__.py b/libmproxy/console/__init__.py
index dd3adaa9..cf4fc988 100644
--- a/libmproxy/console/__init__.py
+++ b/libmproxy/console/__init__.py
@@ -204,15 +204,10 @@ class StatusBar(common.WWrap):
             self.message("")
 
         fc = self.master.state.flow_count()
-        if self.master.currentflow:
-            idx = self.master.state.view.index(self.master.currentflow) + 1
-            t = [
-                ('heading', ("[%s/%s]"%(idx, fc)).ljust(9))
-            ]
-        else:
-            t = [
-                ('heading', ("[%s]"%fc).ljust(9))
-            ]
+        offset = min(self.master.state.focus + 1, fc)
+        t = [
+            ('heading', ("[%s/%s]"%(offset, fc)).ljust(9))
+        ]
 
         if self.master.server.bound:
             boundaddr = "[%s:%s]"%(self.master.server.address or "*", self.master.server.port)
@@ -263,7 +258,10 @@ class ConsoleState(flow.State):
         self.focus = None
         self.follow_focus = None
         self.default_body_view = contentview.get("Auto")
+
+        self.view_mode = common.VIEW_LIST
         self.view_flow_mode = common.VIEW_FLOW_REQUEST
+
         self.last_script = ""
         self.last_saveload = ""
         self.flowsettings = weakref.WeakKeyDictionary()
@@ -308,6 +306,9 @@ class ConsoleState(flow.State):
                 idx = 0
             self.focus = idx
 
+    def set_focus_flow(self, f):
+        self.set_focus(self.view.index(f))
+
     def get_from_pos(self, pos):
         if len(self.view) <= pos or pos < 0:
             return None, None
@@ -320,6 +321,8 @@ class ConsoleState(flow.State):
         return self.get_from_pos(pos-1)
 
     def delete_flow(self, f):
+        if f in self.view and self.view.index(f) <= self.focus:
+            self.focus -= 1
         ret = flow.State.delete_flow(self, f)
         self.set_focus(self.focus)
         return ret
@@ -570,8 +573,6 @@ class ConsoleMaster(flow.FlowMaster):
         self.palette = palettes.palettes[name]
 
     def run(self):
-        self.currentflow = None
-
         self.ui = urwid.raw_display.Screen()
         self.ui.set_terminal_properties(256)
         self.ui.register_palette(self.palette)
@@ -605,13 +606,6 @@ class ConsoleMaster(flow.FlowMaster):
         sys.stderr.flush()
         self.shutdown()
 
-    def focus_current(self):
-        if self.currentflow:
-            try:
-                self.flow_list_walker.set_focus(self.state.view.index(self.currentflow))
-            except (IndexError, ValueError):
-                pass
-
     def make_view(self):
         self.view = urwid.Frame(
                         self.body,
@@ -646,8 +640,6 @@ class ConsoleMaster(flow.FlowMaster):
             self.ui.clear()
         if self.state.follow_focus:
             self.state.set_focus(self.state.flow_count())
-        else:
-            self.focus_current()
 
         if self.eventlog:
             self.body = flowlist.BodyPile(self)
@@ -655,7 +647,7 @@ class ConsoleMaster(flow.FlowMaster):
             self.body = flowlist.FlowListBox(self)
         self.statusbar = StatusBar(self, flowlist.footer)
         self.header = None
-        self.currentflow = None
+        self.state.view_mode = common.VIEW_LIST
 
         self.make_view()
         self.help_context = flowlist.help_context
@@ -664,7 +656,8 @@ class ConsoleMaster(flow.FlowMaster):
         self.body = flowview.FlowView(self, self.state, flow)
         self.header = flowview.FlowViewHeader(self, flow)
         self.statusbar = StatusBar(self, flowview.footer)
-        self.currentflow = flow
+        self.state.set_focus_flow(flow)
+        self.state.view_mode = common.VIEW_FLOW
 
         self.make_view()
         self.help_context = flowview.help_context
@@ -711,7 +704,6 @@ class ConsoleMaster(flow.FlowMaster):
         f.close()
         if self.flow_list_walker:
             self.sync_list_view()
-            self.focus_current()
         return reterr
 
     def path_prompt(self, prompt, text, callback, *args):
@@ -777,8 +769,7 @@ class ConsoleMaster(flow.FlowMaster):
     def change_default_display_mode(self, t):
         v = contentview.get_by_shortcut(t)
         self.state.default_body_view = v
-        if self.currentflow:
-            self.refresh_flow(self.currentflow)
+        self.refresh_focus()
 
     def set_reverse_proxy(self, txt):
         if not txt:
@@ -796,8 +787,8 @@ class ConsoleMaster(flow.FlowMaster):
         return size
 
     def pop_view(self):
-        if self.currentflow:
-            self.view_flow(self.currentflow)
+        if self.state.view_mode == common.VIEW_FLOW:
+            self.view_flow(self.state.view[self.state.focus])
         else:
             self.view_flowlist()
 
@@ -972,7 +963,7 @@ class ConsoleMaster(flow.FlowMaster):
         if a == "h":
             self.showhost = not self.showhost
             self.sync_list_view()
-            self.refresh_flow(self.currentflow)
+            self.refresh_focus()
         elif a == "k":
             self.killextra = not self.killextra
         elif a == "n":
@@ -1003,6 +994,10 @@ class ConsoleMaster(flow.FlowMaster):
         self.state.delete_flow(f)
         self.sync_list_view()
 
+    def refresh_focus(self):
+        if self.state.view:
+            self.refresh_flow(self.state.view[self.state.focus])
+
     def refresh_flow(self, c):
         if hasattr(self.header, "refresh_flow"):
             self.header.refresh_flow(c)
diff --git a/libmproxy/console/common.py b/libmproxy/console/common.py
index d68aba3d..4f3ab090 100644
--- a/libmproxy/console/common.py
+++ b/libmproxy/console/common.py
@@ -18,6 +18,9 @@ import urwid.util
 from .. import utils, flow
 
 
+VIEW_LIST = 0
+VIEW_FLOW = 1
+
 
 VIEW_FLOW_REQUEST = 0
 VIEW_FLOW_RESPONSE = 1
diff --git a/libmproxy/console/flowlist.py b/libmproxy/console/flowlist.py
index 8fd4efce..124e252e 100644
--- a/libmproxy/console/flowlist.py
+++ b/libmproxy/console/flowlist.py
@@ -160,8 +160,7 @@ class ConnectionItem(common.WWrap):
             self.master.sync_list_view()
         elif key == "D":
             f = self.master.duplicate_flow(self.flow)
-            self.master.currentflow = f
-            self.master.focus_current()
+            self.master.view_flow(f)
         elif key == "r":
             self.flow.backup()
             r = self.master.replay_request(self.flow)
diff --git a/libmproxy/console/flowview.py b/libmproxy/console/flowview.py
index da5a6c65..419bfbcd 100644
--- a/libmproxy/console/flowview.py
+++ b/libmproxy/console/flowview.py
@@ -342,7 +342,7 @@ class FlowView(common.WWrap):
         else:
             if not self.flow.response:
                 self.flow.response = flow.Response(
-                    self.flow.request, 
+                    self.flow.request,
                     self.flow.request.httpversion,
                     200, "OK", flow.ODictCaseless(), "", None
                 )
@@ -393,7 +393,7 @@ class FlowView(common.WWrap):
             new_flow, new_idx = self.state.get_next(idx)
         else:
             new_flow, new_idx = self.state.get_prev(idx)
-        if new_idx is None:
+        if new_flow is None:
             self.master.statusbar.message("No more flows!")
             return
         self.master.view_flow(new_flow)
@@ -478,7 +478,6 @@ class FlowView(common.WWrap):
         elif key == "D":
             f = self.master.duplicate_flow(self.flow)
             self.master.view_flow(f)
-            self.master.currentflow = f
             self.master.statusbar.message("Duplicated.")
         elif key == "e":
             if self.state.view_flow_mode == common.VIEW_FLOW_REQUEST:
diff --git a/test/test_flow.py b/test/test_flow.py
index 977cdd4e..32cfb0ca 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -591,7 +591,7 @@ class TestFlowMaster:
         s = flow.State()
         fm = flow.FlowMaster(None, s)
         f = tutils.tflow_full()
-        fm.load_flow(f)
+        f = fm.load_flow(f)
         assert s.flow_count() == 1
         f2 = fm.duplicate_flow(f)
         assert f2.response