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.

10 files changed, 125 insertions, 6 deletions

diff --git a/doc-src/index.html b/doc-src/index.html
index b7561350..904c1550 100644
--- a/doc-src/index.html
+++ b/doc-src/index.html
@@ -4,7 +4,7 @@
     <li><a href="@!urlTo("interception.html")!@">Interception</a></li>
     <li><a href="@!urlTo("clientreplay.html")!@">Client-side replay</a></li>
     <li><a href="@!urlTo("serverreplay.html")!@">Server-side replay</a></li>
-    <li><a href="@!urlTo("stickycookies.html")!@">Sticky cookies</a></li>
+    <li><a href="@!urlTo("sticky.html")!@">Sticky cookies and auth</a></li>
     <li><a href="@!urlTo("anticache.html")!@">Anticache</a></li>
     <li><a href="@!urlTo("filters.html")!@">Filter expressions</a></li>
     <li><a href="@!urlTo("scripts.html")!@">Scripting API</a></li>
diff --git a/doc-src/index.py b/doc-src/index.py
index 626aaaa1..6568195b 100644
--- a/doc-src/index.py
+++ b/doc-src/index.py
@@ -70,7 +70,7 @@ pages = [
     Page("interception.html", "Interception"),
     Page("clientreplay.html", "Client-side replay"),
     Page("serverreplay.html", "Server-side replay"),
-    Page("stickycookies.html", "Sticky cookies"),
+    Page("sticky.html", "Sticky cookies and auth"),
     Page("anticache.html", "Anticache"),
     Page("filters.html", "Filter expressions"),
     Page("scripts.html", "External scripts"),
diff --git a/doc-src/stickycookies.html b/doc-src/sticky.html
index 5aa5045c..32ccdd2d 100644
--- a/doc-src/stickycookies.html
+++ b/doc-src/sticky.html
@@ -15,3 +15,10 @@ process once, and simply replay it on startup every time you need to interact
 with the secured resources.
 
 
+## Sticky auth
+
+The __stickyauth__ option is analagous to the __stickycookie__ option, in that
+HTTP __Authorization__ headers are simply replayed to the server once they have
+been seen. This is enough to allow you to access a server resource using HTTP
+Basic authentication through the proxy. Note that __mitmproxy__ doesn't (yet)
+support replay of HTTP Digest authentication. 
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):
         """
diff --git a/test/test_flow.py b/test/test_flow.py
index 463375aa..a6489462 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -34,6 +34,19 @@ class uStickyCookieState(libpry.AutoTree):
         assert "cookie" in f.request.headers
 
 
+class uStickyAuthState(libpry.AutoTree):
+    def test_handle_response(self):
+        s = flow.StickyAuthState(filt.parse(".*"))
+        f = tutils.tflow_full()
+        f.request.headers["authorization"] = ["foo"]
+        s.handle_request(f)
+        assert "host" in s.hosts
+
+        f = tutils.tflow_full()
+        s.handle_request(f)
+        assert f.request.headers["authorization"] == ["foo"]
+
+
 class uClientPlaybackState(libpry.AutoTree):
     def test_tick(self):
         first = tutils.tflow()
@@ -62,6 +75,9 @@ class uClientPlaybackState(libpry.AutoTree):
         fm.state.clear()
         fm.tick(q)
 
+        fm.stop_client_playback()
+        assert not fm.client_playback
+
 
 class uServerPlaybackState(libpry.AutoTree):
     def test_hash(self):
@@ -441,6 +457,9 @@ class uFlowMaster(libpry.AutoTree):
         fm.tick(q)
         assert controller.exit
 
+        fm.stop_server_playback()
+        assert not fm.server_playback
+
     def test_stickycookie(self):
         s = flow.State()
         fm = flow.FlowMaster(None, s)
@@ -460,9 +479,30 @@ class uFlowMaster(libpry.AutoTree):
         fm.handle_request(tf.request)
         assert tf.request.headers["cookie"] == ["foo=bar"]
 
+    def test_stickyauth(self):
+        s = flow.State()
+        fm = flow.FlowMaster(None, s)
+        assert "Invalid" in fm.set_stickyauth("~h")
+        fm.set_stickyauth(".*")
+        assert fm.stickyauth_state
+        fm.set_stickyauth(None)
+        assert not fm.stickyauth_state
+
+        fm.set_stickyauth(".*")
+        tf = tutils.tflow_full()
+        tf.request.headers["authorization"] = ["foo"]
+        fm.handle_request(tf.request)
+
+        f = tutils.tflow_full()
+        assert fm.stickyauth_state.hosts
+        assert not "authorization" in f.request.headers
+        fm.handle_request(f.request)
+        assert f.request.headers["authorization"] == ["foo"]
+
 
 tests = [
     uStickyCookieState(),
+    uStickyAuthState(),
     uServerPlaybackState(),
     uClientPlaybackState(),
     uFlow(),
diff --git a/todo b/todo
index be2958d6..2a69291a 100644
--- a/todo
+++ b/todo
@@ -20,6 +20,9 @@ Features:
     - No-proxy mode, for when we want to just replay without serving as a proxy ourselves.
     - Upstream proxy support.
     - Strings view for binary responses.
+    - Support HTTP Digest authentication through the stickyauth option. We'll
+    have to save the server nonce, and recalculate the hashes for each request.
+
     - Chunked encoding support for requests (we already support it for responses).
     - Better manipulation of flows in mitmproxy: copying flows, adding
     comments.