firegex-traffic-viewer/backend/modules/regexproxy/utils.py

import secrets
from modules.regexproxy.proxy import Filter, Proxy
import random, socket, asyncio
from base64 import b64decode
from utils.sqlite import SQLite

class STATUS:
    WAIT = "wait"
    STOP = "stop"
    PAUSE = "pause"
    ACTIVE = "active"

class ServiceNotFoundException(Exception): pass

class ServiceManager:
    def __init__(self, id, db):
        self.id = id
        self.db = db
        self.proxy = Proxy(
            internal_host="127.0.0.1",
            callback_blocked_update=self._stats_updater
        )
        self.status = STATUS.STOP
        self.wanted_status = STATUS.STOP
        self.filters = {}
        self._update_port_from_db()
        self._update_filters_from_db()
        self.lock = asyncio.Lock()
        self.starter = None
    
    def _update_port_from_db(self):
        res = self.db.query("""
            SELECT 
                public_port,
                internal_port
            FROM services WHERE service_id = ?;
        """, self.id)
        if len(res) == 0: raise ServiceNotFoundException()
        self.proxy.internal_port = res[0]["internal_port"]
        self.proxy.public_port = res[0]["public_port"]
        
    def _update_filters_from_db(self):
        res = self.db.query("""
            SELECT 
                regex, mode, regex_id `id`, is_blacklist,
                blocked_packets n_packets, is_case_sensitive
            FROM regexes WHERE service_id = ? AND active=1;
        """, self.id)

        #Filter check
        old_filters = set(self.filters.keys())
        new_filters = set([f["id"] for f in res])

        #remove old filters
        for f in old_filters:
            if not f in new_filters:
                del self.filters[f]
        
        for f in new_filters:
            if not f in old_filters:
                filter_info = [ele for ele in res if ele["id"] == f][0]
                self.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
                )
        self.proxy.filters = list(self.filters.values())
    
    def __update_status_db(self, status):
        self.db.query("UPDATE services SET status = ? WHERE service_id = ?;", status, self.id)

    async def next(self,to):
        async with self.lock:
            return await self._next(to)
    
    async def _next(self, to):
        if self.status != to:
            # ACTIVE -> PAUSE or PAUSE -> ACTIVE
            if (self.status, to) in [(STATUS.ACTIVE, STATUS.PAUSE)]:
                await self.proxy.pause()
                self._set_status(to)

            elif (self.status, to) in [(STATUS.PAUSE, STATUS.ACTIVE)]:
                await self.proxy.reload()
                self._set_status(to)

            # ACTIVE -> STOP
            elif (self.status,to) in [(STATUS.ACTIVE, STATUS.STOP), (STATUS.WAIT, STATUS.STOP), (STATUS.PAUSE, STATUS.STOP)]: #Stop proxy
                if self.starter: self.starter.cancel()
                await self.proxy.stop()
                self._set_status(to)

            # STOP -> ACTIVE or STOP -> PAUSE
            elif (self.status, to) in [(STATUS.STOP, STATUS.ACTIVE), (STATUS.STOP, STATUS.PAUSE)]:
                self.wanted_status = to
                self._set_status(STATUS.WAIT)
                self.__proxy_starter(to)


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

    async def update_port(self):
        async with self.lock:
            self._update_port_from_db()
            if self.status in [STATUS.PAUSE, STATUS.ACTIVE]:
                next_status = self.status if self.status != STATUS.WAIT else self.wanted_status
                await self._next(STATUS.STOP)
                await self._next(next_status)

    def _set_status(self,status):
        self.status = status
        self.__update_status_db(status)


    async def update_filters(self):
        async with self.lock:
            self._update_filters_from_db()
            if self.status in [STATUS.PAUSE, STATUS.ACTIVE]:
                await self.proxy.reload()
    
    def __proxy_starter(self,to):
        async def func():
            try:
                while True:
                    if check_port_is_open(self.proxy.public_port):
                        self._set_status(to)
                        await self.proxy.start(in_pause=(to==STATUS.PAUSE))
                        self._set_status(STATUS.STOP)
                        return
                    else:
                        await asyncio.sleep(.5)
            except asyncio.CancelledError:
                self._set_status(STATUS.STOP)
                await self.proxy.stop()
        self.starter = asyncio.create_task(func())

class ProxyManager:
    def __init__(self, db:SQLite):
        self.db = db
        self.proxy_table:dict = {}
        self.lock = asyncio.Lock()

    async def close(self):
        for key in list(self.proxy_table.keys()):
            await self.remove(key)

    async def remove(self,id):
        async with self.lock: 
            if id in self.proxy_table:
                await self.proxy_table[id].next(STATUS.STOP)
                del self.proxy_table[id]
    
    async def reload(self):
        async with self.lock: 
            for srv in self.db.query('SELECT service_id, status FROM services;'):
                srv_id, req_status = srv["service_id"], srv["status"]
                if srv_id in self.proxy_table:
                    continue

                self.proxy_table[srv_id] = ServiceManager(srv_id,self.db)
                await self.proxy_table[srv_id].next(req_status)

    def get(self,id):
        if id in self.proxy_table:
            return self.proxy_table[id]
        else:
            raise ServiceNotFoundException()

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 gen_service_id(db):
    while True:
        res = secrets.token_hex(8)
        if len(db.query('SELECT 1 FROM services WHERE service_id = ?;', res)) == 0:
            break
    return res

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