From dbd75e02f706c7802943d3cd82e1082970c7a0fc Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Mon, 20 Feb 2012 11:29:36 +1300 Subject: Create ODictCaseless for headers, use vanilla ODict for everything else. --- libmproxy/flow.py | 28 +++++++++++++++----- libmproxy/proxy.py | 2 +- test/test_filt.py | 4 +-- test/test_flow.py | 78 ++++++++++++++++++++++++++++++++---------------------- test/tutils.py | 4 +-- 5 files changed, 72 insertions(+), 44 deletions(-) diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 6678771f..63f80dde 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -45,11 +45,14 @@ class ScriptContext: class ODict: + """ + A dictionary-like object for managing ordered (key, value) data. + """ def __init__(self, lst=None): self.lst = lst or [] def _kconv(self, s): - return s.lower() + return s def __eq__(self, other): return self.lst == other.lst @@ -105,7 +108,7 @@ class ODict: Returns a copy of this object. """ lst = copy.deepcopy(self.lst) - return ODict(lst) + return self.__class__(lst) def __repr__(self): elements = [] @@ -143,6 +146,15 @@ class ODict: return count +class ODictCaseless(ODict): + """ + A variant of ODict with "caseless" keys. This version _preserves_ key + case, but does not consider case when setting or getting items. + """ + def _kconv(self, s): + return s.lower() + + class HTTPMsg(controller.Msg): def decode(self): """ @@ -176,7 +188,7 @@ class Request(HTTPMsg): Exposes the following attributes: client_conn: ClientConnection object, or None if this is a replay. - headers: ODict object + headers: ODictCaseless object content: Content of the request, or None scheme: URL scheme (http/https) @@ -188,6 +200,7 @@ class Request(HTTPMsg): method: HTTP method """ def __init__(self, client_conn, host, port, scheme, method, path, headers, content, timestamp=None): + assert isinstance(headers, ODictCaseless) self.client_conn = client_conn self.host, self.port, self.scheme = host, port, scheme self.method, self.path, self.headers, self.content = method, path, headers, content @@ -253,7 +266,7 @@ class Request(HTTPMsg): self.scheme = state["scheme"] self.method = state["method"] self.path = state["path"] - self.headers = ODict._from_state(state["headers"]) + self.headers = ODictCaseless._from_state(state["headers"]) self.content = state["content"] self.timestamp = state["timestamp"] @@ -279,7 +292,7 @@ class Request(HTTPMsg): str(state["scheme"]), str(state["method"]), str(state["path"]), - ODict._from_state(state["headers"]), + ODictCaseless._from_state(state["headers"]), state["content"], state["timestamp"] ) @@ -415,6 +428,7 @@ class Response(HTTPMsg): timestamp: Seconds since the epoch """ def __init__(self, request, code, msg, headers, content, timestamp=None): + assert isinstance(headers, ODictCaseless) self.request = request self.code, self.msg = code, msg self.headers, self.content = headers, content @@ -484,7 +498,7 @@ class Response(HTTPMsg): def _load_state(self, state): self.code = state["code"] self.msg = state["msg"] - self.headers = ODict._from_state(state["headers"]) + self.headers = ODictCaseless._from_state(state["headers"]) self.content = state["content"] self.timestamp = state["timestamp"] @@ -503,7 +517,7 @@ class Response(HTTPMsg): request, state["code"], str(state["msg"]), - ODict._from_state(state["headers"]), + ODictCaseless._from_state(state["headers"]), state["content"], state["timestamp"], ) diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py index ff5e3ec7..8e0aa9cd 100644 --- a/libmproxy/proxy.py +++ b/libmproxy/proxy.py @@ -52,7 +52,7 @@ def read_headers(fp): name = line[:i] value = line[i+1:].strip() ret.append([name, value]) - return flow.ODict(ret) + return flow.ODictCaseless(ret) def read_chunked(fp, limit): diff --git a/test/test_filt.py b/test/test_filt.py index 6c9df232..d9d92ffe 100644 --- a/test/test_filt.py +++ b/test/test_filt.py @@ -74,7 +74,7 @@ class uParsing(libpry.AutoTree): class uMatching(libpry.AutoTree): def req(self): conn = flow.ClientConnect(("one", 2222)) - headers = flow.ODict() + headers = flow.ODictCaseless() headers["header"] = ["qvalue"] return flow.Request( conn, @@ -89,7 +89,7 @@ class uMatching(libpry.AutoTree): def resp(self): q = self.req() - headers = flow.ODict() + headers = flow.ODictCaseless() headers["header_response"] = ["svalue"] return flow.Response( q, diff --git a/test/test_flow.py b/test/test_flow.py index 7ca489c6..52e6781f 100644 --- a/test/test_flow.py +++ b/test/test_flow.py @@ -617,7 +617,7 @@ class uFlowMaster(libpry.AutoTree): class uRequest(libpry.AutoTree): def test_simple(self): - h = flow.ODict() + h = flow.ODictCaseless() h["test"] = ["test"] c = flow.ClientConnect(("addr", 2222)) r = flow.Request(c, "host", 22, "https", "GET", "/", h, "content") @@ -639,7 +639,7 @@ class uRequest(libpry.AutoTree): assert r._assemble(True) def test_getset_form_urlencoded(self): - h = flow.ODict() + h = flow.ODictCaseless() h["content-type"] = [flow.HDR_FORM_URLENCODED] d = flow.ODict([("one", "two"), ("three", "four")]) r = flow.Request(None, "host", 22, "https", "GET", "/", h, utils.urlencode(d.lst)) @@ -653,7 +653,7 @@ class uRequest(libpry.AutoTree): assert not r.get_form_urlencoded() def test_getset_query(self): - h = flow.ODict() + h = flow.ODictCaseless() r = flow.Request(None, "host", 22, "https", "GET", "/foo?x=y&a=b", h, "content") q = r.get_query() @@ -676,7 +676,7 @@ class uRequest(libpry.AutoTree): assert r.get_query() == qv def test_anticache(self): - h = flow.ODict() + h = flow.ODictCaseless() r = flow.Request(None, "host", 22, "https", "GET", "/", h, "content") h["if-modified-since"] = ["test"] h["if-none-match"] = ["test"] @@ -685,7 +685,7 @@ class uRequest(libpry.AutoTree): assert not "if-none-match" in r.headers def test_getset_state(self): - h = flow.ODict() + h = flow.ODictCaseless() h["test"] = ["test"] c = flow.ClientConnect(("addr", 2222)) r = flow.Request(c, "host", 22, "https", "GET", "/", h, "content") @@ -753,7 +753,7 @@ class uRequest(libpry.AutoTree): class uResponse(libpry.AutoTree): def test_simple(self): - h = flow.ODict() + h = flow.ODictCaseless() h["test"] = ["test"] c = flow.ClientConnect(("addr", 2222)) req = flow.Request(c, "host", 22, "https", "GET", "/", h, "content") @@ -798,7 +798,7 @@ class uResponse(libpry.AutoTree): def test_getset_state(self): - h = flow.ODict() + h = flow.ODictCaseless() h["test"] = ["test"] c = flow.ClientConnect(("addr", 2222)) req = flow.Request(c, "host", 22, "https", "GET", "/", h, "content") @@ -886,31 +886,31 @@ class uClientConnect(libpry.AutoTree): class uODict(libpry.AutoTree): def setUp(self): - self.hd = flow.ODict() + self.od = flow.ODict() def test_str_err(self): h = flow.ODict() libpry.raises(ValueError, h.__setitem__, "key", "foo") def test_dictToHeader1(self): - self.hd.add("one", "uno") - self.hd.add("two", "due") - self.hd.add("two", "tre") + self.od.add("one", "uno") + self.od.add("two", "due") + self.od.add("two", "tre") expected = [ "one: uno\r\n", "two: due\r\n", "two: tre\r\n", "\r\n" ] - out = repr(self.hd) + out = repr(self.od) for i in expected: assert out.find(i) >= 0 def test_dictToHeader2(self): - self.hd["one"] = ["uno"] + self.od["one"] = ["uno"] expected1 = "one: uno\r\n" expected2 = "\r\n" - out = repr(self.hd) + out = repr(self.od) assert out.find(expected1) >= 0 assert out.find(expected2) >= 0 @@ -924,36 +924,49 @@ class uODict(libpry.AutoTree): assert not h.match_re("nonono") def test_getset_state(self): - self.hd.add("foo", 1) - self.hd.add("foo", 2) - self.hd.add("bar", 3) - state = self.hd._get_state() + self.od.add("foo", 1) + self.od.add("foo", 2) + self.od.add("bar", 3) + state = self.od._get_state() nd = flow.ODict._from_state(state) - assert nd == self.hd + assert nd == self.od def test_copy(self): - self.hd.add("foo", 1) - self.hd.add("foo", 2) - self.hd.add("bar", 3) - assert self.hd == self.hd.copy() + self.od.add("foo", 1) + self.od.add("foo", 2) + self.od.add("bar", 3) + assert self.od == self.od.copy() def test_del(self): - self.hd.add("foo", 1) - self.hd.add("Foo", 2) - self.hd.add("bar", 3) - del self.hd["foo"] - assert len(self.hd.lst) == 1 + self.od.add("foo", 1) + self.od.add("Foo", 2) + self.od.add("bar", 3) + del self.od["foo"] + assert len(self.od.lst) == 2 def test_replace(self): - self.hd.add("one", "two") - self.hd.add("two", "one") - assert self.hd.replace("one", "vun") == 2 - assert self.hd.lst == [ + self.od.add("one", "two") + self.od.add("two", "one") + assert self.od.replace("one", "vun") == 2 + assert self.od.lst == [ ["vun", "two"], ["two", "vun"], ] +class uODictCaseless(libpry.AutoTree): + def setUp(self): + self.od = flow.ODictCaseless() + + def test_del(self): + self.od.add("foo", 1) + self.od.add("Foo", 2) + self.od.add("bar", 3) + del self.od["foo"] + assert len(self.od) == 1 + + + tests = [ uStickyCookieState(), uStickyAuthState(), @@ -968,4 +981,5 @@ tests = [ uError(), uClientConnect(), uODict(), + uODictCaseless(), ] diff --git a/test/tutils.py b/test/tutils.py index c22eb1a0..b11c997a 100644 --- a/test/tutils.py +++ b/test/tutils.py @@ -7,7 +7,7 @@ import random def treq(conn=None): if not conn: conn = flow.ClientConnect(("address", 22)) - headers = flow.ODict() + headers = flow.ODictCaseless() headers["header"] = ["qvalue"] return flow.Request(conn, "host", 80, "http", "GET", "/path", headers, "content") @@ -15,7 +15,7 @@ def treq(conn=None): def tresp(req=None): if not req: req = treq() - headers = flow.ODict() + headers = flow.ODictCaseless() headers["header_response"] = ["svalue"] return flow.Response(req, 200, "message", headers, "content_response") -- cgit v1.2.3