Merge pull request #223 from droope/improve-backwards-search

Improve search to also search backwards.

3 files changed, 281 insertions, 141 deletions

diff --git a/libmproxy/console/flowview.py b/libmproxy/console/flowview.py
index 6c4a2651..f95f2ded 100644
--- a/libmproxy/console/flowview.py
+++ b/libmproxy/console/flowview.py
@@ -68,7 +68,8 @@ def _mkhelp():
         ("space", "next flow"),
         ("|", "run script on this flow"),
         ("/", "search in response body (case sensitive)"),
-        ("n", "repeat previous search"),
+        ("n", "repeat search forward"),
+        ("N", "repeat search backwards"),
     ]
     text.extend(common.format_keyvals(keys, key="key", val="text", indent=4))
     return text
@@ -255,7 +256,7 @@ class FlowView(common.WWrap):
                 )
         return f
 
-    def search_wrapped_around(self, last_find_line, last_search_index):
+    def search_wrapped_around(self, last_find_line, last_search_index, backwards):
         """
             returns true if search wrapped around the bottom.
         """
@@ -265,15 +266,39 @@ class FlowView(common.WWrap):
         current_search_index = self.state.get_flow_setting(self.flow,
                 "last_search_index")
 
-        if current_find_line <= last_find_line:
-            return True
-        elif current_find_line == last_find_line:
-            if current_search_index <= last_search_index:
-                return True
+        if not backwards:
+            message = "search hit BOTTOM, continuing at TOP"
+            if current_find_line <= last_find_line:
+                return True, message
+            elif current_find_line == last_find_line:
+                if current_search_index <= last_search_index:
+                    return True, message
+        else:
+            message = "search hit TOP, continuing at BOTTOM"
+            if current_find_line >= last_find_line:
+                return True, message
+            elif current_find_line == last_find_line:
+                if current_search_index >= last_search_index:
+                    return True, message
 
-        return False
+        return False, ""
 
-    def search(self, search_string):
+    def search_again(self, backwards=False):
+        """
+            runs the previous search again, forwards or backwards.
+        """
+        last_search_string = self.state.get_flow_setting(self.flow, "last_search_string")
+        if last_search_string:
+            message = self.search(last_search_string, backwards)
+            if message:
+                self.master.statusbar.message(message)
+        else:
+            message = "no previous searches have been made"
+            self.master.statusbar.message(message)
+
+        return message
+
+    def search(self, search_string, backwards=False):
         """
             similar to view_response or view_request, but instead of just
             displaying the conn, it highlights a word that the user is
@@ -301,7 +326,7 @@ class FlowView(common.WWrap):
         # generate the body, highlight the words and get focus
         headers, msg, body = self.conn_text_raw(text)
         try:
-            body, focus_position = self.search_highlight_text(body, search_string)
+            body, focus_position = self.search_highlight_text(body, search_string, backwards=backwards)
         except SearchError:
             return "Search not supported in this view."
 
@@ -318,8 +343,11 @@ class FlowView(common.WWrap):
 
         self.last_displayed_body = list_box
 
-        if self.search_wrapped_around(last_find_line, last_search_index):
-            return "search hit BOTTOM, continuing at TOP"
+        wrapped, wrapped_message = self.search_wrapped_around(last_find_line, last_search_index, backwards)
+
+        if wrapped:
+            print(wrapped, wrapped_message)
+            return wrapped_message
 
     def search_get_start(self, search_string):
         start_line = 0
@@ -344,57 +372,94 @@ class FlowView(common.WWrap):
 
         return (start_line, start_index)
 
-    def search_highlight_text(self, text_objects, search_string, looping = False):
+    def search_get_range(self, len_text_objects, start_line, backwards):
+        if not backwards:
+            loop_range = xrange(start_line, len_text_objects)
+        else:
+            loop_range = xrange(start_line, -1, -1)
+
+        return loop_range
+
+    def search_find(self, text, search_string, start_index, backwards):
+            if backwards == False:
+                find_index = text.find(search_string, start_index)
+            else:
+                if start_index != 0:
+                    start_index -= len(search_string)
+                else:
+                    start_index = None
+
+                find_index = text.rfind(search_string, 0, start_index)
+
+            return find_index
+
+    def search_highlight_text(self, text_objects, search_string, looping = False, backwards = False):
         start_line, start_index = self.search_get_start(search_string)
         i = start_line
 
         found = False
         text_objects = copy.deepcopy(text_objects)
-        for text_object in text_objects[start_line:]:
-            if i != start_line:
-                start_index = 0
+        loop_range = self.search_get_range(len(text_objects), start_line, backwards)
+        for i in loop_range:
+            text_object = text_objects[i]
 
             try:
                 text, style = text_object.get_text()
             except AttributeError:
                 raise SearchError()
-            find_index = text.find(search_string, start_index)
-            if find_index != -1:
-                before = text[:find_index]
-                after = text[find_index+len(search_string):]
-                new_text = urwid.Text(
-                    [
-                        before,
-                        (self.highlight_color, search_string),
-                        after,
-                    ]
-                )
 
-                self.state.add_flow_setting(self.flow, "last_search_index",
-                        find_index)
-                self.state.add_flow_setting(self.flow, "last_find_line", i)
+            if i != start_line:
+                start_index = 0
 
+            find_index = self.search_find(text, search_string, start_index, backwards)
+
+            if find_index != -1:
+                new_text = self.search_highlight_object(text, find_index, search_string)
                 text_objects[i] = new_text
 
                 found = True
+                self.state.add_flow_setting(self.flow, "last_search_index",
+                        find_index)
+                self.state.add_flow_setting(self.flow, "last_find_line", i)
 
                 break
 
-            i += 1
-
+        # handle search WRAP
         if found:
             focus_pos = i
         else :
-            # loop from the beginning, but not forever.
-            if (start_line == 0 and start_index == 0) or looping:
+            if looping:
                 focus_pos = None
             else:
-                self.state.add_flow_setting(self.flow, "last_search_index", 0)
-                self.state.add_flow_setting(self.flow, "last_find_line", 0)
-                text_objects, focus_pos = self.search_highlight_text(text_objects, search_string, True)
+                if not backwards:
+                    self.state.add_flow_setting(self.flow, "last_search_index", 0)
+                    self.state.add_flow_setting(self.flow, "last_find_line", 0)
+                else:
+                    self.state.add_flow_setting(self.flow, "last_search_index", None)
+                    self.state.add_flow_setting(self.flow, "last_find_line", len(text_objects) - 1)
+
+                text_objects, focus_pos = self.search_highlight_text(text_objects,
+                        search_string, looping=True, backwards=backwards)
 
         return text_objects, focus_pos
 
+    def search_highlight_object(self, text_object, find_index, search_string):
+        """
+            just a little abstraction
+        """
+        before = text_object[:find_index]
+        after = text_object[find_index+len(search_string):]
+
+        new_text = urwid.Text(
+            [
+                before,
+                (self.highlight_color, search_string),
+                after,
+            ]
+        )
+
+        return new_text
+
     def view_request(self):
         self.state.view_flow_mode = common.VIEW_FLOW_REQUEST
         body = self.conn_text(self.flow.request)
@@ -761,13 +826,9 @@ class FlowView(common.WWrap):
                     None,
                     self.search)
         elif key == "n":
-            last_search_string = self.state.get_flow_setting(self.flow, "last_search_string")
-            if last_search_string:
-                message = self.search(last_search_string)
-                if message:
-                    self.master.statusbar.message(message)
-            else:
-                self.master.statusbar.message("no previous searches have been made")
+            self.search_again(backwards=False)
+        elif key == "N":
+            self.search_again(backwards=True)
         else:
             return key
 
diff --git a/test/test_console_contentview.py b/test/test_console_contentview.py
index 07ecf1d0..a878ad4e 100644
--- a/test/test_console_contentview.py
+++ b/test/test_console_contentview.py
@@ -276,100 +276,3 @@ if cv.ViewProtobuf.is_available():
 def test_get_by_shortcut():
     assert cv.get_by_shortcut("h")
 
-def test_search_highlights():
-    # Default text in requests is content. We will search for nt once, and
-    # expect the first bit to be highlighted. We will do it again and expect the
-    # second to be.
-    f = tutils.tflowview()
-
-    f.search("nt")
-    text_object = tutils.get_body_line(f.last_displayed_body, 0)
-    assert text_object.get_text() == ('content', [(None, 2), (f.highlight_color, 2)])
-
-    f.search("nt")
-    text_object = tutils.get_body_line(f.last_displayed_body, 1)
-    assert text_object.get_text() == ('content', [(None, 5), (f.highlight_color, 2)])
-
-def test_search_returns_useful_messages():
-    f = tutils.tflowview()
-
-    # original string is content. this string should not be in there.
-    response = f.search("oranges and other fruit.")
-    assert response == "no matches for 'oranges and other fruit.'"
-
-def test_search_highlights_clears_prev():
-    f = tutils.tflowview(request_contents="this is string\nstring is string")
-
-    f.search("string")
-    text_object = tutils.get_body_line(f.last_displayed_body, 0)
-    assert text_object.get_text() == ('this is string', [(None, 8), (f.highlight_color, 6)])
-
-    # search again, it should not be highlighted again.
-    f.search("string")
-    text_object = tutils.get_body_line(f.last_displayed_body, 0)
-    assert text_object.get_text() != ('this is string', [(None, 8), (f.highlight_color, 6)])
-
-def test_search_highlights_multi_line():
-    f = tutils.tflowview(request_contents="this is string\nstring is string")
-
-    # should highlight the first line.
-    f.search("string")
-    text_object = tutils.get_body_line(f.last_displayed_body, 0)
-    assert text_object.get_text() == ('this is string', [(None, 8), (f.highlight_color, 6)])
-
-    # should highlight second line, first appearance of string.
-    f.search("string")
-    text_object = tutils.get_body_line(f.last_displayed_body, 1)
-    assert text_object.get_text() == ('string is string', [(None, 0), (f.highlight_color, 6)])
-
-    # should highlight third line, second appearance of string.
-    f.search("string")
-    text_object = tutils.get_body_line(f.last_displayed_body, 1)
-    assert text_object.get_text() == ('string is string', [(None, 10), (f.highlight_color, 6)])
-
-def test_search_loops():
-    f = tutils.tflowview(request_contents="this is string\nstring is string")
-
-    # get to the end.
-    f.search("string")
-    f.search("string")
-    f.search("string")
-
-    # should highlight the first line.
-    message = f.search("string")
-    text_object = tutils.get_body_line(f.last_displayed_body, 0)
-    assert text_object.get_text() == ('this is string', [(None, 8), (f.highlight_color, 6)])
-    assert message == "search hit BOTTOM, continuing at TOP"
-
-def test_search_focuses():
-    f = tutils.tflowview(request_contents="this is string\nstring is string")
-
-    # should highlight the first line.
-    f.search("string")
-
-    # should be focusing on the 2nd text line.
-    f.search("string")
-    text_object = tutils.get_body_line(f.last_displayed_body, 1)
-    assert f.last_displayed_body.focus == text_object
-
-def test_search_does_not_crash_on_bad():
-    """
-        this used to crash, kept for reference.
-    """
-
-    f = tutils.tflowview(request_contents="this is string\nstring is string\n"+("A" * cv.VIEW_CUTOFF)+"AFTERCUTOFF")
-    f.search("AFTERCUTOFF")
-
-    # pretend F
-    f.state.add_flow_setting(
-        f.flow,
-        (f.state.view_flow_mode, "fullcontents"),
-        True
-    )
-    f.master.refresh_flow(f.flow)
-
-    # text changed, now this string will exist. can happen when user presses F
-    # for full text view
-    f.search("AFTERCUTOFF")
-
-
diff --git a/test/test_console_search.py b/test/test_console_search.py
new file mode 100644
index 00000000..f4e9ff9b
--- /dev/null
+++ b/test/test_console_search.py
@@ -0,0 +1,176 @@
+
+import sys
+import libmproxy.console.contentview as cv
+from libmproxy import utils, flow, encoding
+import tutils
+
+def test_search_highlights():
+    # Default text in requests is content. We will search for nt once, and
+    # expect the first bit to be highlighted. We will do it again and expect the
+    # second to be.
+    f = tutils.tflowview()
+
+    f.search("nt")
+    text_object = tutils.get_body_line(f.last_displayed_body, 0)
+    assert text_object.get_text() == ('content', [(None, 2), (f.highlight_color, 2)])
+
+    f.search("nt")
+    text_object = tutils.get_body_line(f.last_displayed_body, 1)
+    assert text_object.get_text() == ('content', [(None, 5), (f.highlight_color, 2)])
+
+def test_search_returns_useful_messages():
+    f = tutils.tflowview()
+
+    # original string is content. this string should not be in there.
+    test_string = "oranges and other fruit."
+    response = f.search(test_string)
+    assert response == "no matches for '%s'" % test_string
+
+def test_search_highlights_clears_prev():
+    f = tutils.tflowview(request_contents="this is string\nstring is string")
+
+    f.search("string")
+    text_object = tutils.get_body_line(f.last_displayed_body, 0)
+    assert text_object.get_text() == ('this is string', [(None, 8), (f.highlight_color, 6)])
+
+    # search again, it should not be highlighted again.
+    f.search("string")
+    text_object = tutils.get_body_line(f.last_displayed_body, 0)
+    assert text_object.get_text() != ('this is string', [(None, 8), (f.highlight_color, 6)])
+
+def test_search_highlights_multi_line(flow=None):
+    f = flow if flow else tutils.tflowview(request_contents="this is string\nstring is string")
+
+    # should highlight the first line.
+    f.search("string")
+    text_object = tutils.get_body_line(f.last_displayed_body, 0)
+    assert text_object.get_text() == ('this is string', [(None, 8), (f.highlight_color, 6)])
+
+    # should highlight second line, first appearance of string.
+    f.search("string")
+    text_object = tutils.get_body_line(f.last_displayed_body, 1)
+    assert text_object.get_text() == ('string is string', [(None, 0), (f.highlight_color, 6)])
+
+    # should highlight third line, second appearance of string.
+    f.search("string")
+    text_object = tutils.get_body_line(f.last_displayed_body, 1)
+    assert text_object.get_text() == ('string is string', [(None, 10), (f.highlight_color, 6)])
+
+def test_search_loops():
+    f = tutils.tflowview(request_contents="this is string\nstring is string")
+
+    # get to the end.
+    f.search("string")
+    f.search("string")
+    f.search("string")
+
+    # should highlight the first line.
+    message = f.search("string")
+    text_object = tutils.get_body_line(f.last_displayed_body, 0)
+    assert text_object.get_text() == ('this is string', [(None, 8), (f.highlight_color, 6)])
+    assert message == "search hit BOTTOM, continuing at TOP"
+
+def test_search_focuses():
+    f = tutils.tflowview(request_contents="this is string\nstring is string")
+
+    # should highlight the first line.
+    f.search("string")
+
+    # should be focusing on the 2nd text line.
+    f.search("string")
+    text_object = tutils.get_body_line(f.last_displayed_body, 1)
+    assert f.last_displayed_body.focus == text_object
+
+def test_search_does_not_crash_on_bad():
+    """
+        this used to crash, kept for reference.
+    """
+
+    f = tutils.tflowview(request_contents="this is string\nstring is string\n"+("A" * cv.VIEW_CUTOFF)+"AFTERCUTOFF")
+    f.search("AFTERCUTOFF")
+
+    # pretend F
+    f.state.add_flow_setting(
+        f.flow,
+        (f.state.view_flow_mode, "fullcontents"),
+        True
+    )
+    f.master.refresh_flow(f.flow)
+
+    # text changed, now this string will exist. can happen when user presses F
+    # for full text view
+    f.search("AFTERCUTOFF")
+
+def test_search_backwards():
+    f = tutils.tflowview(request_contents="content, content")
+
+    first_match = ('content, content', [(None, 2), (f.highlight_color, 2)])
+
+    f.search("nt")
+    text_object = tutils.get_body_line(f.last_displayed_body, 0)
+    assert text_object.get_text() == first_match
+
+    f.search("nt")
+    text_object = tutils.get_body_line(f.last_displayed_body, 1)
+    assert text_object.get_text() == ('content, content', [(None, 5), (f.highlight_color, 2)])
+
+    f.search_again(backwards=True)
+    text_object = tutils.get_body_line(f.last_displayed_body, 0)
+    assert text_object.get_text() == first_match
+
+def test_search_back_multiline():
+    f = tutils.tflowview(request_contents="this is string\nstring is string")
+
+    # shared assertions. highlight and pointers should now be on the third
+    # 'string' appearance
+    test_search_highlights_multi_line(f)
+
+    # should highlight second line, first appearance of string.
+    f.search_again(backwards=True)
+    text_object = tutils.get_body_line(f.last_displayed_body, 1)
+    assert text_object.get_text() == ('string is string', [(None, 0), (f.highlight_color, 6)])
+
+    # should highlight the first line again.
+    f.search_again(backwards=True)
+    text_object = tutils.get_body_line(f.last_displayed_body, 0)
+    assert text_object.get_text() == ('this is string', [(None, 8), (f.highlight_color, 6)])
+
+def test_search_back_multi_multi_line():
+    """
+        same as above for some bugs the above won't catch.
+    """
+    f = tutils.tflowview(request_contents="this is string\nthis is string\nthis is string")
+
+    f.search("string")
+    f.search_again()
+    f.search_again()
+
+    # should be on second line
+    f.search_again(backwards=True)
+    text_object = tutils.get_body_line(f.last_displayed_body, 1)
+    assert text_object.get_text() == ('this is string', [(None, 8), (f.highlight_color, 6)])
+
+    # first line now
+    f.search_again(backwards=True)
+    text_object = tutils.get_body_line(f.last_displayed_body, 0)
+    assert text_object.get_text() == ('this is string', [(None, 8), (f.highlight_color, 6)])
+
+def test_search_backwards_wraps():
+    """
+        when searching past line 0, it should loop.
+    """
+    f = tutils.tflowview(request_contents="this is string\nthis is string\nthis is string")
+
+    # should be on second line
+    f.search("string")
+    f.search_again()
+    text_object = tutils.get_body_line(f.last_displayed_body, 1)
+    assert text_object.get_text() == ('this is string', [(None, 8), (f.highlight_color, 6)])
+
+    # should be on third now.
+    f.search_again(backwards=True)
+    message = f.search_again(backwards=True)
+    text_object = tutils.get_body_line(f.last_displayed_body, 2)
+    assert text_object.get_text() == ('this is string', [(None, 8), (f.highlight_color, 6)])
+    assert message == "search hit TOP, continuing at BOTTOM"
+