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 pydantic import BaseModel, Field
import traceback
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}}

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 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