Merge pull request #675 from Kriechi/protocol-refactor

HTTP protocol refactoring

7 files changed, 101 insertions, 75 deletions

diff --git a/libmproxy/console/grideditor.py b/libmproxy/console/grideditor.py
index c93ecabf..d32ce5b4 100644
--- a/libmproxy/console/grideditor.py
+++ b/libmproxy/console/grideditor.py
@@ -5,9 +5,11 @@ import re
 import os
 import urwid
 
+from netlib import odict
+from netlib.http import user_agents
+
 from . import common, signals
 from .. import utils, filt, script
-from netlib import http_uastrings, http_cookies, odict
 
 
 FOOTER = [
@@ -516,7 +518,7 @@ class HeaderEditor(GridEditor):
         return text
 
     def set_user_agent(self, k):
-        ua = http_uastrings.get_by_shortcut(k)
+        ua = user_agents.get_by_shortcut(k)
         if ua:
             self.walker.add_value(
                 [
@@ -529,7 +531,7 @@ class HeaderEditor(GridEditor):
         if key == "U":
             signals.status_prompt_onekey.send(
                 prompt = "Add User-Agent header:",
-                keys = [(i[0], i[1]) for i in http_uastrings.UASTRINGS],
+                keys = [(i[0], i[1]) for i in user_agents.UASTRINGS],
                 callback = self.set_user_agent,
             )
             return True
@@ -592,7 +594,7 @@ class SetHeadersEditor(GridEditor):
         return text
 
     def set_user_agent(self, k):
-        ua = http_uastrings.get_by_shortcut(k)
+        ua = user_agents.get_by_shortcut(k)
         if ua:
             self.walker.add_value(
                 [
@@ -606,7 +608,7 @@ class SetHeadersEditor(GridEditor):
         if key == "U":
             signals.status_prompt_onekey.send(
                 prompt = "Add User-Agent header:",
-                keys = [(i[0], i[1]) for i in http_uastrings.UASTRINGS],
+                keys = [(i[0], i[1]) for i in user_agents.UASTRINGS],
                 callback = self.set_user_agent,
             )
             return True
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 1bc73d91..cbb22c3a 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -158,7 +158,7 @@ class StreamLargeBodies(object):
     def run(self, flow, is_request):
         r = flow.request if is_request else flow.response
         code = flow.response.code if flow.response else None
-        expected_size = netlib.http.expected_http_body_size(
+        expected_size = netlib.http.http1.HTTP1Protocol.expected_http_body_size(
             r.headers, is_request, flow.request.method, code
         )
         if not (0 <= expected_size <= self.max_size):
diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py
index da15066b..818b8811 100644
--- a/libmproxy/protocol/http.py
+++ b/libmproxy/protocol/http.py
@@ -1,14 +1,16 @@
 from __future__ import absolute_import
 import Cookie
+import copy
+import threading
+import time
 import urllib
 import urlparse
-import time
-import copy
 from email.utils import parsedate_tz, formatdate, mktime_tz
-import threading
-from netlib import http, tcp, http_status, http_cookies
-import netlib.utils
-from netlib import odict
+
+import netlib
+from netlib import http, tcp, odict, utils
+from netlib.http import cookies
+
 from .tcp import TCPHandler
 from .primitives import KILL, ProtocolHandler, Flow, Error
 from ..proxy.connection import ServerConnection
@@ -303,6 +305,10 @@ class HTTPRequest(HTTPMessage):
         is_replay=bool
     )
 
+    @property
+    def body(self):
+        return self.content
+
     @classmethod
     def from_state(cls, state):
         f = cls(
@@ -354,11 +360,10 @@ class HTTPRequest(HTTPMessage):
         if hasattr(rfile, "reset_timestamps"):
             rfile.reset_timestamps()
 
-        req = http.read_request(
-            rfile,
+        protocol = http.http1.HTTP1Protocol(rfile=rfile, wfile=wfile)
+        req = protocol.read_request(
             include_body = include_body,
             body_size_limit = body_size_limit,
-            wfile = wfile
         )
 
         if hasattr(rfile, "first_byte_timestamp"):
@@ -375,7 +380,7 @@ class HTTPRequest(HTTPMessage):
             req.path,
             req.httpversion,
             req.headers,
-            req.content,
+            req.body,
             timestamp_start,
             timestamp_end
         )
@@ -642,7 +647,7 @@ class HTTPRequest(HTTPMessage):
         """
         ret = odict.ODict()
         for i in self.headers["cookie"]:
-            ret.extend(http_cookies.parse_cookie_header(i))
+            ret.extend(cookies.parse_cookie_header(i))
         return ret
 
     def set_cookies(self, odict):
@@ -650,7 +655,7 @@ class HTTPRequest(HTTPMessage):
             Takes an netlib.odict.ODict object. Over-writes any existing Cookie
             headers.
         """
-        v = http_cookies.format_cookie_header(odict)
+        v = cookies.format_cookie_header(odict)
         self.headers["Cookie"] = [v]
 
     def replace(self, pattern, repl, *args, **kwargs):
@@ -724,6 +729,12 @@ class HTTPResponse(HTTPMessage):
         msg=str
     )
 
+
+    @property
+    def body(self):
+        return self.content
+
+
     @classmethod
     def from_state(cls, state):
         f = cls(None, None, None, None, None)
@@ -760,11 +771,12 @@ class HTTPResponse(HTTPMessage):
         if hasattr(rfile, "reset_timestamps"):
             rfile.reset_timestamps()
 
-        httpversion, code, msg, headers, content = http.read_response(
-            rfile,
+        protocol = http.http1.HTTP1Protocol(rfile=rfile)
+        resp = protocol.read_response(
             request_method,
             body_size_limit,
-            include_body=include_body)
+            include_body=include_body
+        )
 
         if hasattr(rfile, "first_byte_timestamp"):
             # more accurate timestamp_start
@@ -776,11 +788,11 @@ class HTTPResponse(HTTPMessage):
             timestamp_end = None
 
         return HTTPResponse(
-            httpversion,
-            code,
-            msg,
-            headers,
-            content,
+            resp.httpversion,
+            resp.status_code,
+            resp.msg,
+            resp.headers,
+            resp.body,
             timestamp_start,
             timestamp_end
         )
@@ -894,7 +906,7 @@ class HTTPResponse(HTTPMessage):
         """
         ret = []
         for header in self.headers["set-cookie"]:
-            v = http_cookies.parse_set_cookie_header(header)
+            v = http.cookies.parse_set_cookie_header(header)
             if v:
                 name, value, attrs = v
                 ret.append([name, [value, attrs]])
@@ -910,7 +922,7 @@ class HTTPResponse(HTTPMessage):
         values = []
         for i in odict.lst:
             values.append(
-                http_cookies.format_set_cookie_header(
+                http.cookies.format_set_cookie_header(
                     i[0],
                     i[1][0],
                     i[1][1]
@@ -1044,7 +1056,8 @@ class HTTPHandler(ProtocolHandler):
                 self.c.server_conn.send(request_raw)
                 # Only get the headers at first...
                 flow.response = HTTPResponse.from_stream(
-                    self.c.server_conn.rfile, flow.request.method,
+                    self.c.server_conn.rfile,
+                    flow.request.method,
                     body_size_limit=self.c.config.body_size_limit,
                     include_body=False
                 )
@@ -1081,10 +1094,13 @@ class HTTPHandler(ProtocolHandler):
             if flow.response.stream:
                 flow.response.content = CONTENT_MISSING
             else:
-                flow.response.content = http.read_http_body(
-                    self.c.server_conn.rfile, flow.response.headers,
+                protocol = http.http1.HTTP1Protocol(rfile=self.c.server_conn.rfile)
+                flow.response.content = protocol.read_http_body(
+                    flow.response.headers,
                     self.c.config.body_size_limit,
-                    flow.request.method, flow.response.code, False
+                    flow.request.method,
+                    flow.response.code,
+                    False
                 )
         flow.response.timestamp_end = utils.timestamp()
 
@@ -1231,7 +1247,7 @@ class HTTPHandler(ProtocolHandler):
             pass
 
     def send_error(self, code, message, headers):
-        response = http_status.RESPONSES.get(code, "Unknown")
+        response = http.status_codes.RESPONSES.get(code, "Unknown")
         html_content = """
             <html>
                 <head>
@@ -1285,6 +1301,7 @@ class HTTPHandler(ProtocolHandler):
             if not request.host and flow.server_conn:
                 request.host, request.port = flow.server_conn.address.host, flow.server_conn.address.port
 
+
         # Now we can process the request.
         if request.form_in == "authority":
             if self.c.client_conn.ssl_established:
@@ -1364,7 +1381,7 @@ class HTTPHandler(ProtocolHandler):
         # We provide a mostly unified API to the user, which needs to be
         # unfiddled here
         # ( See also: https://github.com/mitmproxy/mitmproxy/issues/337 )
-        address = netlib.tcp.Address((flow.request.host, flow.request.port))
+        address = tcp.Address((flow.request.host, flow.request.port))
 
         ssl = (flow.request.scheme == "https")
 
@@ -1418,8 +1435,8 @@ class HTTPHandler(ProtocolHandler):
             h = flow.response._assemble_head(preserve_transfer_encoding=True)
             self.c.client_conn.send(h)
 
-            chunks = http.read_http_body_chunked(
-                self.c.server_conn.rfile,
+            protocol = http.http1.HTTP1Protocol(rfile=self.c.server_conn.rfile)
+            chunks = protocol.read_http_body_chunked(
                 flow.response.headers,
                 self.c.config.body_size_limit,
                 flow.request.method,
@@ -1441,15 +1458,18 @@ class HTTPHandler(ProtocolHandler):
             semantics. Returns True, if so.
         """
         close_connection = (
-            http.connection_close(
+            http.http1.HTTP1Protocol.connection_close(
                 flow.request.httpversion,
-                flow.request.headers) or http.connection_close(
+                flow.request.headers
+            ) or http.http1.HTTP1Protocol.connection_close(
                 flow.response.httpversion,
-                flow.response.headers) or http.expected_http_body_size(
+                flow.response.headers
+            ) or http.http1.HTTP1Protocol.expected_http_body_size(
                 flow.response.headers,
                 False,
                 flow.request.method,
-                flow.response.code) == -1)
+                flow.response.code) == -1
+            )
         if close_connection:
             if flow.request.form_in == "authority" and flow.response.code == 200:
                 # Workaround for
diff --git a/libmproxy/proxy/config.py b/libmproxy/proxy/config.py
index c5306b4a..ec91a6e0 100644
--- a/libmproxy/proxy/config.py
+++ b/libmproxy/proxy/config.py
@@ -2,7 +2,11 @@ from __future__ import absolute_import
 import os
 import re
 from OpenSSL import SSL
-from netlib import http_auth, certutils, tcp
+
+import netlib
+from netlib import http, certutils, tcp
+from netlib.http import authentication
+
 from .. import utils, platform, version
 from .primitives import RegularProxyMode, SpoofMode, SSLSpoofMode, TransparentProxyMode, UpstreamProxyMode, ReverseProxyMode, Socks5ProxyMode
 
@@ -103,7 +107,7 @@ class ProxyConfig:
             self.openssl_method_server = ssl_version_server
         else:
             self.openssl_method_server = tcp.SSL_VERSIONS[ssl_version_server]
-        
+
         if ssl_verify_upstream_cert:
             self.openssl_verification_mode_server = SSL.VERIFY_PEER
         else:
@@ -164,18 +168,18 @@ def process_proxy_options(parser, options):
                 return parser.error(
                     "Invalid single-user specification. Please use the format username:password")
             username, password = options.auth_singleuser.split(':')
-            password_manager = http_auth.PassManSingleUser(username, password)
+            password_manager = authentication.PassManSingleUser(username, password)
         elif options.auth_nonanonymous:
-            password_manager = http_auth.PassManNonAnon()
+            password_manager = authentication.PassManNonAnon()
         elif options.auth_htpasswd:
             try:
-                password_manager = http_auth.PassManHtpasswd(
+                password_manager = authentication.PassManHtpasswd(
                     options.auth_htpasswd)
             except ValueError as v:
                 return parser.error(v.message)
-        authenticator = http_auth.BasicProxyAuth(password_manager, "mitmproxy")
+        authenticator = authentication.BasicProxyAuth(password_manager, "mitmproxy")
     else:
-        authenticator = http_auth.NullProxyAuth(None)
+        authenticator = authentication.NullProxyAuth(None)
 
     certs = []
     for i in options.certs:
diff --git a/test/test_protocol_http.py b/test/test_protocol_http.py
index d8489d4d..747fdc1e 100644
--- a/test/test_protocol_http.py
+++ b/test/test_protocol_http.py
@@ -327,11 +327,11 @@ class TestInvalidRequests(tservers.HTTPProxTest):
         p = self.pathoc()
         r = p.request("connect:'%s:%s'" % ("127.0.0.1", self.server2.port))
         assert r.status_code == 400
-        assert "Must not CONNECT on already encrypted connection" in r.content
+        assert "Must not CONNECT on already encrypted connection" in r.body
 
     def test_relative_request(self):
         p = self.pathoc_raw()
         p.connect()
         r = p.request("get:/p/200")
         assert r.status_code == 400
-        assert "Invalid HTTP request form" in r.content
+        assert "Invalid HTTP request form" in r.body
diff --git a/test/test_proxy.py b/test/test_proxy.py
index 77051edd..01fbe953 100644
--- a/test/test_proxy.py
+++ b/test/test_proxy.py
@@ -31,7 +31,9 @@ class TestServerConnection:
         f.server_conn = sc
         f.request.path = "/p/200:da"
         sc.send(f.request.assemble())
-        assert http.read_response(sc.rfile, f.request.method, 1000)
+
+        protocol = http.http1.HTTP1Protocol(rfile=sc.rfile)
+        assert protocol.read_response(f.request.method, 1000)
         assert self.d.last_log()
 
         sc.finish()
diff --git a/test/test_server.py b/test/test_server.py
index 2805e936..10822ae4 100644
--- a/test/test_server.py
+++ b/test/test_server.py
@@ -1,15 +1,17 @@
 import socket
 import time
-from libmproxy.proxy.config import HostMatcher
-import libpathod
-from netlib import tcp, http_auth, http, socks
-from libpathod import pathoc, pathod
+from OpenSSL import SSL
+
+from netlib import tcp, http, socks
 from netlib.certutils import SSLCert
-import tutils
-import tservers
+from netlib.http import authentication
+from libpathod import pathoc, pathod
+
+from libmproxy.proxy.config import HostMatcher
 from libmproxy.protocol import KILL, Error
 from libmproxy.protocol.http import CONTENT_MISSING
-from OpenSSL import SSL
+import tutils
+import tservers
 
 """
     Note that the choice of response code in these tests matters more than you
@@ -182,6 +184,9 @@ class TcpMixin:
 class AppMixin:
     def test_app(self):
         ret = self.app("/")
+        print(ret)
+        print(ret.status_code)
+        print(ret.content)
         assert ret.status_code == 200
         assert "mitmproxy" in ret.content
 
@@ -295,8 +300,8 @@ class TestHTTP(tservers.HTTPProxTest, CommonMixin, AppMixin):
 
 
 class TestHTTPAuth(tservers.HTTPProxTest):
-    authenticator = http_auth.BasicProxyAuth(
-        http_auth.PassManSingleUser(
+    authenticator = http.authentication.BasicProxyAuth(
+        http.authentication.PassManSingleUser(
             "test",
             "test"),
         "realm")
@@ -310,8 +315,8 @@ class TestHTTPAuth(tservers.HTTPProxTest):
             h'%s'='%s'
         """ % (
             self.server.port,
-            http_auth.BasicProxyAuth.AUTH_HEADER,
-            http.assemble_http_basic_auth("basic", "test", "test")
+            http.authentication.BasicProxyAuth.AUTH_HEADER,
+            authentication.assemble_http_basic_auth("basic", "test", "test")
         ))
         assert ret.status_code == 202
 
@@ -526,7 +531,7 @@ class TestHttps2Http(tservers.ReverseProxTest):
         """
             Returns a connected Pathoc instance.
         """
-        p = libpathod.pathoc.Pathoc(
+        p = pathoc.Pathoc(
             ("localhost", self.proxy.port), ssl=ssl, sni=sni, fp=None
         )
         p.connect()
@@ -765,22 +770,15 @@ class TestStreamRequest(tservers.HTTPProxTest):
             (self.server.urlbase, spec))
         connection.send("\r\n")
 
-        httpversion, code, msg, headers, content = http.read_response(
-            fconn, "GET", None, include_body=False)
+        protocol = http.http1.HTTP1Protocol(rfile=fconn)
+        resp = protocol.read_response("GET", None, include_body=False)
 
-        assert headers["Transfer-Encoding"][0] == 'chunked'
-        assert code == 200
+        assert resp.headers["Transfer-Encoding"][0] == 'chunked'
+        assert resp.status_code == 200
 
         chunks = list(
-            content for _,
-                        content,
-                        _ in http.read_http_body_chunked(
-                fconn,
-                headers,
-                None,
-                "GET",
-                200,
-                False))
+            content for _, content, _ in protocol.read_http_body_chunked(
+                resp.headers, None, "GET", 200, False))
         assert chunks == ["this", "isatest", ""]
 
         connection.close()