r'''Manage saved webshell connections or create a connection.
'''

from api.exploit import Exploit, Cmdline
from api.webshell import WebshellSession, Webshell
from api import MainSession, config, logger, Session, tablor, colour, utils
import argparse
import re
import xml.etree.ElementTree as ET
import os

class exploit(Exploit):

    manager = None # 保存主session的实例

    def __init__(self):
        self.parse = argparse.ArgumentParser(prog='exploit', description='Manage saved webshell connections or active session info.', 
            epilog="Create a webshell connection from current webshell client, if no option set or `nosave` option set only.\n\
                When the session does not have the `nosave` mark, it will be automatically saved when the session is deleted or the program exits.")
        self.parse.add_argument('-n', '--nosave', help="Exploit but not save session,when session is build successfully.And set `nosave` mark on created session.", action='store_true')
        self.parse.add_argument('-l', '--list', help="List all saved webshell connections.", action='store_true')
        self.parse.add_argument('-d', '--delete', nargs='+', help="Provide a ID and delete the saved webshell connections", type=int)
        self.parse.add_argument('-c', '--create', help="Create a connection from saved webshell connections by id.", type=int)
        self.parse.add_argument('-s', '--save', nargs='*', type=int, help="Save the webshell client information corresponding to the session id.")
        self.help = __doc__ + self.parse.format_help()


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

        return True

    def _save(self, session:Session)-> int:
        '''保存session对应的webshell客户端信息
        '''
        path = os.path.join(os.path.dirname(__file__), "connections_example.xml")
        webshell = session.client
        if webshell is None:
            logger.error(f"No webshell client is using for session `{'manager' if isinstance(session, MainSession) else session.id}`.")
            return self.STOP
        if os.path.exists(config.webshell_save_path):
            path = config.webshell_save_path
            
        try:
            tree = ET.parse(path)
            root = tree.getroot()
            ID = session.state.get('save_id')# 如果session创建时添加了-n选项，则不会有save_id字段（指定保存连接的ID）
            item = root.find(f"./item[@id='{ID}']")
            if item is None:
                ID = 0
                for c in root:
                    i = int(c.attrib.get('id', 0))
                    if ID < i:
                        ID = i
                ID += 1
                item = ET.Element("item", {'id':str(ID)})
                root.append(item)
            
            w = item.find("./webshell")
            if w is None:
                w = ET.Element('webshell', {'value':webshell.type})
                item.append(w)
            else:
                w.set('value', webshell.type)
            
            old_options = item.find("./options")
            if old_options:
                item.remove(old_options)
            options = ET.Element('options')
            for o in webshell.options.option_list()[1:]:
                attr = {'name':o[0], 'value':str(o[1])}
                if o[1] is None: # 如果参数值为None则加一个NoneType属性
                    attr['NoneType'] = 'true'
                option = ET.Element('option', attr)
                options.append(option)
            item.append(options)

            tree.write(config.webshell_save_path)
            session.state['save_id'] = ID # 保存成功后添加save_id字段
            logger.info(f"Save a webshell connection `{ID}`.", True)
        except ET.ParseError:
            logger.error(f"Xml parse error on file `{path}`!")
            return self.STOP

        return self.SUCCESS


    def _create_from_client(self, session:Session, nosave: bool, save_id:int)-> int:
        '''从指定session中的webshell client建立连接
        '''
        webshell = session.client
        if webshell is None:
            logger.error('You must use a webshell client first.')
            return self.STOP
        if len(webshell.options.unset_list()) > 0: # 检测选项设置
            logger.error("There are some unset options!")
            return self.STOP
        # explot 当前目标
        if webshell.connect():
            ID = 1 if len(self.manager.session_map) == 0 else max(self.manager.session_map)+1
            s = WebshellSession(webshell, ID)
            self.manager.session_map[ID] = s
            if save_id:
                s.state['save_id'] = save_id # 表示该session是从保存的连接建立的
            if isinstance(session, MainSession):
                session.client = None
                session['Prompt'] = lambda :f"{colour.colorize(config.app_name, 'underline')} {colour.colorize(config.prompt, fore='blue')} "
            logger.info(f'A session is created, id is `{ID}`', True)

            if not nosave:# 存储建立的连接
                self._save(s)
            else:
                s.state['nosave'] = True # 给session增加一个不保存的标志

            webshell.hook_start(s)
            s['Prompt'] = lambda :colour.colorize(s.server_info.domain+'\\'+s.server_info.user, 'bold', 'cyan')+'@'+\
                colour.colorize(s.state['name'], fore='purple')+\
                ':'+f"({colour.colorize(s.pwd, fore='red')})"+' \n'+colour.colorize(config.prompt, 'bold', 'blue')+' '
            self.manager.loadexploit_to_session(self.manager.modules.get('exploit', {}), 'exploit', s) # 载入命令
            # 全部命令加载成功后调用命令回调
            s._hook_loaded()
            webshell.hook_loaded()
            self.manager.exec(["sessions", '-i', str(ID)])
            return self.SUCCESS

        logger.warning("Session create failed!")
        return self.STOP

    def _create_from_file(self, ID: int, nosave: bool)-> int:
        '''从保存的连接中创建连接
        '''
        try:
            tree = ET.parse(config.webshell_save_path)
            root = tree.getroot()
            item = None
            for c in root:
                if int(c.attrib.get('id', 0)) == ID:
                    item = c
                    break
            if item is None:
                logger.error(f"No this webshell connection for ID `{ID}`!")
                return self.STOP
            w = item.find("./webshell").get('value', '')
            if self.manager.exec(['use', w]) == self.SUCCESS:
                for o in item.findall("./options/option"):
                    if o.attrib.get("NoneType"):
                        continue
                    if self.manager.exec(['set', o.attrib.get('name'), o.attrib.get('value')]) != self.SUCCESS:
                        logger.warning(f"Set option `{o.attrib.get('name')}` faild with value `{o.attrib.get('value')}`!This will use its default value.")
                    
                return self._create_from_client(self.manager, nosave, ID)
        except ET.ParseError:
            logger.error(f"Xml parse error!Check this file `{config.webshell_save_path}`")
        except FileNotFoundError:
            logger.error(f"This file `{config.webshell_save_path}` not found!")

        return self.STOP

    def _list(self)-> int:
        table = [['ID', 'Webshell Type', 'Target']]
        try:
            tree = ET.parse(config.webshell_save_path)
            root = tree.getroot()
            for c in root:
                l = []
                l.append(c.attrib.get('id', 1))
                l.append(c.find("./webshell").attrib.get("value"))
                l.append(c.find("./options/option[@name='target']").attrib.get('value'))
                table.append(l)
        except (FileNotFoundError, ET.ParseError):
            logger.warning(f"Load or parse file `{config.webshell_save_path}` failed!")
            return self.STOP

        print(tablor(table, border=False, title="Saved Webshell Connections"))
        return self.SUCCESS

    def _delete(self, ID: int, force: bool)-> int:
        try:
            tree = ET.parse(config.webshell_save_path)
            root = tree.getroot()
            item = None
            for c in root:
                if int(c.attrib.get('id', 0)) == ID:
                    item = c
                    break
            if item is None:
                logger.error(f"No saved webshell connection for ID `{ID}`!")
                return self.STOP
            if force or utils.input("Are you sure to delete this webshell connection? (y/n) ").lower() == 'y':
                root.remove(item)
                tree.write(config.webshell_save_path)
                logger.info(f"Remove a saved webshell connections `{ID}`", False)
                return self.SUCCESS
        except ET.ParseError:
            logger.error(f"Xml parse error!Check this file `{config.webshell_save_path}`")
        except FileNotFoundError:
            logger.error(f"This file `{config.webshell_save_path}` not found!")

        return self.STOP

    def run(self, args: Cmdline)-> int:
        args = self.parse.parse_args(args.options)
        ret = self.SUCCESS
        if args.create is not None:
            ret = self._create_from_file(args.create, args.nosave)
        elif args.list:
            ret = self._list()
        elif args.delete:
            ids = ', '.join([str(i) for i in args.delete])
            if utils.input(f"Are you sure to delete these webshell connection `{ids}`? (y/n) ").lower() == 'y':
                for i in args.delete:
                    self._delete(i, True)
        elif args.save is not None:
            if len(args.save) == 0:
                if self.session.client is None:
                    logger.error("No webshell client is using.")
                    return self.STOP
                else:
                    ret = self._save(self.session)
            else:
                for i in args.save:
                    if i in self.manager.session_map:
                        self._save(self.manager.session_map.get(i))
        else:
            if isinstance(self.session, MainSession):
                ret = self._create_from_client(self.session, args.nosave, None)
            else:
                logger.error("Current session is a webshellsession!If you want to get a copy, it is recommended to create it from a file")
                ret = self.STOP
        return ret

    def hook_destroy(self):
        if 'nosave' not in self.session.state and isinstance(self.session, WebshellSession):
            self._save(self.session)