r'''Run a code on server.

该命令会调用当前webshell客户端的eval方法执行代码


'''

from api.exploit import WebshellExploit, Cmdline
from api.webshell import WebshellSession
from api import Session, logger
import argparse
import os
import re

class exploit(WebshellExploit):
    
    def __init__(self):
        super().__init__()
        self.parse = argparse.ArgumentParser(prog='eval', description="Run a dynamic script code.")
        self.parse.add_argument('code', help="A dynamic script code", nargs='?', default=None)
        self.parse.add_argument('-f', '--file', help="Specify a local file to execute on the server (the file must use the language code supported by webshell).")
        self.parse.add_argument('-o', '--option', nargs='*', help="Pass global variables to the script.Like `-o pwd=/tmp id=123`.Only supports strings.")
        self.help = __doc__+self.parse.format_help()

    def run(self, args: Cmdline)-> int:
        args = self.parse.parse_args(args.options)
        ret = self.SUCCESS
        kw = [] if args.option is None else args.option
        if args.code:
            ret = self._eval_code(args.code, kw)
        elif args.file:
            ret = self._eval_file(args.file, kw)
        else:
            print(self.parse.format_help())
            return self.STOP
            
        return ret

    def _eval_code(self, code: str, kw: list)-> int:
        d = {}
        for w in kw:
            l = w.split('=')
            k = l[0]
            v = '='.join(l[1:])
            d[k] = v
        ret = self.eval(code, **d)
        if not ret.is_success():
            logger.error('Eval code failed!')
            return self.STOP
        ret = ret.data
        if ret[-2:] == '\r\n':
            ret = ret[:-2]
        elif ret[-1:] == '\n':
            ret = ret[:-1]
        print(ret)
        return self.SUCCESS

    def _eval_file(self, fname: str, kw: dict)-> int:
        try:
            code = ''
            with open(fname, 'rb') as f:
                code = f.read().decode(errors='ignore')
            return self._eval_code(code, kw)
        except FileNotFoundError:
            logger.error(f"File read failed, check file `{fname}`")

        return self.STOP
        
    def docomplete(self, text: str)-> dict:# 本地文件路径补全
        result = []
        pre = ''
        match = re.compile(r'''^(eval +-f +)([\w\-/\\.]*)$''', re.M).search(text)
        root = os.getcwd()
        if match:
            dirname, name = os.path.split(match.group(2))
            path = os.path.join(root, dirname)
            pre = match.group(1)+os.path.join(dirname, 'a')[:-1]
            if os.path.exists(path) and os.path.isdir(path):
                for filename in os.listdir(path):
                    if os.path.isdir(os.path.join(path, filename)):
                        filename += os.sep
                    else:
                        filename += ' '
                    if filename.startswith(name):
                        result.append(filename)

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