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