From 79039eb5d23b6f7076664a3383988cd6b51e377e Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Thu, 24 Feb 2011 15:15:51 +1300 Subject: More mature sticky cookie primitive. Use it in console.py. --- libmproxy/console.py | 19 ---------------- libmproxy/flow.py | 62 +++++++++++++++++++++++++++++++++++++++------------- libmproxy/utils.py | 8 +++++++ test/test_flow.py | 56 ++++++++++++++++++++++++++++++++++++++++------- 4 files changed, 103 insertions(+), 42 deletions(-) diff --git a/libmproxy/console.py b/libmproxy/console.py index 187502a6..f63115aa 100644 --- a/libmproxy/console.py +++ b/libmproxy/console.py @@ -1154,15 +1154,6 @@ class ConsoleMaster(flow.FlowMaster): else: self.state.beep = None - def set_stickycookie(self, txt): - if txt: - self.stickycookie = filt.parse(txt) - if not self.stickycookie: - return "Invalid filter expression." - else: - self.stickyhosts = {} - self.stickycookie = None - def drawscreen(self): size = self.ui.get_cols_rows() canvas = self.view.render(size, focus=1) @@ -1311,20 +1302,10 @@ class ConsoleMaster(flow.FlowMaster): def handle_request(self, r): f = flow.FlowMaster.handle_request(self, r) if f: - if f.match(self.stickycookie): - hid = (f.request.host, f.request.port) - if f.request.headers.has_key("cookie"): - self.stickyhosts[hid] = f.request.headers["cookie"] - elif hid in self.stickyhosts: - f.request.headers["cookie"] = self.stickyhosts[hid] self.process_flow(f, r) def handle_response(self, r): f = flow.FlowMaster.handle_response(self, r) if f: - if f.match(self.stickycookie): - hid = (f.request.host, f.request.port) - if f.response.headers.has_key("set-cookie"): - self.stickyhosts[hid] = f.response.headers["set-cookie"] self.process_flow(f, r) diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 9636c3bd..9b083036 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -3,7 +3,7 @@ with their responses, and provide filtering and interception facilities. """ import subprocess, base64, sys, json, hashlib, Cookie, cookielib, copy -import proxy, threading, netstring +import proxy, threading, netstring, filt import controller class RunException(Exception): @@ -90,26 +90,44 @@ class ServerPlaybackState: class StickyCookieState: - def __init__(self): + def __init__(self, flt): + """ + flt: A compiled filter. + """ self.jar = {} + self.flt = flt - def ckey(self, c): - c = copy.copy(c) - del c["expires"] - return str(c) + def ckey(self, m, f): + """ + Returns a (domain, port, path) tuple. + """ + return ( + m["domain"] or f.request.host, + f.request.port, + m["path"] or "/" + ) - def add_cookies(self, headers): - for i in headers: + def handle_response(self, f): + for i in f.response.headers.get("set-cookie", []): c = Cookie.SimpleCookie(i) m = c.values()[0] - self.jar[self.ckey(m)] = m + k = self.ckey(m, f) + if cookielib.domain_match(f.request.host, k[0]): + self.jar[self.ckey(m, f)] = m + + def handle_request(self, f): + if f.match(self.flt): + cs = [] + for i in self.jar.keys(): + match = [ + cookielib.domain_match(i[0], f.request.host), + f.request.port == i[1], + f.request.path.startswith(i[2]) + ] + if all(match): + l = f.request.headers.setdefault("cookie", []) + l.append(self.jar[i].output(header="").strip()) - def get_cookies(self, domain, path): - cs = [] - for i in self.jar.values(): - if cookielib.domain_match(domain, i["domain"]) and path.startswith(i.get("path", "/")): - cs.append(i) - return cs class Flow: @@ -369,6 +387,7 @@ class FlowMaster(controller.Master): self.playback = None self.scripts = {} self.kill_nonreplay = False + self.stickycookie_state = False def _runscript(self, f, script): return f.run_script(script) @@ -379,6 +398,15 @@ class FlowMaster(controller.Master): def set_request_script(self, s): self.scripts["request"] = s + def set_stickycookie(self, txt): + if txt: + flt = filt.parse(txt) + if not flt: + return "Invalid filter expression." + self.stickycookie_state = StickyCookieState(flt) + else: + self.stickycookie_state = None + def start_playback(self, flows, kill, headers): """ flows: A list of flows. @@ -419,6 +447,8 @@ class FlowMaster(controller.Master): def handle_request(self, r): f = self.state.add_request(r) + if self.stickycookie_state: + self.stickycookie_state.handle_request(f) if "request" in self.scripts: self._runscript(f, self.scripts["request"]) if self.playback: @@ -434,6 +464,8 @@ class FlowMaster(controller.Master): f = self.state.add_response(r) if not f: r.ack() + if self.stickycookie_state: + self.stickycookie_state.handle_response(f) if "response" in self.scripts: self._runscript(f, self.scripts["response"]) return f diff --git a/libmproxy/utils.py b/libmproxy/utils.py index afef8e63..c67b9397 100644 --- a/libmproxy/utils.py +++ b/libmproxy/utils.py @@ -172,6 +172,10 @@ class MultiDict: key = self._helper[0](key) return self._d.get(key, d) + def __contains__(self, key): + key = self._helper[0](key) + return self._d.__contains__(key) + def __eq__(self, other): return dict(self) == dict(other) @@ -192,6 +196,10 @@ class MultiDict: key = self._helper[0](key) return self._d.has_key(key) + def setdefault(self, key, default=None): + key = self._helper[0](key) + return self._d.setdefault(key, default) + def keys(self): return self._d.keys() diff --git a/test/test_flow.py b/test/test_flow.py index f58b1dd0..10b445db 100644 --- a/test/test_flow.py +++ b/test/test_flow.py @@ -5,14 +5,33 @@ import libpry class uStickyCookieState(libpry.AutoTree): - def test_simple(self): - s = flow.StickyCookieState() - s.add_cookies( - ["SSID=mooo, FOO=bar; Domain=.google.com; Path=/; Expires=Wed, 13-Jan-2021 22:23:01 GMT; Secure; "] - ) - assert len(s.jar) == 1 - assert len(s.get_cookies("www.google.com", "/foo")) == 1 - assert len(s.get_cookies("www.foo.com", "/foo")) == 0 + def _response(self, cookie, host): + s = flow.StickyCookieState(filt.parse(".*")) + f = utils.tflow_full() + f.request.host = host + f.response.headers["Set-Cookie"] = [cookie] + s.handle_response(f) + return s, f + + def test_handle_response(self): + c = "SSID=mooo, FOO=bar; Domain=.google.com; Path=/; "\ + "Expires=Wed, 13-Jan-2021 22:23:01 GMT; Secure; " + + s, f = self._response(c, "host") + assert not s.jar.keys() + + s, f = self._response(c, "www.google.com") + assert s.jar.keys() + + s, f = self._response("SSID=mooo", "www.google.com") + assert s.jar.keys()[0] == ('www.google.com', 80, '/') + + def test_handle_request(self): + s, f = self._response("SSID=mooo", "www.google.com") + assert "cookie" not in f.request.headers + s.handle_request(f) + assert "cookie" in f.request.headers + class uServerPlaybackState(libpry.AutoTree): @@ -353,6 +372,27 @@ class uFlowMaster(libpry.AutoTree): r.request.content = "gibble" assert not fm.do_playback(r) + def test_stickycookie(self): + s = flow.State() + fm = flow.FlowMaster(None, s) + assert "Invalid" in fm.set_stickycookie("~h") + fm.set_stickycookie(".*") + assert fm.stickycookie_state + fm.set_stickycookie(None) + assert not fm.stickycookie_state + + fm.set_stickycookie(".*") + tf = utils.tflow_full() + tf.response.headers["set-cookie"] = ["foo=bar"] + fm.handle_request(tf.request) + f = fm.handle_response(tf.response) + assert fm.stickycookie_state.jar + assert not "cookie" in tf.request.headers + fm.handle_request(tf.request) + assert tf.request.headers["cookie"] == ["foo=bar"] + + + tests = [ -- cgit v1.2.3