r'''Create an alias for command.

'''
from api.exploit import Exploit, Cmdline, WebshellExploit
from api.webshell import WebshellSession
from api import logger, config, utils, colour, tablor
import argparse
import re

class AliasCommand(Exploit):
    '''
    __init__:
        name        指定命令的别名
        value       别名对应的实际命令行字符串
    '''
    def __init__(self, name: str, value: str, help: str, session):
        super().__init__()
        self.name = name
        self.value = value
        self.session = session
        self.help = help

    def get_origin_exploit(self):
        '''获取源命令对象
        '''
        cmdline = Cmdline(self.value)
        return self.session.command_map.get(cmdline.cmd)
    
    def run(self, args:Cmdline)-> int:
        args = args.cmdline.replace(self.name, self.value)
        return self.session.exec(args)

class exploit(WebshellExploit):
    def __init__(self):
        self.parse = argparse.ArgumentParser(prog='alias', description='Create alias for command.')
        self.parse.add_argument('-d', '--description', help="A description for alias.")
        self.parse.add_argument('alias', help="A alias for cmd", type=self._check)
        self.parse.add_argument('cmd', help="A command name or expression")
        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"Options must be fullmatch by regexp `{test}`.")

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

    def run(self, args: Cmdline)-> int:
        args = self.parse.parse_args(args.options)
        if self.session.has_cmd(args.alias) and utils.input("Alias ​​already exists!Do you want to overwrite it? (y/n) ").lower() != 'y':
            return self.STOP
        self._add(args.alias, args.cmd, args.description)
        logger.info(f"Add alias command: {args.alias} => {cmd}", True)
        return self.SUCCESS

    def _add(self, name: str, value: str, des: str=None):
        self.session.alias_command[name] = AliasCommand(name, value, des if des else f"{name} => {value}", self.session)

    def docomplete(self, text: str)-> dict: # alias命令补全
        result = []
        pre = ''
        match = re.compile(r'^(alias(?: -D .*?|-p .*?)? [\w\-]* )([\w\-]*)$', re.M).search(text)
        if match:
            pre = match.group(1)
            for cmd in self.session.command_map:
                if cmd.startswith(match.group(2).lower()):
                    result.append(cmd+' ')

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

    def hook_loaded(self):
        # 常用别名定义
        self._add('lls', f"lrun {'dir' if 'win' in config.platform else 'ls'}", "List local files")
        if isinstance(self.session, WebshellSession):
            self._add('bg', "sessions -i 0", "Background current session and switch to main session.")

        # 从配置文件中恢复保存的别名
        alias_list = self.load_config('alias-config')
        if alias_list is not None:
            for a in alias_list:
                self._add(*a)

    def hook_destroy(self):
        if isinstance(self.session, WebshellSession):
            alias_list = []
            for name, cmd in self.session.alias_command.items():
                alias_list.append((name, cmd.value, cmd.help))
            self.save_config('alias-config', alias_list)
