#!/usr/bin/env python

"""SocksiPy - Python SOCKS module.
Version 1.00

Copyright 2006 Dan-Haim. All rights reserved.

Redistribution and use in source and binary forms, with or without modification,
are permitted provided that the following conditions are met:
1. Redistributions of source code must retain the above copyright notice, this
   list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright notice,
   this list of conditions and the following disclaimer in the documentation
   and/or other materials provided with the distribution.
3. Neither the name of Dan Haim nor the names of his contributors may be used
   to endorse or promote products derived from this software without specific
   prior written permission.

THIS SOFTWARE IS PROVIDED BY DAN HAIM "AS IS" AND ANY EXPRESS OR IMPLIED
WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO
EVENT SHALL DAN HAIM OR HIS CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA
OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMANGE.


This module provides a standard socket-like interface for Python
for tunneling connections through SOCKS proxies.

"""

"""
Minor modifications made by Miroslav Stampar (https://sqlmap.org/)
for patching DNS-leakage occuring in socket.create_connection()

Minor modifications made by Christopher Gilbert (http://motomastyle.com/)
for use in PyLoris (http://pyloris.sourceforge.net/)

Minor modifications made by Mario Vilas (http://breakingcode.wordpress.com/)
mainly to merge bug fixes found in Sourceforge

"""

import socket
import struct

PROXY_TYPE_SOCKS4 = 1
PROXY_TYPE_SOCKS5 = 2
PROXY_TYPE_HTTP = 3

_defaultproxy = None
socket._orig_socket = _orgsocket = _orig_socket = socket.socket
_orgcreateconnection = socket.create_connection

class ProxyError(Exception): pass
class GeneralProxyError(ProxyError): pass
class Socks5AuthError(ProxyError): pass
class Socks5Error(ProxyError): pass
class Socks4Error(ProxyError): pass
class HTTPError(ProxyError): pass

_generalerrors = ("success",
    "invalid data",
    "not connected",
    "not available",
    "bad proxy type",
    "bad input")

_socks5errors = ("succeeded",
    "general SOCKS server failure",
    "connection not allowed by ruleset",
    "Network unreachable",
    "Host unreachable",
    "Connection refused",
    "TTL expired",
    "Command not supported",
    "Address type not supported",
    "Unknown error")

_socks5autherrors = ("succeeded",
    "authentication is required",
    "all offered authentication methods were rejected",
    "unknown username or invalid password",
    "unknown error")

_socks4errors = ("request granted",
    "request rejected or failed",
    "request rejected because SOCKS server cannot connect to identd on the client",
    "request rejected because the client program and identd report different user-ids",
    "unknown error")

def setdefaultproxy(proxytype=None, addr=None, port=None, rdns=True, username=None, password=None):
    """setdefaultproxy(proxytype, addr[, port[, rdns[, username[, password]]]])
    Sets a default proxy which all further socksocket objects will use,
    unless explicitly changed.
    """
    global _defaultproxy
    _defaultproxy = (proxytype, addr, port, rdns, username, password)

def wrapmodule(module):
    """wrapmodule(module)
    Attempts to replace a module's socket library with a SOCKS socket. Must set
    a default proxy using setdefaultproxy(...) first.
    This will only work on modules that import socket directly into the namespace;
    most of the Python Standard Library falls into this category.
    """
    if _defaultproxy != None:
        module.socket.socket = socksocket
        if _defaultproxy[0] == PROXY_TYPE_SOCKS4:
            # Note: unable to prevent DNS leakage in SOCKS4 (Reference: https://security.stackexchange.com/a/171280)
            pass
        else:
            module.socket.create_connection = create_connection
    else:
        raise GeneralProxyError((4, "no proxy specified"))

def unwrapmodule(module):
    module.socket.socket = _orgsocket
    module.socket.create_connection = _orgcreateconnection


class socksocket(socket.socket):
    """socksocket([family[, type[, proto]]]) -> socket object
    Open a SOCKS enabled socket. The parameters are the same as
    those of the standard socket init. In order for SOCKS to work,
    you must specify family=AF_INET, type=SOCK_STREAM and proto=0.
    """

    def __init__(self, family=socket.AF_INET, type=socket.SOCK_STREAM, proto=0, _sock=None):
        # 初始化父类 socket
        _orgsocket.__init__(self, family, type, proto, _sock)

        # 根据默认代理设置代理
        if _defaultproxy != None:
            self.__proxy = _defaultproxy
        else:
            # 初始化代理信息，如果没有默认代理则设为空元组
            self.__proxy = (None, None, None, None, None, None)

        # 用于保存代理关联的本地socket信息
        self.__proxysockname = None
        self.__proxypeername = None

    def __recvall(self, count):
        """__recvall(count) -> data
        Receive EXACTLY the number of bytes requested from the socket.
        Blocks until the required number of bytes have been received.
        """
        # 从socket接收指定字节数数据
        data = self.recv(count)
        # 当接收的数据少于请求的字节数时继续接收
        while len(data) < count:
            d = self.recv(count - len(data))
            if not d:
                raise GeneralProxyError((0, "connection closed unexpectedly"))  # 如果连接关闭则抛出异常
            data = data + d  # 累加接收到的数据
        return data

    def setproxy(self, proxytype=None, addr=None, port=None, rdns=True, username=None, password=None):
        """setproxy(proxytype, addr[, port[, rdns[, username[, password]]]])
        Sets the proxy to be used.
        """
        # 设置代理类型、地址、端口等参数
        self.__proxy = (proxytype, addr, port, rdns, username, password)

    def __negotiatesocks5(self, destaddr, destport):
        """__negotiatesocks5(self,destaddr,destport)
        Negotiates a connection through a SOCKS5 server.
        """
        # 发送支持的认证方法
        if (self.__proxy[4] != None) and (self.__proxy[5] != None):
            # 如果提供了用户名和密码,则支持用户名/密码认证
            self.sendall(struct.pack('BBBB', 0x05, 0x02, 0x00, 0x02))  # 0x02 代表用户名/密码认证
        else:
            # 仅支持无认证连接
            self.sendall(struct.pack('BBB', 0x05, 0x01, 0x00))  # 0x01 代表无认证

        # 接收服务器的响应以确定选定的身份验证方法
        chosenauth = self.__recvall(2)
        if chosenauth[0:1] != b'\x05':
            self.close()
            raise GeneralProxyError((1, _generalerrors[1]))  # 无效的协议版本

        # 检查所选的身份验证方法
        if chosenauth[1:2] == b'\x00':
            # 无需身份验证
            pass
        elif chosenauth[1:2] == b'\x02':
            # 需要进行用户名/密码身份验证
            self.sendall(b'\x01' + chr(len(self.__proxy[4])).encode() + self.__proxy[4].encode() + chr(
                len(self.__proxy[5])).encode() + self.__proxy[5].encode())
            authstat = self.__recvall(2)  # 接收身份验证状态
            if authstat[0:1] != b'\x01':
                self.close()
                raise GeneralProxyError((1, _generalerrors[1]))  # 无效的认证响应
            if authstat[1:2] != b'\x00':
                self.close()
                raise Socks5AuthError((3, _socks5autherrors[3]))  # 身份验证失败
        else:
            # 非法的身份验证方法
            self.close()
            if chosenauth[1:2] == b'\xff':
                raise Socks5AuthError((2, _socks5autherrors[2]))  # 不支持的身份验证
            else:
                raise GeneralProxyError((1, _generalerrors[1]))

        # 请求实际连接
        req = struct.pack('BBB', 0x05, 0x01, 0x00)  # 0x01 表示连接请求
        try:
            # 尝试将目标地址作为IP地址处理
            ipaddr = socket.inet_aton(destaddr)
            req = req + b'\x01' + ipaddr  # IPv4地址请求
        except socket.error:
            # 目标地址不是IP地址，可能是DNS名称
            if self.__proxy[3]:  # 检查是否需要远程解析
                ipaddr = None
                req = req + chr(0x03).encode() + chr(len(destaddr)).encode() + (
                    destaddr if isinstance(destaddr, bytes) else destaddr.encode())
            else:
                # 本地解析
                ipaddr = socket.inet_aton(socket.gethostbyname(destaddr))
                req = req + chr(0x01).encode() + ipaddr
        req = req + struct.pack(">H", destport)  # 添加目标端口
        self.sendall(req)  # 发送请求

        # 获取响应
        resp = self.__recvall(4)
        if resp[0:1] != chr(0x05).encode():
            self.close()
            raise GeneralProxyError((1, _generalerrors[1]))  # 协议错误
        elif resp[1:2] != chr(0x00).encode():
            # 连接失败
            self.close()
            if ord(resp[1:2]) <= 8:
                raise Socks5Error((ord(resp[1:2]), _socks5errors[ord(resp[1:2])]))
            else:
                raise Socks5Error((9, _socks5errors[9]))

        # 解析绑定地址/端口
        elif resp[3:4] == chr(0x01).encode():
            boundaddr = self.__recvall(4)  # IPv4绑定地址
        elif resp[3:4] == chr(0x03).encode():
            resp = resp + self.recv(1)  # 读取后续字节
            boundaddr = self.__recvall(ord(resp[4:5]))  # DNS名
        else:
            self.close()
            raise GeneralProxyError((1, _generalerrors[1]))

        # 获取绑定端口
        boundport = struct.unpack(">H", self.__recvall(2))[0]
        self.__proxysockname = (boundaddr, boundport)  # 保存绑定的socket信息
        if ipaddr != None:
            self.__proxypeername = (socket.inet_ntoa(ipaddr), destport)  # 保存peer名称
        else:
            self.__proxypeername = (destaddr, destport)

    def getproxysockname(self):
        """getsockname() -> address info
        Returns the bound IP address and port number at the proxy.
        """
        return self.__proxysockname  # 返回代理的socket信息

    def getproxypeername(self):
        """getproxypeername() -> address info
        Returns the IP and port number of the proxy.
        """
        return _orgsocket.getpeername(self)  # 返回代理的peer信息

    def getpeername(self):
        """getpeername() -> address info
        Returns the IP address and port number of the destination
        machine (note: getproxypeername returns the proxy)
        """
        return self.__proxypeername  # 返回目标机器的peer名称

    def __negotiatesocks4(self, destaddr, destport):
        """__negotiatesocks4(self,destaddr,destport)
        Negotiates a connection through a SOCKS4 server.
        """
        # 检查目标地址是否为IP地址
        rmtrslv = False
        try:
            ipaddr = socket.inet_aton(destaddr)
        except socket.error:
            # 目标是DNS名称, 检查是否应进行远程解析
            if self.__proxy[3]:
                ipaddr = struct.pack("BBBB", 0x00, 0x00, 0x00, 0x01)  # 默认IP地址
                rmtrslv = True
            else:
                ipaddr = socket.inet_aton(socket.gethostbyname(destaddr))  # 本地解析

        # 构造请求报文
        req = struct.pack(">BBH", 0x04, 0x01, destport) + ipaddr  # SOCKS4 请求
        # 用户名参数视为userid
        if self.__proxy[4] != None:
            req = req + self.__proxy[4]
        req = req + chr(0x00).encode()  # 请求结束

        # 需要远程解析的情况下使用DNS名称
        if rmtrslv:
            req = req + destaddr + chr(0x00).encode()  # SOCKS4A扩展
        self.sendall(req)  # 发送请求

        # 获取服务器响应
        resp = self.__recvall(8)
        if resp[0:1] != chr(0x00).encode():
            # 出现错误
            self.close()
            raise GeneralProxyError((1, _generalerrors[1]))  # 无效数据
        if resp[1:2] != chr(0x5A).encode():
            # 服务器返回错误
            self.close()
            if ord(resp[1:2]) in (91, 92, 93):
                self.close()
                raise Socks4Error((ord(resp[1:2]), _socks4errors[ord(resp[1:2]) - 90]))  # 处理错误代码
            else:
                raise Socks4Error((94, _socks4errors[4]))  # 一般性错误

        # 获取绑定的地址和端口
        self.__proxysockname = (socket.inet_ntoa(resp[4:]), struct.unpack(">H", resp[2:4])[0])  # 保存绑定信息
        if rmtrslv != None:
            self.__proxypeername = (socket.inet_ntoa(ipaddr), destport)  # 目标地址
        else:
            self.__proxypeername = (destaddr, destport)  # 目标地址

    def __negotiatehttp(self, destaddr, destport):
        """__negotiatehttp(self,destaddr,destport)
        Negotiates a connection through an HTTP server.
        """
        # 本地解析
        if not self.__proxy[3]:
            addr = socket.gethostbyname(destaddr)
        else:
            addr = destaddr

        # 发送HTTP的CONNECT请求
        self.sendall(
            ("CONNECT " + addr + ":" + str(destport) + " HTTP/1.1\r\n" + "Host: " + destaddr + "\r\n\r\n").encode())

        # 读取响应直到遇到"\r\n\r\n"
        resp = self.recv(1)
        while resp.find("\r\n\r\n".encode()) == -1:
            resp = resp + self.recv(1)

        # 仅需第一行检查连接是否成功
        statusline = resp.splitlines()[0].split(" ".encode(), 2)
        if statusline[0] not in ("HTTP/1.0".encode(), "HTTP/1.1".encode()):
            self.close()
            raise GeneralProxyError((1, _generalerrors[1]))  # 无效的HTTP协议
        try:
            statuscode = int(statusline[1])  # 获取状态码
        except ValueError:
            self.close()
            raise GeneralProxyError((1, _generalerrors[1]))  # 状态码无效
        if statuscode != 200:
            self.close()
            raise HTTPError((statuscode, statusline[2]))  # 显示状态码错误

        # 连接成功后的代理信息
        self.__proxysockname = ("0.0.0.0", 0)  # 绑定在0.0.0.0
        self.__proxypeername = (addr, destport)  # 保存peer信息

    def connect(self, destpair):
        """connect(self, destpair)
        Connects to the specified destination through a proxy.
        destpair - A tuple of the IP/DNS address and the port number.
        (identical to socket's connect).
        To select the proxy server use setproxy().
        """
        # 校验输入参数
        if (not type(destpair) in (list, tuple)) or (len(destpair) < 2) or (type(destpair[0]) != type('')) or (
                type(destpair[1]) != int):
            raise GeneralProxyError((5, _generalerrors[5]))

        # 根据代理类型选择相应的握手协议
        if self.__proxy[0] == PROXY_TYPE_SOCKS5:
            if self.__proxy[2] != None:
                portnum = self.__proxy[2]
            else:
                portnum = 1080  # 默认端口

            _orgsocket.connect(self, (self.__proxy[1], portnum))  # 连接到代理
            self.__negotiatesocks5(destpair[0], destpair[1])  # 握手

        elif self.__proxy[0] == PROXY_TYPE_SOCKS4:
            if self.__proxy[2] != None:
                portnum = self.__proxy[2]
            else:
                portnum = 1080  # 默认端口

            _orgsocket.connect(self, (self.__proxy[1], portnum))  # 连接到代理
            self.__negotiatesocks4(destpair[0], destpair[1])  # 握手

        elif self.__proxy[0] == PROXY_TYPE_HTTP:
            if self.__proxy[2] != None:
                portnum = self.__proxy[2]
            else:
                portnum = 8080  # 默认端口

            _orgsocket.connect(self, (self.__proxy[1], portnum))  # 连接到代理
            self.__negotiatehttp(destpair[0], destpair[1])  # 握手

        elif self.__proxy[0] == None:
            # 没有代理时直接连接目标
            _orgsocket.connect(self, (destpair[0], destpair[1]))
        else:
            raise GeneralProxyError((4, _generalerrors[4]))  # 不支持的代理类型

def create_connection(address, timeout=socket._GLOBAL_DEFAULT_TIMEOUT,
                    source_address=None):
    # Patched for a DNS-leakage
    host, port = address
    sock = None
    try:
        sock = socksocket(socket.AF_INET, socket.SOCK_STREAM)
        if timeout is not socket._GLOBAL_DEFAULT_TIMEOUT:
            sock.settimeout(timeout)
        if source_address:
            sock.bind(source_address)
        sock.connect(address)
    except socket.error:
        if sock is not None:
            sock.close()
        raise
    return sock
