r'''Use a webshell client.
'''

from api.exploit import Exploit, Cmdline
from api.webshell import Webshell
from api import MainSession, config, load_module, colour, logger
import argparse
import re
import os

class exploit(Exploit):
    def __init__(self):
        self.parse = argparse.ArgumentParser(prog='use', description='Use a webshell client.')
        self.parse.add_argument('client', help="A webshell client.like `php/sample_one_word`", type=self._check)
        self.help = __doc__ + self.parse.format_help()
    
    def _check(self, param: str)-> str:
        test = r'[\w\-/]+'
        if re.fullmatch(test, param):
            return param
        
        raise TypeError(f'Webshell client like `{test}`')

    def support(self, session)-> bool:
        '''判断当前命令是否支持指定的session， 一般来说session即为运行改命令时的session
        '''
        self.session = session
        return isinstance(self.session, MainSession)

    def run(self, args: Cmdline)-> int:
        args = self.parse.parse_args(args.options)
        if args.client:
            webshell = self._getwebshell(args.client)
            if webshell is None:
                logger.error(f'No webshell client named `{args.client}`.')
                return self.STOP
            self.session['Prompt'] = lambda :f"{colour.colorize(config.app_name, 'underline')} webshell({colour.colorize(args.client, fore='red')}) {colour.colorize(config.prompt, 'bold', 'blue')} "
            webshell.type = args.client
            self.session.client = webshell
            logger.info(f"Use a webshell client `{args.client}`", True)

        return self.SUCCESS

    def _getwebshell(self, path: str)->Webshell:
        type = 'webshell'
        ml = path.split('/')
        ml.insert(0, type)
        tmp = self.session.modules
        for m in ml:
            if tmp is None:
                return None
            tmp = tmp.get(m)
        if tmp is not None:
            try:
                tmp = load_module(tmp.__name__, tmp.__file__) # 重载模块
                tmp = tmp.__dict__[type]() # 创建实例
                if not isinstance(tmp, Webshell):
                    raise Exception('Webshell client is not a Webshell instance, it is a wrong implementation!')
            except Exception as e:
                logger.error(e)
                return None
        return tmp

    def docomplete(self, text: str)-> dict:
        result = []
        pre = ''
        match = re.compile(r'^(use +)([\w\-/\\]*)$', re.M).search(text)
        if match:
            tmp = self.session.modules.get('webshell')
            k = ''
            s = ''
            for k in match.group(2).split('/'):
                if k not in tmp:
                    break
                tmp = tmp[k]
                s = os.path.join(s, k)
            pre = match.group(1)+os.path.join(s, 'a')[:-1]
            pre = pre.replace(os.sep, '/')
            if isinstance(tmp, dict):
                for key in tmp:
                    if key.startswith(k):
                        if isinstance(tmp[key], dict):
                            key += '/'
                        else:
                            key += ' '
                        result.append(key)
                if len(result) == 0 and s != '':
                    result.append('')

        return {'pre':pre, 'matchs':result}