console: first pass of a Set-cookie editor for responses

1 files changed, 149 insertions, 50 deletions

diff --git a/libmproxy/console/grideditor.py b/libmproxy/console/grideditor.py
index 493268ba..6d112e42 100644
--- a/libmproxy/console/grideditor.py
+++ b/libmproxy/console/grideditor.py
@@ -7,7 +7,7 @@ import urwid
 
 from . import common, signals
 from .. import utils, filt, script
-from netlib import http_uastrings
+from netlib import http_uastrings, http_cookies
 
 
 FOOTER = [
@@ -19,17 +19,36 @@ FOOTER_EDITING = [
 ]
 
 
-class SText(urwid.WidgetWrap):
-    def __init__(self, txt, focused, error):
+class TextColumn:
+    subeditor = None
+
+    def __init__(self, heading):
+        self.heading = heading
+
+    def text(self, obj):
+        return SEscaped(obj or "")
+
+    def blank(self):
+        return ""
+
+
+class SubgridColumn:
+    def __init__(self, heading, subeditor):
+        self.heading = heading
+        self.subeditor = subeditor
+
+    def text(self, obj):
+        p = http_cookies._format_pairs(obj, sep="\n")
+        return urwid.Text(p)
+
+    def blank(self):
+        return []
+
+
+class SEscaped(urwid.WidgetWrap):
+    def __init__(self, txt):
         txt = txt.encode("string-escape")
         w = urwid.Text(txt, wrap="any")
-        if focused:
-            if error:
-                w = urwid.AttrWrap(w, "focusfield_error")
-            else:
-                w = urwid.AttrWrap(w, "focusfield")
-        elif error:
-            w = urwid.AttrWrap(w, "field_error")
         urwid.WidgetWrap.__init__(self, w)
 
     def get_text(self):
@@ -50,7 +69,7 @@ class SEdit(urwid.WidgetWrap):
         urwid.WidgetWrap.__init__(self, w)
 
     def get_text(self):
-        return self._w.get_text()[0]
+        return self._w.get_text()[0].strip()
 
     def selectable(self):
         return True
@@ -67,9 +86,15 @@ class GridRow(urwid.WidgetWrap):
                 self.editing = SEdit(v)
                 self.fields.append(self.editing)
             else:
-                self.fields.append(
-                    SText(v, True if focused == i else False, i in errors)
-                )
+                w = self.editor.columns[i].text(v)
+                if focused == i:
+                    if i in errors:
+                        w = urwid.AttrWrap(w, "focusfield_error")
+                    else:
+                        w = urwid.AttrWrap(w, "focusfield")
+                elif i in errors:
+                    w = urwid.AttrWrap(w, "field_error")
+                self.fields.append(w)
 
         fspecs = self.fields[:]
         if len(self.fields) > 1:
@@ -126,7 +151,9 @@ class GridWalker(urwid.ListWalker):
                 val = val.decode("string-escape")
             except ValueError:
                 signals.status_message.send(
-                    self, message = "Invalid Python-style string encoding.", expure = 1000
+                    self,
+                    message = "Invalid Python-style string encoding.",
+                    expire = 1000
                 )
                 return
         errors = self.lst[self.focus][1]
@@ -136,10 +163,15 @@ class GridWalker(urwid.ListWalker):
             errors.add(self.focus_col)
         else:
             errors.discard(self.focus_col)
-
-        row = list(self.lst[self.focus][0])
-        row[self.focus_col] = val
-        self.lst[self.focus] = [tuple(row), errors]
+        self.set_value(val, self.focus, self.focus_col, errors)
+
+    def set_value(self, val, focus, focus_col, errors=None):
+        if not errors:
+            errors = set([])
+        row = list(self.lst[focus][0])
+        row[focus_col] = val
+        self.lst[focus] = [tuple(row), errors]
+        self._modified()
 
     def delete_focus(self):
         if self.lst:
@@ -149,7 +181,10 @@ class GridWalker(urwid.ListWalker):
 
     def _insert(self, pos):
         self.focus = pos
-        self.lst.insert(self.focus, [[""]*self.editor.columns, set([])])
+        self.lst.insert(
+            self.focus, [
+            [c.blank() for c in self.editor.columns], set([])]
+        )
         self.focus_col = 0
         self.start_edit()
 
@@ -179,12 +214,12 @@ class GridWalker(urwid.ListWalker):
         self._modified()
 
     def right(self):
-        self.focus_col = min(self.focus_col + 1, self.editor.columns-1)
+        self.focus_col = min(self.focus_col + 1, len(self.editor.columns)-1)
         self._modified()
 
     def tab_next(self):
         self.stop_edit()
-        if self.focus_col < self.editor.columns-1:
+        if self.focus_col < len(self.editor.columns)-1:
             self.focus_col += 1
         elif self.focus != len(self.lst)-1:
             self.focus_col = 0
@@ -231,7 +266,6 @@ FIRST_WIDTH_MIN = 20
 class GridEditor(urwid.WidgetWrap):
     title = None
     columns = None
-    headings = None
 
     def __init__(self, master, value, callback, *cb_args, **cb_kwargs):
         value = copy.deepcopy(value)
@@ -241,7 +275,7 @@ class GridEditor(urwid.WidgetWrap):
         first_width = 20
         if value:
             for r in value:
-                assert len(r) == self.columns
+                assert len(r) == len(self.columns)
                 first_width = max(len(r), first_width)
         self.first_width = min(first_width, FIRST_WIDTH_MAX)
 
@@ -250,9 +284,9 @@ class GridEditor(urwid.WidgetWrap):
         title = urwid.AttrWrap(title, "heading")
 
         headings = []
-        for i, h in enumerate(self.headings):
-            c = urwid.Text(h)
-            if i == 0 and len(self.headings) > 1:
+        for i, col in enumerate(self.columns):
+            c = urwid.Text(col.heading)
+            if i == 0 and len(self.columns) > 1:
                 headings.append(("fixed", first_width + 2, c))
             else:
                 headings.append(c)
@@ -303,6 +337,9 @@ class GridEditor(urwid.WidgetWrap):
             except IOError, v:
                 return str(v)
 
+    def set_subeditor_value(self, val, focus, focus_col):
+        self.walker.set_value(val, focus, focus_col)
+
     def keypress(self, size, key):
         if self.walker.editing:
             if key in ["esc"]:
@@ -317,10 +354,11 @@ class GridEditor(urwid.WidgetWrap):
             return None
 
         key = common.shortcuts(key)
+        column = self.columns[self.walker.focus_col]
         if key in ["q", "esc"]:
             res = []
             for i in self.walker.lst:
-                if not i[1] and any([x.strip() for x in i[0]]):
+                if not i[1] and any([x for x in i[0]]):
                     res.append(i[0])
             self.callback(res, *self.cb_args, **self.cb_kwargs)
             signals.pop_view_state.send(self)
@@ -337,6 +375,13 @@ class GridEditor(urwid.WidgetWrap):
         elif key == "d":
             self.walker.delete_focus()
         elif key == "r":
+            if column.subeditor:
+                signals.status_message.send(
+                    self,
+                    message = "Press enter to edit this field.",
+                    expire = 1000
+                )
+                return
             if self.walker.get_current_value() is not None:
                 signals.status_prompt_path.send(
                     self,
@@ -344,6 +389,13 @@ class GridEditor(urwid.WidgetWrap):
                     callback = self.read_file
                 )
         elif key == "R":
+            if column.subeditor:
+                signals.status_message.send(
+                    self,
+                    message = "Press enter to edit this field.",
+                    expire = 1000
+                )
+                return
             if self.walker.get_current_value() is not None:
                 signals.status_prompt_path.send(
                     self,
@@ -352,6 +404,13 @@ class GridEditor(urwid.WidgetWrap):
                     args = (True,)
                 )
         elif key == "e":
+            if column.subeditor:
+                signals.status_message.send(
+                    self,
+                    message = "Press enter to edit this field.",
+                    expire = 1000
+                )
+                return
             o = self.walker.get_current_value()
             if o is not None:
                 n = self.master.spawn_editor(o.encode("string-escape"))
@@ -359,7 +418,19 @@ class GridEditor(urwid.WidgetWrap):
                 self.walker.set_current_value(n, False)
                 self.walker._modified()
         elif key in ["enter"]:
-            self.walker.start_edit()
+            if column.subeditor:
+                self.master.view_grideditor(
+                    self.columns[self.walker.focus_col].subeditor(
+                        self.master,
+                        self.walker.get_current_value(),
+                        self.set_subeditor_value,
+                        self.walker.focus,
+                        self.walker.focus_col
+                    )
+                )
+            else:
+                self.walker.start_edit()
+
         elif not self.handle_key(key):
             return self._w.keypress(size, key)
 
@@ -403,14 +474,18 @@ class GridEditor(urwid.WidgetWrap):
 
 class QueryEditor(GridEditor):
     title = "Editing query"
-    columns = 2
-    headings = ("Key", "Value")
+    columns = [
+        TextColumn("Key"),
+        TextColumn("Value")
+    ]
 
 
 class HeaderEditor(GridEditor):
     title = "Editing headers"
-    columns = 2
-    headings = ("Key", "Value")
+    columns = [
+        TextColumn("Key"),
+        TextColumn("Value")
+    ]
 
     def make_help(self):
         h = GridEditor.make_help(self)
@@ -448,14 +523,19 @@ class HeaderEditor(GridEditor):
 
 class URLEncodedFormEditor(GridEditor):
     title = "Editing URL-encoded form"
-    columns = 2
-    headings = ("Key", "Value")
+    columns = [
+        TextColumn("Key"),
+        TextColumn("Value")
+    ]
 
 
 class ReplaceEditor(GridEditor):
     title = "Editing replacement patterns"
-    columns = 3
-    headings = ("Filter", "Regex", "Replacement")
+    columns = [
+        TextColumn("Filter"),
+        TextColumn("Regex"),
+        TextColumn("Replacement"),
+    ]
 
     def is_error(self, col, val):
         if col == 0:
@@ -471,8 +551,11 @@ class ReplaceEditor(GridEditor):
 
 class SetHeadersEditor(GridEditor):
     title = "Editing header set patterns"
-    columns = 3
-    headings = ("Filter", "Header", "Value")
+    columns = [
+        TextColumn("Filter"),
+        TextColumn("Header"),
+        TextColumn("Value"),
+    ]
 
     def is_error(self, col, val):
         if col == 0:
@@ -517,14 +600,16 @@ class SetHeadersEditor(GridEditor):
 
 class PathEditor(GridEditor):
     title = "Editing URL path components"
-    columns = 1
-    headings = ("Component",)
+    columns = [
+        TextColumn("Component"),
+    ]
 
 
 class ScriptEditor(GridEditor):
     title = "Editing scripts"
-    columns = 1
-    headings = ("Command",)
+    columns = [
+        TextColumn("Command"),
+    ]
 
     def is_error(self, col, val):
         try:
@@ -535,8 +620,9 @@ class ScriptEditor(GridEditor):
 
 class HostPatternEditor(GridEditor):
     title = "Editing host patterns"
-    columns = 1
-    headings = ("Regex (matched on hostname:port / ip:port)",)
+    columns = [
+        TextColumn("Regex (matched on hostname:port / ip:port)")
+    ]
 
     def is_error(self, col, val):
         try:
@@ -547,11 +633,24 @@ class HostPatternEditor(GridEditor):
 
 class CookieEditor(GridEditor):
     title = "Editing request Cookie header"
-    columns = 2
-    headings = ("Name", "Value")
+    columns = [
+        TextColumn("Name"),
+        TextColumn("Value"),
+    ]
+
+
+class CookieAttributeEditor(GridEditor):
+    title = "Editing Set-Cookie attributes"
+    columns = [
+        TextColumn("Name"),
+        TextColumn("Value"),
+    ]
 
 
 class SetCookieEditor(GridEditor):
-    title = "Editing request SetCookie header"
-    columns = 2
-    headings = ("Name", "Value")
+    title = "Editing response SetCookie header"
+    columns = [
+        TextColumn("Name"),
+        TextColumn("Value"),
+        SubgridColumn("Attributes", CookieAttributeEditor),
+    ]