Add SOCKS5 mode, fix #305

1 files changed, 75 insertions, 6 deletions

diff --git a/libmproxy/proxy/primitives.py b/libmproxy/proxy/primitives.py
index 23d089d3..c0ae424d 100644
--- a/libmproxy/proxy/primitives.py
+++ b/libmproxy/proxy/primitives.py
@@ -1,5 +1,5 @@
 from __future__ import absolute_import
-
+from netlib import socks
 
 class ProxyError(Exception):
     def __init__(self, code, message, headers=None):
@@ -15,7 +15,7 @@ class ProxyMode(object):
     http_form_in = None
     http_form_out = None
 
-    def get_upstream_server(self, conn):
+    def get_upstream_server(self, client_conn):
         """
         Returns the address of the server to connect to.
         Returns None if the address needs to be determined on the protocol level (regular proxy mode)
@@ -46,7 +46,7 @@ class RegularProxyMode(ProxyMode):
     http_form_in = "absolute"
     http_form_out = "relative"
 
-    def get_upstream_server(self, conn):
+    def get_upstream_server(self, client_conn):
         return None
 
 
@@ -58,9 +58,9 @@ class TransparentProxyMode(ProxyMode):
         self.resolver = resolver
         self.sslports = sslports
 
-    def get_upstream_server(self, conn):
+    def get_upstream_server(self, client_conn):
         try:
-            dst = self.resolver.original_addr(conn)
+            dst = self.resolver.original_addr(client_conn.connection)
         except Exception, e:
             raise ProxyError(502, "Transparent mode failure: %s" % str(e))
 
@@ -71,11 +71,80 @@ class TransparentProxyMode(ProxyMode):
         return [ssl, ssl] + list(dst)
 
 
+class Socks5ProxyMode(ProxyMode):
+    http_form_in = "relative"
+    http_form_out = "relative"
+
+    def __init__(self, sslports):
+        self.sslports = sslports
+
+    @staticmethod
+    def _assert_socks5(msg):
+        if msg.ver != socks.VERSION.SOCKS5:
+            if msg.ver == ord("G") and len(msg.methods) == ord("E"):
+                guess = "Probably not a SOCKS request but a regular HTTP request. "
+            else:
+                guess = ""
+            raise socks.SocksError(
+                socks.REP.GENERAL_SOCKS_SERVER_FAILURE,
+                guess + "Invalid SOCKS version. Expected 0x05, got 0x%x" % msg.ver)
+
+    def get_upstream_server(self, client_conn):
+        try:
+            # Parse Client Greeting
+            client_greet = socks.ClientGreeting.from_file(client_conn.rfile)
+            self._assert_socks5(client_greet)
+            if socks.METHOD.NO_AUTHENTICATION_REQUIRED not in client_greet.methods:
+                raise socks.SocksError(
+                    socks.METHOD.NO_ACCEPTABLE_METHODS,
+                    "mitmproxy only supports SOCKS without authentication"
+                )
+
+            # Send Server Greeting
+            server_greet = socks.ServerGreeting(
+                socks.VERSION.SOCKS5,
+                socks.METHOD.NO_AUTHENTICATION_REQUIRED
+            )
+            server_greet.to_file(client_conn.wfile)
+            client_conn.wfile.flush()
+
+            # Parse Connect Request
+            connect_request = socks.Message.from_file(client_conn.rfile)
+            self._assert_socks5(connect_request)
+            if connect_request.msg != socks.CMD.CONNECT:
+                raise socks.SocksError(
+                    socks.REP.COMMAND_NOT_SUPPORTED,
+                    "mitmproxy only supports SOCKS5 CONNECT."
+                )
+
+            # We do not connect here yet, as the clientconnect event has not been handled yet.
+
+            connect_reply = socks.Message(
+                socks.VERSION.SOCKS5,
+                socks.REP.SUCCEEDED,
+                socks.ATYP.DOMAINNAME,
+                client_conn.address  # dummy value, we don't have an upstream connection yet.
+            )
+            connect_reply.to_file(client_conn.wfile)
+            client_conn.wfile.flush()
+
+            ssl = bool(connect_request.addr.port in self.sslports)
+            return ssl, ssl, connect_request.addr.host, connect_request.addr.port
+
+        except socks.SocksError as e:
+            msg = socks.Message(5, e.code, socks.ATYP.DOMAINNAME, repr(e))
+            try:
+                msg.to_file(client_conn.wfile)
+            except:
+                pass
+            raise ProxyError(502, "SOCKS5 mode failure: %s" % str(e))
+
+
 class _ConstDestinationProxyMode(ProxyMode):
     def __init__(self, dst):
         self.dst = dst
 
-    def get_upstream_server(self, conn):
+    def get_upstream_server(self, client_conn):
         return self.dst