Add a stickyauth option.

This allows us to replay an HTTP Authorization header, in the same way as we
replay cookies using stickycookies. This lets us conveniently get at HTTP Basic
Auth protected resources through the proxy, but is not enough to do the same
for HTTP Digest auth. We'll put that on the todo list.

5 files changed, 73 insertions, 4 deletions

diff --git a/libmproxy/cmdline.py b/libmproxy/cmdline.py
index 78a88e9e..4ff7cbe0 100644
--- a/libmproxy/cmdline.py
+++ b/libmproxy/cmdline.py
@@ -3,11 +3,17 @@ import optparse
 
 
 def get_common_options(options):
-    stickycookie = None
+    stickycookie, stickyauth = None, None
     if options.stickycookie_all:
         stickycookie = ".*"
     elif options.stickycookie_filt:
         stickycookie = options.stickycookie_filt
+
+    if options.stickyauth_all:
+        stickyauth = ".*"
+    elif options.stickyauth_filt:
+        stickyauth = options.stickyauth_filt
+
     return dict(
         verbosity = options.verbose,
         wfile = options.wfile,
@@ -18,6 +24,7 @@ def get_common_options(options):
         rheaders = options.rheaders,
         client_replay = options.client_replay,
         stickycookie = stickycookie,
+        stickyauth = stickyauth,
         anticache = options.anticache,
         refresh_server_playback = not options.norefresh,
     )
@@ -70,6 +77,16 @@ def common_options(parser):
         help="Set sticky cookie filter. Matched against requests."
     )
     parser.add_option(
+        "-u",
+        action="store_true", dest="stickyauth_all", default=None,
+        help="Set sticky auth for all requests."
+    )
+    parser.add_option(
+        "-U",
+        action="store", dest="stickyauth_filt", default=None, metavar="FILTER",
+        help="Set sticky auth filter. Matched against requests."
+    )
+    parser.add_option(
         "-v",
         action="count", dest="verbose", default=1,
         help="Increase verbosity. Can be passed multiple times."
diff --git a/libmproxy/console.py b/libmproxy/console.py
index 46bd0078..bbf69d2d 100644
--- a/libmproxy/console.py
+++ b/libmproxy/console.py
@@ -635,6 +635,10 @@ class StatusBar(WWrap):
             r.append("[")
             r.append(("statusbar_highlight", "t"))
             r.append(":%s]"%self.master.stickycookie_txt)
+        if self.master.stickyauth_txt:
+            r.append("[")
+            r.append(("statusbar_highlight", "u"))
+            r.append(":%s]"%self.master.stickyauth_txt)
 
         opts = []
         if self.master.anticache:
@@ -773,6 +777,7 @@ class Options(object):
         "rheaders",
         "server_replay",
         "stickycookie",
+        "stickyauth",
         "verbosity",
         "wfile",
     ]
@@ -828,8 +833,10 @@ class ConsoleMaster(flow.FlowMaster):
             print >> sys.stderr, "Sticky cookies error:", r
             sys.exit(1)
 
-        self.stickycookie = None
-        self.stickyhosts = {}
+        r = self.set_stickyauth(options.stickyauth)
+        if r:
+            print >> sys.stderr, "Sticky auth error:", r
+            sys.exit(1)
 
         self.refresh_server_playback = options.refresh_server_playback
         self.anticache = options.anticache
@@ -1136,6 +1143,7 @@ class ConsoleMaster(flow.FlowMaster):
             ("S", "save all flows matching current limit"),
             ("s", "server replay"),
             ("t", "set sticky cookie expression"),
+            ("u", "set sticky auth expression"),
             ("w", "save this flow"),
             ("page up/down", "page up/down"),
             ("enter", "view connection"),
@@ -1410,6 +1418,13 @@ class ConsoleMaster(flow.FlowMaster):
                                 self.set_stickycookie
                             )
                             k = None
+                        elif k == "u":
+                            self.prompt(
+                                "Sticky auth filter: ",
+                                self.stickyauth_txt,
+                                self.set_stickyauth
+                            )
+                            k = None
                     if k:
                         self.view.keypress(size, k)
         except (Stop, KeyboardInterrupt):
diff --git a/libmproxy/dump.py b/libmproxy/dump.py
index 97232ac7..6f18c8fe 100644
--- a/libmproxy/dump.py
+++ b/libmproxy/dump.py
@@ -16,6 +16,7 @@ class Options(object):
         "rheaders",
         "server_replay",
         "stickycookie",
+        "stickyauth",
         "verbosity",
         "wfile",
     ]
@@ -64,6 +65,9 @@ class DumpMaster(flow.FlowMaster):
         if options.stickycookie:
             self.set_stickycookie(options.stickycookie)
 
+        if options.stickyauth:
+            self.set_stickyauth(options.stickyauth)
+
         if options.wfile:
             path = os.path.expanduser(options.wfile)
             try:
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index eed006b4..387c49f0 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -163,6 +163,21 @@ class StickyCookieState:
                     l.append(self.jar[i].output(header="").strip())
 
 
+class StickyAuthState:
+    def __init__(self, flt):
+        """
+            flt: A compiled filter.
+        """
+        self.flt = flt
+        self.hosts = {}
+
+    def handle_request(self, f):
+        if "authorization" in f.request.headers:
+            self.hosts[f.request.host] = f.request.headers["authorization"]
+        elif f.match(self.flt):
+            if f.request.host in self.hosts:
+                f.request.headers["authorization"] = self.hosts[f.request.host]
+
 
 class Flow:
     def __init__(self, request):
@@ -433,6 +448,9 @@ class FlowMaster(controller.Master):
         self.stickycookie_state = False
         self.stickycookie_txt = None
 
+        self.stickyauth_state = False
+        self.stickyauth_txt = None
+
         self.anticache = False
         self.refresh_server_playback = False
 
@@ -458,6 +476,17 @@ class FlowMaster(controller.Master):
             self.stickycookie_state = None
             self.stickycookie_txt = None
 
+    def set_stickyauth(self, txt):
+        if txt:
+            flt = filt.parse(txt)
+            if not flt:
+                return "Invalid filter expression."
+            self.stickyauth_state = StickyAuthState(flt)
+            self.stickyauth_txt = txt
+        else:
+            self.stickyauth_state = None
+            self.stickyauth_txt = None
+
     def start_client_playback(self, flows, exit):
         """
             flows: A list of flows.
@@ -516,6 +545,9 @@ class FlowMaster(controller.Master):
     def process_new_request(self, f):
         if self.stickycookie_state:
             self.stickycookie_state.handle_request(f)
+        if self.stickyauth_state:
+            self.stickyauth_state.handle_request(f)
+
         if "request" in self.scripts:
             self._runscript(f, self.scripts["request"])
         if self.anticache:
diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py
index 1f6dafa8..e9561848 100644
--- a/libmproxy/proxy.py
+++ b/libmproxy/proxy.py
@@ -137,8 +137,9 @@ class Request(controller.Msg):
         self.close = False
         controller.Msg.__init__(self)
 
-        # Have this request's cookies been modified by sticky cookies?
+        # Have this request's cookies been modified by sticky cookies or auth?
         self.stickycookie = False
+        self.stickyauth = False
 
     def anticache(self):
         """