from asyncore import file_dispatcher
from imp import reload
from pstats import Stats
from proxy import Filter, Proxy
import random, string, os, threading, sqlite3, time, atexit, socket
from kthread import KThread
from base64 import b64decode

LOCALHOST_IP = socket.gethostbyname(os.getenv("LOCALHOST_IP","127.0.0.1"))

class SQLite():
    def __init__(self, db_name) -> None:
        self.conn = None
        self.cur = None
        self.db_name = db_name
        self.lock = threading.Lock()

    def connect(self) -> None:
        if not os.path.exists("db"): os.mkdir("db")
        try:
            self.conn = sqlite3.connect("db/" + self.db_name + '.db', check_same_thread = False)
        except Exception:
            with open("db/" + self.db_name + '.db', 'x') as f:
                pass
            self.conn = sqlite3.connect("db/" + self.db_name + '.db', check_same_thread = False)

    def disconnect(self) -> None:
        self.conn.close()

    def check_integrity(self, tables = {}) -> None:
        cur = self.conn.cursor()
        for t in tables:
            cur.execute('''
                SELECT name FROM sqlite_master WHERE type='table' AND name='{}';
            '''.format(t))

            if len(cur.fetchall()) == 0:
                cur.execute('''CREATE TABLE main.{}({});'''.format(t, ''.join([(c + ' ' + tables[t][c] + ', ') for c in tables[t]])[:-2]))
        cur.close()
    
    def query(self, query, values = ()):
        cur = self.conn.cursor()
        try:
            with self.lock:
                cur.execute(query, values)
                return cur.fetchall()
        finally:
            cur.close()
            try: self.conn.commit()
            except Exception: pass

class KeyValueStorage:
    def __init__(self, db):
        self.db = db

    def get(self, key):
        q = self.db.query('SELECT value FROM keys_values WHERE key = ?', (key,))
        if len(q) == 0:
            return None
        else:
            return q[0][0]

    def put(self, key, value):
        if self.get(key) is None:
            self.db.query('INSERT INTO keys_values (key, value) VALUES (?, ?);', (key,str(value)))
        else:
            self.db.query('UPDATE keys_values SET value=? WHERE key = ?;', (str(value), key))

class STATUS:
    WAIT = "wait"
    STOP = "stop"
    PAUSE = "pause"
    ACTIVE = "active"
 
class ProxyManager:
    def __init__(self, db:SQLite):
        self.db = db
        self.proxy_table = {}
        self.lock = threading.Lock()
        atexit.register(self.clear)

    def __clear_proxy_table(self):
        with self.lock:
            for key in list(self.proxy_table.keys()):
                if not self.proxy_table[key]["thread"].is_alive():
                    del self.proxy_table[key]

    def clear(self):
        with self.lock:
            for key in list(self.proxy_table.keys()):
                if self.proxy_table[key]["thread"].is_alive():
                    self.proxy_table[key]["thread"].kill()
                del self.proxy_table[key]

    def reload(self):
        self.__clear_proxy_table()
        with self.lock: 
            for srv_id in self.db.query('SELECT service_id, status FROM services;'):
                srv_id, n_status = srv_id
                if srv_id in self.proxy_table:
                    continue
                update_signal = threading.Event()
                req_status = [n_status]
                thread = KThread(target=self.service_manager, args=(srv_id, req_status, update_signal))
                self.proxy_table[srv_id] = {
                    "thread":thread,
                    "event":update_signal,
                    "next_status":req_status
                }
                thread.start()

    def get_service_data(self, id):
        q = self.db.query('SELECT * FROM services WHERE service_id=?;',(id,))
        if len(q) == 0: return None
        srv = q[0]
        filters = [{
            'id': row[5],
            'regex': row[0],
            'is_blacklist': True if row[3] == "1" else False,
            'is_case_sensitive' : True if row[6] == "1" else False,
            'mode': row[1],
            'n_packets': row[4],
        } for row in self.db.query('SELECT * FROM regexes WHERE service_id = ?;', (id,))]
        return {
            'id': srv[1],
            'status': srv[0],
            'public_port': srv[3],
            'internal_port': srv[2],
            'filters':filters
        }

    def change_status(self, id, to):
        with self.lock:
            if id in self.proxy_table:
                if self.proxy_table[id]["thread"].is_alive():
                    self.proxy_table[id]["next_status"][0] = to
                    self.proxy_table[id]["event"].set()
                else:
                    del self.proxy_table[id]
    
    def fire_update(self, id):
        with self.lock:
            if id in self.proxy_table:
                if self.proxy_table[id]["thread"].is_alive():
                    self.proxy_table[id]["event"].set()
                else:
                    del self.proxy_table[id]
    
    def __update_status_db(self, id, status):
        self.db.query("UPDATE services SET status = ? WHERE service_id = ?;", (status, id))

    def __proxy_starter(self, id, proxy:Proxy, next_status, saved_status):
        def func():
            while True:
                if check_port_is_open(proxy.public_port):
                    self.__update_status_db(id, next_status)
                    if saved_status[0] == "wait": saved_status[0] = next_status
                    proxy_status = proxy.start(in_pause=(next_status==STATUS.PAUSE))
                    if proxy_status != 0:
                        self.__update_status_db(id, STATUS.STOP)
                    return
                else:
                    time.sleep(.5)
                    
        thread = KThread(target=func)
        thread.start()
        return thread

    def service_manager(self, id, next_status, signal:threading.Event):
    
        proxy = None
        thr_starter:KThread = None
        previous_status = STATUS.STOP
        filters = {}

        while True:
            restart_required = False
            reload_required = False
            
            data = self.get_service_data(id)

            #Close thread
            if data is None:
                if proxy and proxy.isactive():
                    proxy.stop()
                return
               
            #Filter check
            old_filters = set(filters.keys())
            new_filters = set([f["id"] for f in data["filters"]])

            #remove old filters
            for f in old_filters:
                if not f in new_filters:
                    reload_required = True
                    del filters[f]
            
            for f in new_filters:
                if not f in old_filters:
                    reload_required = True
                    filter_info = [ele for ele in data['filters'] if ele["id"] == f][0]
                    filters[f] = Filter(
                        is_case_sensitive=filter_info["is_case_sensitive"],
                        c_to_s=filter_info["mode"] in ["C","B"],
                        s_to_c=filter_info["mode"] in ["S","B"],
                        is_blacklist=filter_info["is_blacklist"],
                        regex=b64decode(filter_info["regex"]),
                        blocked_packets=filter_info["n_packets"],
                        code=f
                    )
            

            def stats_updater(filter:Filter):
                self.db.query("UPDATE regexes SET blocked_packets = ? WHERE regex_id = ?;", (filter.blocked, filter.code))

            if not proxy:
                proxy = Proxy(
                    internal_port=data['internal_port'],
                    public_port=data['public_port'],
                    filters=list(filters.values()),
                    internal_host=LOCALHOST_IP,
                    callback_blocked_update=stats_updater
                )
            
            #Port checks
            if proxy.internal_port != data['internal_port'] or proxy.public_port != data['public_port']:
                proxy.internal_port = data['internal_port']
                proxy.public_port = data['public_port']
                restart_required = True
            
            #Update filters
            if reload_required:
                proxy.filters = list(filters.values())

            #proxy status managment
            if previous_status != next_status[0]:
                # ACTIVE -> PAUSE or PAUSE -> ACTIVE
                if (previous_status, next_status[0]) in [(STATUS.ACTIVE, STATUS.PAUSE), (STATUS.PAUSE, STATUS.ACTIVE)]:
                    if restart_required:
                        proxy.restart(in_pause=next_status[0])
                    else:
                        if next_status[0] == STATUS.ACTIVE: proxy.reload()
                        else: proxy.pause()
                    previous_status = next_status[0]
                    self.__update_status_db(id, next_status[0])
                    reload_required = restart_required = False

                # ACTIVE -> STOP
                elif (previous_status,next_status[0]) in [(STATUS.ACTIVE, STATUS.STOP), (STATUS.WAIT, STATUS.STOP), (STATUS.PAUSE, STATUS.STOP)]: #Stop proxy
                    if thr_starter and thr_starter.is_alive(): thr_starter.kill()
                    proxy.stop()
                    previous_status = next_status[0] = STATUS.STOP
                    self.__update_status_db(id, STATUS.STOP)
                    reload_required = restart_required = False

                # STOP -> ACTIVE or STOP -> PAUSE
                elif (previous_status, next_status[0]) in [(STATUS.STOP, STATUS.ACTIVE), (STATUS.STOP, STATUS.PAUSE)]:    
                    previous_status = next_status[0] 
                    self.__update_status_db(id, STATUS.WAIT)
                    thr_starter = self.__proxy_starter(id, proxy, next_status[0], [previous_status])
                    reload_required = restart_required = False
                else:
                    self.__update_status_db(id, previous_status)
            
            if previous_status != STATUS.STOP:
                if restart_required: proxy.restart(in_pause=(previous_status == STATUS.PAUSE))
                elif reload_required and previous_status != STATUS.PAUSE: proxy.reload()

            signal.wait()
            signal.clear()
        

def check_port_is_open(port):
    try:
        sock = socket.socket()
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock.bind(('0.0.0.0',port))
        sock.close()
        return True
    except Exception:
        return False

def from_name_get_id(name):
    serv_id = name.strip().replace(" ","-")
    serv_id = "".join([c for c in serv_id if c in (string.ascii_uppercase + string.ascii_lowercase + string.digits + "-")])
    return serv_id.lower()

def gen_internal_port(db):
    while True:
        res = random.randint(30000, 45000)
        if len(db.query('SELECT 1 FROM services WHERE internal_port = ?;', (res,))) == 0:
            break
    return res