Add socks_socket kwarg to TorClientEndpoint's constructor

txtorcon/endpoints.py

@@ -32,6 +32,7 @@
 from twisted.internet.endpoints import serverFromString
 from twisted.internet.endpoints import clientFromString
 from twisted.internet.endpoints import TCP4ClientEndpoint
+from twisted.internet.endpoints import UNIXClientEndpoint
 from twisted.internet import error
 from twisted.plugin import IPlugin
 from twisted.python.util import FancyEqMixin
@@ -667,6 +668,7 @@ class TorClientEndpoint(object):
     socks_ports_to_try = [9050, 9150]
 
     def __init__(self, host, port,
+                 socks_socket=None,
                  socks_hostname=None, socks_port=None,
                  socks_username=None, socks_password=None,
                  _proxy_endpoint_generator=default_tcp4_endpoint_generator):
@@ -676,6 +678,7 @@ def __init__(self, host, port,
         self.host = host
         self.port = int(port)
         self._proxy_endpoint_generator = _proxy_endpoint_generator
+        self.socks_socket = socks_socket
         self.socks_hostname = socks_hostname
         self.socks_port = int(socks_port) if socks_port is not None else None
         self.socks_username = socks_username
@@ -691,31 +694,39 @@ def __init__(self, host, port,
     @defer.inlineCallbacks
     def connect(self, protocolfactory):
         last_error = None
-        for socks_port in self._socks_port_iter:
-            self.socks_port = socks_port
-            tor_ep = self._proxy_endpoint_generator(
-                reactor,
-                self.socks_hostname,
-                self.socks_port,
-            )
-
+        if self.socks_socket is not None:
+            tor_ep = UNIXClientEndpoint(reactor, self.socks_socket)
             args = (self.host, self.port, tor_ep)
             kwargs = dict()
-            if self.socks_username is not None and self.socks_password is not None:
-                kwargs['methods'] = dict(
-                    login=(self.socks_username, self.socks_password),
+            socks_ep = SOCKS5ClientEndpoint(*args, **kwargs)
+            proto = yield socks_ep.connect(protocolfactory)
+            defer.returnValue(proto)
+        else:
+            for socks_port in self._socks_port_iter:
+                self.socks_port = socks_port
+                tor_ep = self._proxy_endpoint_generator(
+                    reactor,
+                    self.socks_hostname,
+                    self.socks_port,
                 )
 
-            socks_ep = SOCKS5ClientEndpoint(*args, **kwargs)
+                args = (self.host, self.port, tor_ep)
+                kwargs = dict()
+                if self.socks_username is not None and self.socks_password is not None:
+                    kwargs['methods'] = dict(
+                        login=(self.socks_username, self.socks_password),
+                    )
 
-            try:
-                proto = yield socks_ep.connect(protocolfactory)
-                defer.returnValue(proto)
+                socks_ep = SOCKS5ClientEndpoint(*args, **kwargs)
+
+                try:
+                    proto = yield socks_ep.connect(protocolfactory)
+                    defer.returnValue(proto)
 
-            except error.ConnectError as e0:
-                last_error = e0
-        if last_error is not None:
-            raise last_error
+                except error.ConnectError as e0:
+                    last_error = e0
+            if last_error is not None:
+                raise last_error
 
 
 @implementer(IPlugin, IStreamClientEndpointStringParserWithReactor)