Add a size() method to flow.Request and flow.Response.

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):
         """

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()