From 2616f490fee1b732aa0853318d67a550fc561cc4 Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Mon, 20 Feb 2012 10:34:32 +1300 Subject: Rename Headers class to ODict ODict is an ordered dictionary class that will be useful in many other parts of our API. --- libmproxy/console/connview.py | 4 +-- libmproxy/flow.py | 64 +++++++++++++------------------------------ libmproxy/proxy.py | 31 ++++++++++++++++++--- test/test_filt.py | 4 +-- test/test_flow.py | 64 +++++++++---------------------------------- test/test_proxy.py | 48 ++++++++++++++++++++++++++++++-- test/tutils.py | 4 +-- 7 files changed, 110 insertions(+), 109 deletions(-) diff --git a/libmproxy/console/connview.py b/libmproxy/console/connview.py index ed038a11..b722f78e 100644 --- a/libmproxy/console/connview.py +++ b/libmproxy/console/connview.py @@ -371,7 +371,7 @@ class ConnectionView(common.WWrap): self.master.refresh_connection(self.flow) def set_headers(self, lst, conn): - conn.headers = flow.Headers(lst) + conn.headers = flow.ODict(lst) def set_query(self, lst, conn): conn.set_query(lst) @@ -384,7 +384,7 @@ class ConnectionView(common.WWrap): conn = self.flow.request else: if not self.flow.response: - self.flow.response = flow.Response(self.flow.request, 200, "OK", flow.Headers(), "") + self.flow.response = flow.Response(self.flow.request, 200, "OK", flow.ODict(), "") conn = self.flow.response self.flow.backup() diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 024d23e6..f29dcf27 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -44,7 +44,7 @@ class ScriptContext: self._master.replay_request(f) -class Headers: +class ODict: def __init__(self, lst=None): self.lst = lst or [] @@ -69,12 +69,12 @@ class Headers: new.append(i) return new - def __setitem__(self, k, hdrs): - if isinstance(hdrs, basestring): - raise ValueError("Header values should be lists.") + def __setitem__(self, k, values): + if isinstance(values, basestring): + raise ValueError("ODict values should be lists.") k = self._kconv(k) new = self._filter_lst(k, self.lst) - for i in hdrs: + for i in values: new.append((k, i)) self.lst = new @@ -102,22 +102,19 @@ class Headers: Returns a copy of this object. """ lst = copy.deepcopy(self.lst) - return Headers(lst) + return ODict(lst) def __repr__(self): - """ - Returns a string containing a formatted header string. - """ - headerElements = [] + elements = [] for itm in self.lst: - headerElements.append(itm[0] + ": " + itm[1]) - headerElements.append("") - return "\r\n".join(headerElements) + elements.append(itm[0] + ": " + itm[1]) + elements.append("") + return "\r\n".join(elements) def match_re(self, expr): """ - Match the regular expression against each header. For each (key, - value) pair a string of the following format is matched against: + Match the regular expression against each (key, value) pair. For + each pair a string of the following format is matched against: "key: value" """ @@ -127,32 +124,9 @@ class Headers: return True return False - def read(self, fp): - """ - Read a set of headers from a file pointer. Stop once a blank line - is reached. - """ - ret = [] - name = '' - while 1: - line = fp.readline() - if not line or line == '\r\n' or line == '\n': - break - if line[0] in ' \t': - # continued header - ret[-1][1] = ret[-1][1] + '\r\n ' + line.strip() - else: - i = line.find(':') - # We're being liberal in what we accept, here. - if i > 0: - name = line[:i] - value = line[i+1:].strip() - ret.append([name, value]) - self.lst = ret - def replace(self, pattern, repl, *args, **kwargs): """ - Replaces a regular expression pattern with repl in both header keys + Replaces a regular expression pattern with repl in both keys and values. Returns the number of replacements made. """ nlst, count = [], 0 @@ -199,7 +173,7 @@ class Request(HTTPMsg): Exposes the following attributes: client_conn: ClientConnection object, or None if this is a replay. - headers: Headers object + headers: ODict object content: Content of the request, or None scheme: URL scheme (http/https) @@ -276,7 +250,7 @@ class Request(HTTPMsg): self.scheme = state["scheme"] self.method = state["method"] self.path = state["path"] - self.headers = Headers._from_state(state["headers"]) + self.headers = ODict._from_state(state["headers"]) self.content = state["content"] self.timestamp = state["timestamp"] @@ -302,7 +276,7 @@ class Request(HTTPMsg): str(state["scheme"]), str(state["method"]), str(state["path"]), - Headers._from_state(state["headers"]), + ODict._from_state(state["headers"]), state["content"], state["timestamp"] ) @@ -434,7 +408,7 @@ class Response(HTTPMsg): request: Request object. code: HTTP response code msg: HTTP response message - headers: Headers object + headers: ODict object content: Response content timestamp: Seconds since the epoch """ @@ -508,7 +482,7 @@ class Response(HTTPMsg): def _load_state(self, state): self.code = state["code"] self.msg = state["msg"] - self.headers = Headers._from_state(state["headers"]) + self.headers = ODict._from_state(state["headers"]) self.content = state["content"] self.timestamp = state["timestamp"] @@ -527,7 +501,7 @@ class Response(HTTPMsg): request, state["code"], str(state["msg"]), - Headers._from_state(state["headers"]), + ODict._from_state(state["headers"]), state["content"], state["timestamp"], ) diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py index 43277a6c..ff5e3ec7 100644 --- a/libmproxy/proxy.py +++ b/libmproxy/proxy.py @@ -30,6 +30,31 @@ class ProxyConfig: self.reverse_proxy = reverse_proxy + +def read_headers(fp): + """ + Read a set of headers from a file pointer. Stop once a blank line + is reached. Return a ODict object. + """ + ret = [] + name = '' + while 1: + line = fp.readline() + if not line or line == '\r\n' or line == '\n': + break + if line[0] in ' \t': + # continued header + ret[-1][1] = ret[-1][1] + '\r\n ' + line.strip() + else: + i = line.find(':') + # We're being liberal in what we accept, here. + if i > 0: + name = line[:i] + value = line[i+1:].strip() + ret.append([name, value]) + return flow.ODict(ret) + + def read_chunked(fp, limit): content = "" total = 0 @@ -224,8 +249,7 @@ class ServerConnection: code = int(code) except ValueError: raise ProxyError(502, "Invalid server response: %s."%line) - headers = flow.Headers() - headers.read(self.rfile) + headers = read_headers(self.rfile) if code >= 100 and code <= 199: return self.read_response() if self.request.method == "HEAD" or code == 204 or code == 304: @@ -350,8 +374,7 @@ class ProxyHandler(SocketServer.StreamRequestHandler): method, scheme, host, port, path, httpminor = parse_request_line(self.rfile.readline()) if scheme is None: scheme = "https" - headers = flow.Headers() - headers.read(self.rfile) + headers = read_headers(self.rfile) if host is None and "host" in headers: netloc = headers["host"][0] if ':' in netloc: diff --git a/test/test_filt.py b/test/test_filt.py index dfb2b3e1..6c9df232 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.Headers() + headers = flow.ODict() headers["header"] = ["qvalue"] return flow.Request( conn, @@ -89,7 +89,7 @@ class uMatching(libpry.AutoTree): def resp(self): q = self.req() - headers = flow.Headers() + headers = flow.ODict() headers["header_response"] = ["svalue"] return flow.Response( q, diff --git a/test/test_flow.py b/test/test_flow.py index 3e126263..250127a9 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.Headers() + h = flow.ODict() 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.Headers() + h = flow.ODict() h["content-type"] = [flow.HDR_FORM_URLENCODED] d = [("one", "two"), ("three", "four")] r = flow.Request(None, "host", 22, "https", "GET", "/", h, utils.urlencode(d)) @@ -653,7 +653,7 @@ class uRequest(libpry.AutoTree): assert not r.get_form_urlencoded() def test_getset_query(self): - h = flow.Headers() + h = flow.ODict() 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.Headers() + h = flow.ODict() 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.Headers() + h = flow.ODict() 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.Headers() + h = flow.ODict() 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.Headers() + h = flow.ODict() h["test"] = ["test"] c = flow.ClientConnect(("addr", 2222)) req = flow.Request(c, "host", 22, "https", "GET", "/", h, "content") @@ -884,52 +884,14 @@ class uClientConnect(libpry.AutoTree): assert c3 == c -class uHeaders(libpry.AutoTree): +class uODict(libpry.AutoTree): def setUp(self): - self.hd = flow.Headers() + self.hd = flow.ODict() def test_str_err(self): - h = flow.Headers() + h = flow.ODict() libpry.raises(ValueError, h.__setitem__, "key", "foo") - def test_read_simple(self): - data = """ - Header: one - Header2: two - \r\n - """ - data = textwrap.dedent(data) - data = data.strip() - s = StringIO(data) - self.hd.read(s) - assert self.hd["header"] == ["one"] - assert self.hd["header2"] == ["two"] - - def test_read_multi(self): - data = """ - Header: one - Header: two - \r\n - """ - data = textwrap.dedent(data) - data = data.strip() - s = StringIO(data) - self.hd.read(s) - assert self.hd["header"] == ["one", "two"] - - def test_read_continued(self): - data = """ - Header: one - \ttwo - Header2: three - \r\n - """ - data = textwrap.dedent(data) - data = data.strip() - s = StringIO(data) - self.hd.read(s) - assert self.hd["header"] == ['one\r\n two'] - def test_dictToHeader1(self): self.hd.add("one", "uno") self.hd.add("two", "due") @@ -953,7 +915,7 @@ class uHeaders(libpry.AutoTree): assert out.find(expected2) >= 0 def test_match_re(self): - h = flow.Headers() + h = flow.ODict() h.add("one", "uno") h.add("two", "due") h.add("two", "tre") @@ -966,7 +928,7 @@ class uHeaders(libpry.AutoTree): self.hd.add("foo", 2) self.hd.add("bar", 3) state = self.hd._get_state() - nd = flow.Headers._from_state(state) + nd = flow.ODict._from_state(state) assert nd == self.hd def test_copy(self): @@ -1005,5 +967,5 @@ tests = [ uResponse(), uError(), uClientConnect(), - uHeaders(), + uODict(), ] diff --git a/test/test_proxy.py b/test/test_proxy.py index f0b54681..8120bbbe 100644 --- a/test/test_proxy.py +++ b/test/test_proxy.py @@ -1,4 +1,5 @@ -import cStringIO, time +import cStringIO, time, textwrap +from cStringIO import StringIO import libpry from libmproxy import proxy, controller, utils, dump, flow @@ -27,7 +28,7 @@ class u_read_http_body(libpry.AutoTree): def test_all(self): d = Dummy() - h = flow.Headers() + h = flow.ODict() s = cStringIO.StringIO("testing") assert proxy.read_http_body(s, d, h, False, None) == "" @@ -42,7 +43,7 @@ class u_read_http_body(libpry.AutoTree): libpry.raises(proxy.ProxyError, proxy.read_http_body, s, d, h, False, 4) - h = flow.Headers() + h = flow.ODict() s = cStringIO.StringIO("testing") assert len(proxy.read_http_body(s, d, h, True, 4)) == 4 s = cStringIO.StringIO("testing") @@ -94,10 +95,51 @@ class uProxyError(libpry.AutoTree): assert repr(p) +class u_read_headers(libpry.AutoTree): + def test_read_simple(self): + data = """ + Header: one + Header2: two + \r\n + """ + data = textwrap.dedent(data) + data = data.strip() + s = StringIO(data) + headers = proxy.read_headers(s) + assert headers["header"] == ["one"] + assert headers["header2"] == ["two"] + + def test_read_multi(self): + data = """ + Header: one + Header: two + \r\n + """ + data = textwrap.dedent(data) + data = data.strip() + s = StringIO(data) + headers = proxy.read_headers(s) + assert headers["header"] == ["one", "two"] + + def test_read_continued(self): + data = """ + Header: one + \ttwo + Header2: three + \r\n + """ + data = textwrap.dedent(data) + data = data.strip() + s = StringIO(data) + headers = proxy.read_headers(s) + assert headers["header"] == ['one\r\n two'] + + tests = [ uProxyError(), uFileLike(), u_parse_request_line(), u_read_chunked(), u_read_http_body(), + u_read_headers() ] diff --git a/test/tutils.py b/test/tutils.py index bed99caf..87977b1c 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.Headers() + headers = flow.ODict() 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.Headers() + headers = flow.ODict() headers["header_response"] = ["svalue"] return flow.Response(req, 200, "message", headers, "content_response") -- cgit v1.2.3