firegex/backend/utils/__init__.py

import asyncio
from ipaddress import ip_address, ip_interface
import os
import re
import socket
import psutil
import sys
try:
    import nftables
except ImportError:
    nftables = None
from socketio import AsyncServer
from typing import Annotated, List, Union
from functools import wraps
from pydantic import BaseModel, ValidationError, Field
import traceback
from utils.models import StatusMessageModel
from pathlib import Path

from fastapi import HTTPException, status

_localhost_env = os.getenv("LOCALHOST_IP", "127.0.0.1")
if _localhost_env == "None":
    _localhost_env = "127.0.0.1"
LOCALHOST_IP = socket.gethostbyname(_localhost_env)

socketio:AsyncServer = None
sid_list:set = set()

ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
ROUTERS_DIR = os.path.join(ROOT_DIR,"routers")
ON_DOCKER = "DOCKER" in sys.argv
DEBUG = "DEBUG" in sys.argv
NORELOAD = "NORELOAD" in sys.argv
_port_env = os.getenv("PORT", "4444")
FIREGEX_PORT = int(_port_env) if _port_env and _port_env != "None" else 4444
_host_env = os.getenv("HOST", "0.0.0.0")
FIREGEX_HOST = _host_env if _host_env and _host_env != "None" else "0.0.0.0"
_socket_dir_env = os.getenv("SOCKET_DIR", None)
FIREGEX_SOCKET_DIR = None if _socket_dir_env == "None" else _socket_dir_env
FIREGEX_SOCKET = os.path.join(FIREGEX_SOCKET_DIR, "firegex.sock") if FIREGEX_SOCKET_DIR else None
JWT_ALGORITHM: str = "HS256"
def _get_version():
    v = "{{VERSION_PLACEHOLDER}}" if "{" not in "{{VERSION_PLACEHOLDER}}" else "0.0.0"
    if v == "0.0.0":
        env_v = os.getenv("FIREGEX_VERSION")
        if env_v:
            return env_v
        try:
            import re
            
            def parse_version(ver):
                ver = ver.lstrip("v")
                parts = re.split(r'[^0-9]+', ver)
                return tuple(int(p) if p.isdigit() else 0 for p in parts if p)

            def find_git_dir(start_path):
                current = start_path
                while current and current != "/" and current != "":
                    gd = os.path.join(current, ".git")
                    if os.path.isdir(gd):
                        return gd
                    parent = os.path.dirname(current)
                    if parent == current:
                        break
                    current = parent
                return None
                
            git_dir = find_git_dir(os.path.abspath(os.path.dirname(__file__)))
            if not git_dir:
                return v
                
            branch = None
            with open(os.path.join(git_dir, "HEAD"), "r") as f:
                head = f.read().strip()
                if head.startswith("ref: refs/heads/"):
                    branch = head.split("ref: refs/heads/")[1]
                    
            if branch != "main":
                return v
                
            tags = []
            try:
                with open(os.path.join(git_dir, "packed-refs"), "r") as f:
                    for line in f:
                        if " refs/tags/" in line:
                            tags.append(line.strip().split(" refs/tags/")[1])
            except Exception:
                pass
                
            tags_dir = os.path.join(git_dir, "refs", "tags")
            try:
                for root, _, files in os.walk(tags_dir):
                    for file in files:
                        tags.append(os.path.relpath(os.path.join(root, file), tags_dir).replace(os.path.sep, "/"))
            except Exception:
                pass
                
            version_tags = [t for t in set(tags) if any(c.isdigit() for c in t)]
            if version_tags:
                return max(version_tags, key=parse_version)
        except Exception:
            pass
    return v

API_VERSION = _get_version()

PortType = Annotated[int, Field(gt=0, lt=65536)]   

def safe_join(base_dir: Union[str, Path], *paths: str) -> Path:
    """
    Safely join a base directory with one or more path components.
    
    Returns the resolved Path object if safe.
    Raises HTTPException 403 if a directory traversal attack is detected.
    """
    # 1. Convert base_dir to a resolved, absolute Path
    base_path = Path(base_dir).resolve()
    
    # 2. Join the parts and resolve the result (removes symbols like '../')
    # If the user input is an absolute path (starting with /), 
    # joinpath handles it correctly or we can strip leading slashes.
    clean_paths = [p.lstrip("/") for p in paths]
    target_path = base_path.joinpath(*clean_paths).resolve()

    # 3. Security check
    if not target_path.is_relative_to(base_path):
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Access denied: Path is invalid."
        )

    return target_path

async def run_func(func, *args, **kwargs):
    if asyncio.iscoroutinefunction(func): 
        return await func(*args, **kwargs)
    else: 
        return func(*args, **kwargs)

async def socketio_emit(elements:list[str]):
    await socketio.emit("update",elements)

def refactor_name(name:str):
    name = name.strip()
    while "  " in name:
        name = name.replace("  "," ")
    return name

class SysctlManager:
    def __init__(self, ctl_table):
        self.old_table = {}
        self.new_table = {}
        if os.path.isdir("/sys_host/"):
            self.old_table = dict()
            self.new_table = dict(ctl_table)
            for name in ctl_table.keys():
                self.old_table[name] = read_sysctl(name)
    
    def write_table(self, table) -> bool:
        for name, value in table.items():
            if read_sysctl(name) != value:
                write_sysctl(name, value)
                
    def set(self):
        self.write_table(self.new_table)

    def reset(self):
        self.write_table(self.old_table)

def read_sysctl(name:str):
    with open(f"/sys_host/{name}", "rt") as f:
        return "1" in f.read()

def write_sysctl(name:str, value:bool):
    with open(f"/sys_host/{name}", "wt") as f:
        f.write("1" if value else "0")

def list_files(mypath):
    from os import listdir
    from os.path import isfile, join
    return [f for f in listdir(mypath) if isfile(join(mypath, f))]

def ip_parse(ip:str):
    return str(ip_interface(ip).network)

def is_ip_parse(ip:str):
    try:
        ip_parse(ip)
        return True
    except Exception:
        return False

def is_interface_name(name: str) -> bool:
    if not isinstance(name, str):
        return False
    name = name.strip()
    return bool(re.match(r'^[a-zA-Z0-9_.:-]+$', name)) and 0 < len(name) <= 15

def parse_ip_or_int(ip: str) -> str:
    if is_ip_parse(ip):
        return ip_parse(ip)
    ip_str = str(ip).strip()
    if is_interface_name(ip_str):
        return ip_str
    raise ValueError(f"'{ip}' is neither a valid IP address nor a valid interface name")

def addr_parse(ip:str):
    return str(ip_address(ip))

def ip_family(ip:str):
    return "ip6" if ip_interface(ip).version == 6 else "ip"

def get_interfaces():
    def _get_interfaces():
        for int_name, interfs in psutil.net_if_addrs().items():
            for interf in interfs:
                if interf.family in [socket.AF_INET, socket.AF_INET6]:
                    yield {"name": int_name, "addr":interf.address}
    return list(_get_interfaces())

def get_interface_ips(iface_name: str) -> list[str]:
    ips = []
    for int_name, interfs in psutil.net_if_addrs().items():
        if int_name == iface_name:
            for interf in interfs:
                if interf.family in [socket.AF_INET, socket.AF_INET6]:
                    ips.append(interf.address)
    return ips

def nftables_int_to_json(ip_int):
    ip_int = ip_parse(ip_int)
    ip_addr = str(ip_int).split("/")[0]
    ip_addr_cidr = int(str(ip_int).split("/")[1])
    return {"prefix": {"addr": ip_addr, "len": ip_addr_cidr}}

def nftables_json_to_int(ip_json_int):
    if isinstance(ip_json_int,str):
        return str(ip_parse(ip_json_int))
    else:
        return f'{ip_json_int["prefix"]["addr"]}/{ip_json_int["prefix"]["len"]}'
    
class Singleton(object):
    __instance = None
    def __new__(class_, *args, **kwargs):
        if not isinstance(class_.__instance, class_):
            class_.__instance = object.__new__(class_, *args, **kwargs)
        return class_.__instance

class NFTableManager(Singleton):
    
    table_name = "fgex"
    
    def __init__(self, init_cmd, reset_cmd):
        self.__init_cmds = init_cmd
        self.__reset_cmds = reset_cmd
        self.nft = nftables.Nftables() if nftables is not None else None
    
    def raw_cmd(self, *cmds):
        return self.nft.json_cmd({"nftables": list(cmds)})

    def cmd(self, *cmds):
        code, out, err = self.raw_cmd(*cmds)
        if code == 0:
            return out
        else:
            raise Exception(err)
    
    def init(self):
        self.reset()
        self.raw_cmd({"add":{"table":{"name":self.table_name,"family":"inet"}}})
        self.cmd(*self.__init_cmds)
            
    def reset(self):
        self.raw_cmd(*self.__reset_cmds)

    def list_rules(self, tables = None, chains = None):
        for filter in [ele["rule"] for ele in self.raw_list() if "rule" in ele ]:
            if tables and filter["table"] not in tables:
                continue
            if chains and filter["chain"] not in chains:
                continue
            yield filter
    
    def raw_list(self):
        return self.cmd({"list": {"ruleset": None}})["nftables"]

def _json_like(obj: BaseModel|List[BaseModel], unset=False, convert_keys:dict[str, str]=None, exclude:list[str]=None, mode:str="json"):
    res = obj.model_dump(mode=mode, exclude_unset=not unset)
    if convert_keys:
        for from_k, to_k in convert_keys.items():
            if from_k in res:
                res[to_k] = res.pop(from_k)
    if exclude:
        for ele in exclude:
            if ele in res:
                del res[ele]
    return res

def json_like(obj: BaseModel|List[BaseModel], unset=False, convert_keys:dict[str, str]=None, exclude:list[str]=None, mode:str="json") -> dict:
    if isinstance(obj, list):
        return [_json_like(ele, unset=unset, convert_keys=convert_keys, exclude=exclude, mode=mode) for ele in obj]
    return _json_like(obj, unset=unset, convert_keys=convert_keys, exclude=exclude, mode=mode)

def register_event(sio_server: AsyncServer, event_name: str, model: BaseModel, response_model: BaseModel|None = None):
    def decorator(func):
        @sio_server.on(event_name)  # Automatically registers the event
        @wraps(func)
        async def wrapper(sid, data):
            try:
                # Parse and validate incoming data
                parsed_data = model.model_validate(data)
            except ValidationError:
                return json_like(StatusMessageModel(status=f"Invalid {event_name} request"))
            
            # Call the original function with the parsed data
            result = await func(sid, parsed_data)
            # If a response model is provided, validate the output
            if response_model:
                try:
                    parsed_result = response_model.model_validate(result)
                except ValidationError:
                    traceback.print_exc()
                    return json_like(StatusMessageModel(status=f"SERVER ERROR: Invalid {event_name} response"))
            else:
                parsed_result = result
            # Emit the validated result
            if parsed_result:
                if isinstance(parsed_result, BaseModel):
                    return json_like(parsed_result)
                return parsed_result
        return wrapper
    return decorator

def nicenessify(priority:int, pid:int|None=None):
    try:
        pid = os.getpid() if pid is None else pid
        ps = psutil.Process(pid)
        if os.name == 'posix':
            ps.nice(priority)
    except (psutil.AccessDenied, PermissionError) as e:
        print(f"Permission denied setting priority (pid={pid}, priority={priority}): {e}")
    except Exception as e:
        print(f"Error setting priority: {e} {traceback.format_exc()}")
        pass