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

#: What a backup entry may be called: a plain basename of one of the two shapes
#: `export_db()` produces — `<name>.db` for a database and `<id>.py` for a filter's code.
#: The charset forbids path separators and `..`, so a crafted key cannot escape the
#: directory it is joined onto.
#:
#: Here rather than in `app.py`, beside `safe_join`, because the two are one defence in
#: two halves: this decides what a name may contain and that one decides where it may
#: land. Keeping them apart meant the charset could only be exercised by importing the
#: whole application — a server, a JWT library and everything else `app.py` pulls in —
#: to check two regular expressions.
#:
#: Matched with `fullmatch`, never `match`: in Python `$` also matches *before a trailing
#: newline*, so `"services.db\n"` satisfied a pattern written to mean "this and nothing
#: else".
SAFE_DB_NAME = re.compile(r'[A-Za-z0-9_-]+\.db')
SAFE_PY_NAME = re.compile(r'[A-Za-z0-9_-]+\.py')


def boot_auth_mode(held: str | None, fresh_boot: bool, env_disabled: bool) -> tuple[bool, bool]:
    """Whether authentication starts off disabled, and whether the host's stored decision
    is kept — `(disabled, keep_held)`.

    `held` is what `run.py config` last stored for the host (`"1"`, `"0"` or `None`). A
    fresh boot — a container or a standalone start `run.py` has just written the
    environment for — takes the environment, which already says what the configuration
    says, and the stored decision is dropped. Any other boot is Docker starting an existing
    container again with the environment it was created with, which may be older than the
    stored decision, and the stored decision wins.
    """
    if fresh_boot or held is None:
        return env_disabled, False
    return held == "1", True


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:
    """Normalise what an operator typed into what the rules and the key will use.

    **Stripped once, before either question is asked.** `is_interface_name` strips and
    `is_ip_parse` does not, so surrounding whitespace used to decide *which of the two an
    address was*: `" 10.0.0.1 "` failed the IP test, passed the interface-name test —
    every character in it is in that charset — and was stored verbatim. Nothing then
    broke loudly, which is the problem: the rules still matched (both sides re-parse the
    value), but `"10.0.0.1"` and `"10.0.0.1/32"` are different strings, so the
    `(ip_int, port, proto)` uniqueness key no longer saw them as the same address and two
    services could each believe they were protecting it.
    """
    ip_str = str(ip).strip()
    if is_ip_parse(ip_str):
        return ip_parse(ip_str)
    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, family: str | None = None):
        for filter in [ele["rule"] for ele in self.raw_list(tables, family) 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, tables = None, family: str | None = None):
        # One table of ours when that is all that is asked for, rather than the host's
        # whole ruleset: on a node running Docker or Kubernetes that is thousands of rules
        # of somebody else's, serialised and parsed on every lookup — and a lookup runs
        # inside the backend's one event loop, on every start, stop, address change and
        # statistics request. A table that is not there is simply no rules.
        if tables and family:
            out = []
            for table in tables:
                code, listed, _ = self.raw_cmd({"list": {"table": {"family": family, "name": table}}})
                if code == 0:
                    out.extend(listed["nftables"])
            return out
        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