bitbot-3.11-fork/modules/proxy.py

import typing, urllib.parse
import socks
from src import ModuleManager, utils

TYPES = {
    "socks4": socks.SOCKS4,
    "socks5": socks.SOCKS5,
    "http": socks.HTTP
}

class ProxySetting(utils.Setting):
    def parse(self, value: str) -> typing.Any:
        parsed = urllib.parse.urlparse(value)
        if parsed.scheme in TYPES and parsed.hostname:
            return value

@utils.export("serverset", ProxySetting("proxy",
    "Proxy configuration for the current server",
    example="socks5://localhost:9050"))
class Module(ModuleManager.BaseModule):
    @utils.hook("preprocess.connect")
    def new_server(self, event):
        proxy = event["server"].get_setting("proxy", None)
        if proxy:
            proxy_parsed = urllib.parse.urlparse(proxy)
            type = TYPES.get(proxy_parsed.scheme)

            if type == None:
                raise ValueError("Invalid proxy type '%s' for '%s'" %
                    (proxy_parsed.scheme, str(event["server"])))

            event["server"].socket._make_socket = self._socket_factory(
                type, proxy_parsed.hostname, proxy_parsed.port)

    def _socket_factory(self, ptype, phost, pport):
        return lambda host, port, bind, timeout: self._make_socket(
            ptype, phost, pport, host, port, bind, timeout)
    def _make_socket(self, ptype, phost, pport, host, port, bind, timeout):
        return socks.create_connection((host, port), timeout, bind,
            ptype, phost, pport)