ProxyC/Builder/__init__.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import copy
from Basis.Logger import logging
from Basis.Process import Process
from Basis.Constant import WorkDir, PathEnv
from Basis.Functions import hostFormat, genFlag, getAvailablePort

from Builder import Brook
from Builder import VMess
from Builder import VLESS
from Builder import Trojan
from Builder import TrojanGo
from Builder import Hysteria
from Builder import Shadowsocks
from Builder import ShadowsocksR

clientEntry = {
    'ss': [Shadowsocks.load, '.json'],
    'ssr': [ShadowsocksR.load, '.json'],
    'vmess': [VMess.load, '.json'],
    'vless': [VLESS.load, '.json'],
    'trojan': [Trojan.load, '.json'],
    'trojan-go': [TrojanGo.load, '.json'],
    'brook': [Brook.load, ''],
    'hysteria': [Hysteria.load, '.json'],
}


class Builder(object):
    """ Build the proxy client process and expose socks5 port.

    Arguments:
      proxyType: Proxy node type.

      proxyInfo: Proxy node information.

      bindAddr: Socks5 proxy bind address.

      taskId: Task ID, defaults to 12 random characters length.

    Attributes:
        id, proxyType, proxyInfo, socksAddr, socksPort, output
    """
    output = None

    def __loadClient(self):
        logging.info('[%s] Builder load %s proxy client at %s -> %s' % (
            self.id, self.proxyType,
            'socks5://%s:%i' % (hostFormat(self.socksAddr, v6Bracket = True), self.socksPort), self.proxyInfo
        ))
        configFile = os.path.join(  # config file path
            WorkDir, self.id + clientEntry[self.proxyType][1]  # workDir + taskId + file suffix
        )
        logging.debug('[%s] Builder config file -> %s' % (self.id, configFile))
        command, fileContent, envVar = clientEntry[self.proxyType][0](self.proxyInfo, {  # load client boot info
            'addr': self.socksAddr,
            'port': self.socksPort,
        }, configFile)
        envVar['PATH'] = PathEnv  # add PATH env (some programs need it)
        self.__process = Process(WorkDir, taskId = self.id, cmd = command, env = envVar, file = {  # start process
            'path': configFile,
            'content': fileContent
        })

    def __init__(self, proxyType: str, proxyInfo: dict, bindAddr: str, taskId: str = '') -> None:  # init proxy client
        self.id = genFlag(length = 12) if taskId == '' else taskId  # load task ID
        if proxyType not in clientEntry:
            logging.error('[%s] Builder receive unknown proxy type %s' % (self.id, proxyType))
            raise RuntimeError('Unknown proxy type')
        self.proxyType = proxyType  # proxy type -> ss / ssr / vmess ...
        self.proxyInfo = copy.copy(proxyInfo)  # connection info
        self.socksAddr = bindAddr
        self.socksPort = getAvailablePort()  # random port for socks5 exposed
        self.__loadClient()

    def status(self) -> bool:  # check if the sub process is still running
        return self.__process.status()

    def destroy(self) -> None:  # kill sub process and remove config file
        logging.debug('[%s] Builder destroy' % self.id)
        self.__process.quit()
        self.output = self.__process.output