1"""SocksiPy - Python SOCKS module.
2
3Version 1.00
4
5Copyright 2006 Dan-Haim. All rights reserved.
6
7Redistribution and use in source and binary forms, with or without modification,
8are permitted provided that the following conditions are met:
91. Redistributions of source code must retain the above copyright notice, this
10   list of conditions and the following disclaimer.
112. Redistributions in binary form must reproduce the above copyright notice,
12   this list of conditions and the following disclaimer in the documentation
13   and/or other materials provided with the distribution.
143. Neither the name of Dan Haim nor the names of his contributors may be used
15   to endorse or promote products derived from this software without specific
16   prior written permission.
17
18THIS SOFTWARE IS PROVIDED BY DAN HAIM "AS IS" AND ANY EXPRESS OR IMPLIED
19WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
20MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO
21EVENT SHALL DAN HAIM OR HIS CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
22INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
23LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA
24OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
25LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
26OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMANGE.
27
28This module provides a standard socket-like interface for Python
29for tunneling connections through SOCKS proxies.
30
31Minor modifications made by Christopher Gilbert (http://motomastyle.com/) for
32use in PyLoris (http://pyloris.sourceforge.net/).
33
34Minor modifications made by Mario Vilas (http://breakingcode.wordpress.com/)
35mainly to merge bug fixes found in Sourceforge.
36"""
37
38import base64
39import socket
40import struct
41import sys
42
43if getattr(socket, "socket", None) is None:
44    raise ImportError("socket.socket missing, proxy support unusable")
45
46PROXY_TYPE_SOCKS4 = 1
47PROXY_TYPE_SOCKS5 = 2
48PROXY_TYPE_HTTP = 3
49PROXY_TYPE_HTTP_NO_TUNNEL = 4
50
51_defaultproxy = None
52_orgsocket = socket.socket
53
54
55class ProxyError(Exception):
56    pass
57
58
59class GeneralProxyError(ProxyError):
60    pass
61
62
63class Socks5AuthError(ProxyError):
64    pass
65
66
67class Socks5Error(ProxyError):
68    pass
69
70
71class Socks4Error(ProxyError):
72    pass
73
74
75class HTTPError(ProxyError):
76    pass
77
78
79_generalerrors = (
80    "success",
81    "invalid data",
82    "not connected",
83    "not available",
84    "bad proxy type",
85    "bad input",
86)
87
88_socks5errors = (
89    "succeeded",
90    "general SOCKS server failure",
91    "connection not allowed by ruleset",
92    "Network unreachable",
93    "Host unreachable",
94    "Connection refused",
95    "TTL expired",
96    "Command not supported",
97    "Address type not supported",
98    "Unknown error",
99)
100
101_socks5autherrors = (
102    "succeeded",
103    "authentication is required",
104    "all offered authentication methods were rejected",
105    "unknown username or invalid password",
106    "unknown error",
107)
108
109_socks4errors = (
110    "request granted",
111    "request rejected or failed",
112    "request rejected because SOCKS server cannot connect to identd on the client",
113    "request rejected because the client program and identd report different "
114    "user-ids",
115    "unknown error",
116)
117
118
119def setdefaultproxy(
120    proxytype=None, addr=None, port=None, rdns=True, username=None, password=None
121):
122    """setdefaultproxy(proxytype, addr[, port[, rdns[, username[, password]]]])
123    Sets a default proxy which all further socksocket objects will use,
124    unless explicitly changed.
125    """
126    global _defaultproxy
127    _defaultproxy = (proxytype, addr, port, rdns, username, password)
128
129
130def wrapmodule(module):
131    """wrapmodule(module)
132
133    Attempts to replace a module's socket library with a SOCKS socket. Must set
134    a default proxy using setdefaultproxy(...) first.
135    This will only work on modules that import socket directly into the
136    namespace;
137    most of the Python Standard Library falls into this category.
138    """
139    if _defaultproxy != None:
140        module.socket.socket = socksocket
141    else:
142        raise GeneralProxyError((4, "no proxy specified"))
143
144
145class socksocket(socket.socket):
146    """socksocket([family[, type[, proto]]]) -> socket object
147    Open a SOCKS enabled socket. The parameters are the same as
148    those of the standard socket init. In order for SOCKS to work,
149    you must specify family=AF_INET, type=SOCK_STREAM and proto=0.
150    """
151
152    def __init__(
153        self, family=socket.AF_INET, type=socket.SOCK_STREAM, proto=0, _sock=None
154    ):
155        _orgsocket.__init__(self, family, type, proto, _sock)
156        if _defaultproxy != None:
157            self.__proxy = _defaultproxy
158        else:
159            self.__proxy = (None, None, None, None, None, None)
160        self.__proxysockname = None
161        self.__proxypeername = None
162        self.__httptunnel = True
163
164    def __recvall(self, count):
165        """__recvall(count) -> data
166        Receive EXACTLY the number of bytes requested from the socket.
167        Blocks until the required number of bytes have been received.
168        """
169        data = self.recv(count)
170        while len(data) < count:
171            d = self.recv(count - len(data))
172            if not d:
173                raise GeneralProxyError((0, "connection closed unexpectedly"))
174            data = data + d
175        return data
176
177    def sendall(self, content, *args):
178        """ override socket.socket.sendall method to rewrite the header
179        for non-tunneling proxies if needed
180        """
181        if not self.__httptunnel:
182            content = self.__rewriteproxy(content)
183        return super(socksocket, self).sendall(content, *args)
184
185    def __rewriteproxy(self, header):
186        """ rewrite HTTP request headers to support non-tunneling proxies
187        (i.e. those which do not support the CONNECT method).
188        This only works for HTTP (not HTTPS) since HTTPS requires tunneling.
189        """
190        host, endpt = None, None
191        hdrs = header.split("\r\n")
192        for hdr in hdrs:
193            if hdr.lower().startswith("host:"):
194                host = hdr
195            elif hdr.lower().startswith("get") or hdr.lower().startswith("post"):
196                endpt = hdr
197        if host and endpt:
198            hdrs.remove(host)
199            hdrs.remove(endpt)
200            host = host.split(" ")[1]
201            endpt = endpt.split(" ")
202            if self.__proxy[4] != None and self.__proxy[5] != None:
203                hdrs.insert(0, self.__getauthheader())
204            hdrs.insert(0, "Host: %s" % host)
205            hdrs.insert(0, "%s http://%s%s %s" % (endpt[0], host, endpt[1], endpt[2]))
206        return "\r\n".join(hdrs)
207
208    def __getauthheader(self):
209        auth = self.__proxy[4] + ":" + self.__proxy[5]
210        return "Proxy-Authorization: Basic " + base64.b64encode(auth)
211
212    def setproxy(
213        self,
214        proxytype=None,
215        addr=None,
216        port=None,
217        rdns=True,
218        username=None,
219        password=None,
220        headers=None,
221    ):
222        """setproxy(proxytype, addr[, port[, rdns[, username[, password]]]])
223
224        Sets the proxy to be used.
225        proxytype -    The type of the proxy to be used. Three types
226                are supported: PROXY_TYPE_SOCKS4 (including socks4a),
227                PROXY_TYPE_SOCKS5 and PROXY_TYPE_HTTP
228        addr -        The address of the server (IP or DNS).
229        port -        The port of the server. Defaults to 1080 for SOCKS
230                servers and 8080 for HTTP proxy servers.
231        rdns -        Should DNS queries be preformed on the remote side
232                (rather than the local side). The default is True.
233                Note: This has no effect with SOCKS4 servers.
234        username -    Username to authenticate with to the server.
235                The default is no authentication.
236        password -    Password to authenticate with to the server.
237                Only relevant when username is also provided.
238        headers -     Additional or modified headers for the proxy connect
239        request.
240        """
241        self.__proxy = (
242            proxytype,
243            addr,
244            port,
245            rdns,
246            username.encode() if username else None,
247            password.encode() if password else None,
248            headers,
249        )
250
251    def __negotiatesocks5(self, destaddr, destport):
252        """__negotiatesocks5(self,destaddr,destport)
253        Negotiates a connection through a SOCKS5 server.
254        """
255        # First we'll send the authentication packages we support.
256        if (self.__proxy[4] != None) and (self.__proxy[5] != None):
257            # The username/password details were supplied to the
258            # setproxy method so we support the USERNAME/PASSWORD
259            # authentication (in addition to the standard none).
260            self.sendall(struct.pack("BBBB", 0x05, 0x02, 0x00, 0x02))
261        else:
262            # No username/password were entered, therefore we
263            # only support connections with no authentication.
264            self.sendall(struct.pack("BBB", 0x05, 0x01, 0x00))
265        # We'll receive the server's response to determine which
266        # method was selected
267        chosenauth = self.__recvall(2)
268        if chosenauth[0:1] != chr(0x05).encode():
269            self.close()
270            raise GeneralProxyError((1, _generalerrors[1]))
271        # Check the chosen authentication method
272        if chosenauth[1:2] == chr(0x00).encode():
273            # No authentication is required
274            pass
275        elif chosenauth[1:2] == chr(0x02).encode():
276            # Okay, we need to perform a basic username/password
277            # authentication.
278            self.sendall(
279                chr(0x01).encode()
280                + chr(len(self.__proxy[4]))
281                + self.__proxy[4]
282                + chr(len(self.__proxy[5]))
283                + self.__proxy[5]
284            )
285            authstat = self.__recvall(2)
286            if authstat[0:1] != chr(0x01).encode():
287                # Bad response
288                self.close()
289                raise GeneralProxyError((1, _generalerrors[1]))
290            if authstat[1:2] != chr(0x00).encode():
291                # Authentication failed
292                self.close()
293                raise Socks5AuthError((3, _socks5autherrors[3]))
294            # Authentication succeeded
295        else:
296            # Reaching here is always bad
297            self.close()
298            if chosenauth[1] == chr(0xFF).encode():
299                raise Socks5AuthError((2, _socks5autherrors[2]))
300            else:
301                raise GeneralProxyError((1, _generalerrors[1]))
302        # Now we can request the actual connection
303        req = struct.pack("BBB", 0x05, 0x01, 0x00)
304        # If the given destination address is an IP address, we'll
305        # use the IPv4 address request even if remote resolving was specified.
306        try:
307            ipaddr = socket.inet_aton(destaddr)
308            req = req + chr(0x01).encode() + ipaddr
309        except socket.error:
310            # Well it's not an IP number,  so it's probably a DNS name.
311            if self.__proxy[3]:
312                # Resolve remotely
313                ipaddr = None
314                req = (
315                    req
316                    + chr(0x03).encode()
317                    + chr(len(destaddr)).encode()
318                    + destaddr.encode()
319                )
320            else:
321                # Resolve locally
322                ipaddr = socket.inet_aton(socket.gethostbyname(destaddr))
323                req = req + chr(0x01).encode() + ipaddr
324        req = req + struct.pack(">H", destport)
325        self.sendall(req)
326        # Get the response
327        resp = self.__recvall(4)
328        if resp[0:1] != chr(0x05).encode():
329            self.close()
330            raise GeneralProxyError((1, _generalerrors[1]))
331        elif resp[1:2] != chr(0x00).encode():
332            # Connection failed
333            self.close()
334            if ord(resp[1:2]) <= 8:
335                raise Socks5Error((ord(resp[1:2]), _socks5errors[ord(resp[1:2])]))
336            else:
337                raise Socks5Error((9, _socks5errors[9]))
338        # Get the bound address/port
339        elif resp[3:4] == chr(0x01).encode():
340            boundaddr = self.__recvall(4)
341        elif resp[3:4] == chr(0x03).encode():
342            resp = resp + self.recv(1)
343            boundaddr = self.__recvall(ord(resp[4:5]))
344        else:
345            self.close()
346            raise GeneralProxyError((1, _generalerrors[1]))
347        boundport = struct.unpack(">H", self.__recvall(2))[0]
348        self.__proxysockname = (boundaddr, boundport)
349        if ipaddr != None:
350            self.__proxypeername = (socket.inet_ntoa(ipaddr), destport)
351        else:
352            self.__proxypeername = (destaddr, destport)
353
354    def getproxysockname(self):
355        """getsockname() -> address info
356        Returns the bound IP address and port number at the proxy.
357        """
358        return self.__proxysockname
359
360    def getproxypeername(self):
361        """getproxypeername() -> address info
362        Returns the IP and port number of the proxy.
363        """
364        return _orgsocket.getpeername(self)
365
366    def getpeername(self):
367        """getpeername() -> address info
368        Returns the IP address and port number of the destination
369        machine (note: getproxypeername returns the proxy)
370        """
371        return self.__proxypeername
372
373    def __negotiatesocks4(self, destaddr, destport):
374        """__negotiatesocks4(self,destaddr,destport)
375        Negotiates a connection through a SOCKS4 server.
376        """
377        # Check if the destination address provided is an IP address
378        rmtrslv = False
379        try:
380            ipaddr = socket.inet_aton(destaddr)
381        except socket.error:
382            # It's a DNS name. Check where it should be resolved.
383            if self.__proxy[3]:
384                ipaddr = struct.pack("BBBB", 0x00, 0x00, 0x00, 0x01)
385                rmtrslv = True
386            else:
387                ipaddr = socket.inet_aton(socket.gethostbyname(destaddr))
388        # Construct the request packet
389        req = struct.pack(">BBH", 0x04, 0x01, destport) + ipaddr
390        # The username parameter is considered userid for SOCKS4
391        if self.__proxy[4] != None:
392            req = req + self.__proxy[4]
393        req = req + chr(0x00).encode()
394        # DNS name if remote resolving is required
395        # NOTE: This is actually an extension to the SOCKS4 protocol
396        # called SOCKS4A and may not be supported in all cases.
397        if rmtrslv:
398            req = req + destaddr + chr(0x00).encode()
399        self.sendall(req)
400        # Get the response from the server
401        resp = self.__recvall(8)
402        if resp[0:1] != chr(0x00).encode():
403            # Bad data
404            self.close()
405            raise GeneralProxyError((1, _generalerrors[1]))
406        if resp[1:2] != chr(0x5A).encode():
407            # Server returned an error
408            self.close()
409            if ord(resp[1:2]) in (91, 92, 93):
410                self.close()
411                raise Socks4Error((ord(resp[1:2]), _socks4errors[ord(resp[1:2]) - 90]))
412            else:
413                raise Socks4Error((94, _socks4errors[4]))
414        # Get the bound address/port
415        self.__proxysockname = (
416            socket.inet_ntoa(resp[4:]),
417            struct.unpack(">H", resp[2:4])[0],
418        )
419        if rmtrslv != None:
420            self.__proxypeername = (socket.inet_ntoa(ipaddr), destport)
421        else:
422            self.__proxypeername = (destaddr, destport)
423
424    def __negotiatehttp(self, destaddr, destport):
425        """__negotiatehttp(self,destaddr,destport)
426        Negotiates a connection through an HTTP server.
427        """
428        # If we need to resolve locally, we do this now
429        if not self.__proxy[3]:
430            addr = socket.gethostbyname(destaddr)
431        else:
432            addr = destaddr
433        headers = ["CONNECT ", addr, ":", str(destport), " HTTP/1.1\r\n"]
434        wrote_host_header = False
435        wrote_auth_header = False
436        if self.__proxy[6] != None:
437            for key, val in self.__proxy[6].iteritems():
438                headers += [key, ": ", val, "\r\n"]
439                wrote_host_header = key.lower() == "host"
440                wrote_auth_header = key.lower() == "proxy-authorization"
441        if not wrote_host_header:
442            headers += ["Host: ", destaddr, "\r\n"]
443        if not wrote_auth_header:
444            if self.__proxy[4] != None and self.__proxy[5] != None:
445                headers += [self.__getauthheader(), "\r\n"]
446        headers.append("\r\n")
447        self.sendall("".join(headers).encode())
448        # We read the response until we get the string "\r\n\r\n"
449        resp = self.recv(1)
450        while resp.find("\r\n\r\n".encode()) == -1:
451            resp = resp + self.recv(1)
452        # We just need the first line to check if the connection
453        # was successful
454        statusline = resp.splitlines()[0].split(" ".encode(), 2)
455        if statusline[0] not in ("HTTP/1.0".encode(), "HTTP/1.1".encode()):
456            self.close()
457            raise GeneralProxyError((1, _generalerrors[1]))
458        try:
459            statuscode = int(statusline[1])
460        except ValueError:
461            self.close()
462            raise GeneralProxyError((1, _generalerrors[1]))
463        if statuscode != 200:
464            self.close()
465            raise HTTPError((statuscode, statusline[2]))
466        self.__proxysockname = ("0.0.0.0", 0)
467        self.__proxypeername = (addr, destport)
468
469    def connect(self, destpair):
470        """connect(self, despair)
471        Connects to the specified destination through a proxy.
472        destpar - A tuple of the IP/DNS address and the port number.
473        (identical to socket's connect).
474        To select the proxy server use setproxy().
475        """
476        # Do a minimal input check first
477        if (
478            (not type(destpair) in (list, tuple))
479            or (len(destpair) < 2)
480            or (not isinstance(destpair[0], basestring))
481            or (type(destpair[1]) != int)
482        ):
483            raise GeneralProxyError((5, _generalerrors[5]))
484        if self.__proxy[0] == PROXY_TYPE_SOCKS5:
485            if self.__proxy[2] != None:
486                portnum = self.__proxy[2]
487            else:
488                portnum = 1080
489            _orgsocket.connect(self, (self.__proxy[1], portnum))
490            self.__negotiatesocks5(destpair[0], destpair[1])
491        elif self.__proxy[0] == PROXY_TYPE_SOCKS4:
492            if self.__proxy[2] != None:
493                portnum = self.__proxy[2]
494            else:
495                portnum = 1080
496            _orgsocket.connect(self, (self.__proxy[1], portnum))
497            self.__negotiatesocks4(destpair[0], destpair[1])
498        elif self.__proxy[0] == PROXY_TYPE_HTTP:
499            if self.__proxy[2] != None:
500                portnum = self.__proxy[2]
501            else:
502                portnum = 8080
503            _orgsocket.connect(self, (self.__proxy[1], portnum))
504            self.__negotiatehttp(destpair[0], destpair[1])
505        elif self.__proxy[0] == PROXY_TYPE_HTTP_NO_TUNNEL:
506            if self.__proxy[2] != None:
507                portnum = self.__proxy[2]
508            else:
509                portnum = 8080
510            _orgsocket.connect(self, (self.__proxy[1], portnum))
511            if destpair[1] == 443:
512                self.__negotiatehttp(destpair[0], destpair[1])
513            else:
514                self.__httptunnel = False
515        elif self.__proxy[0] == None:
516            _orgsocket.connect(self, (destpair[0], destpair[1]))
517        else:
518            raise GeneralProxyError((4, _generalerrors[4]))
519