diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 20020e842..309281778 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -174,6 +174,18 @@ class HTTPMsg(controller.Msg): self.content = encoding.encode(e, self.content) self.headers["content-encoding"] = [e] + def size(self, **kwargs): + """ + Size in bytes of a fully rendered message, including headers and + HTTP lead-in. + """ + hl = len(self._assemble_head(**kwargs)) + if self.content: + return hl + len(self.content) + else: + return hl + + class Request(HTTPMsg): """ @@ -374,17 +386,9 @@ class Request(HTTPMsg): self.scheme, self.host, self.port, self.path = parts return True - def _assemble(self, _proxy = False): - """ - Assembles the request for transmission to the server. We make some - modifications to make sure interception works properly. - - Returns None if the request cannot be assembled. - """ - if self.content == CONTENT_MISSING: - return None - FMT = '%s %s HTTP/%s.%s\r\n%s\r\n%s' - FMT_PROXY = '%s %s://%s:%s%s HTTP/%s.%s\r\n%s\r\n%s' + def _assemble_head(self, proxy=False): + FMT = '%s %s HTTP/%s.%s\r\n%s\r\n' + FMT_PROXY = '%s %s://%s:%s%s HTTP/%s.%s\r\n%s\r\n' headers = self.headers.copy() utils.del_all( @@ -405,14 +409,13 @@ class Request(HTTPMsg): content = "" if self.close: headers["connection"] = ["close"] - if not _proxy: + if not proxy: return FMT % ( self.method, self.path, self.httpversion[0], self.httpversion[1], - str(headers), - content + str(headers) ) else: return FMT_PROXY % ( @@ -423,10 +426,24 @@ class Request(HTTPMsg): self.path, self.httpversion[0], self.httpversion[1], - str(headers), - content + str(headers) ) + def _assemble(self, _proxy = False): + """ + Assembles the request for transmission to the server. We make some + modifications to make sure interception works properly. + + Returns None if the request cannot be assembled. + """ + if self.content == CONTENT_MISSING: + return None + head = self._assemble_head(_proxy) + if self.content: + return head + self.content + else: + return head + def replace(self, pattern, repl, *args, **kwargs): """ Replaces a regular expression pattern with repl in both the headers @@ -576,6 +593,19 @@ class Response(HTTPMsg): c.headers = self.headers.copy() return c + def _assemble_head(self): + FMT = '%s\r\n%s\r\n' + headers = self.headers.copy() + utils.del_all( + headers, + ['proxy-connection', 'transfer-encoding'] + ) + if self.content: + headers["content-length"] = [str(len(self.content))] + proto = "HTTP/%s.%s %s %s"%(self.httpversion[0], self.httpversion[1], self.code, str(self.msg)) + data = (proto, str(headers)) + return FMT%data + def _assemble(self): """ Assembles the response for transmission to the client. We make some @@ -585,20 +615,11 @@ class Response(HTTPMsg): """ if self.content == CONTENT_MISSING: return None - FMT = '%s\r\n%s\r\n%s' - headers = self.headers.copy() - utils.del_all( - headers, - ['proxy-connection', 'transfer-encoding'] - ) - content = self.content - if content: - headers["content-length"] = [str(len(content))] + head = self._assemble_head() + if self.content: + return head + self.content else: - content = "" - proto = "HTTP/%s.%s %s %s"%(self.httpversion[0], self.httpversion[1], self.code, str(self.msg)) - data = (proto, str(headers), content) - return FMT%data + return head def replace(self, pattern, repl, *args, **kwargs): """ diff --git a/test/test_flow.py b/test/test_flow.py index a2051eae5..eccd11f4f 100644 --- a/test/test_flow.py +++ b/test/test_flow.py @@ -724,12 +724,14 @@ class TestRequest: assert not r.set_url("") assert r.get_url() == u assert r._assemble() + assert r.size() == len(r._assemble()) r2 = r.copy() assert r == r2 r.content = None assert r._assemble() + assert r.size() == len(r._assemble()) r.close = True assert "connection: close" in r._assemble() @@ -739,6 +741,8 @@ class TestRequest: r.content = flow.CONTENT_MISSING assert not r._assemble() + + def test_getset_form_urlencoded(self): h = flow.ODictCaseless() h["content-type"] = [flow.HDR_FORM_URLENCODED] @@ -860,12 +864,15 @@ class TestResponse: req = flow.Request(c, (1, 1), "host", 22, "https", "GET", "/", h, "content") resp = flow.Response(req, (1, 1), 200, "msg", h.copy(), "content", None) assert resp._assemble() + assert resp.size() == len(resp._assemble()) + resp2 = resp.copy() assert resp2 == resp resp.content = None assert resp._assemble() + assert resp.size() == len(resp._assemble()) resp.content = flow.CONTENT_MISSING assert not resp._assemble()