push: code changes

backend/utils/__init__.py

@@ -8,15 +8,22 @@ import nftables
 from socketio import AsyncServer
 from fastapi import Path
 from typing import Annotated
+from functools import wraps
+from pydantic import BaseModel, ValidationError
+import traceback
+from utils.models import StatusMessageModel
+from typing import List
 
 LOCALHOST_IP = socket.gethostbyname(os.getenv("LOCALHOST_IP","127.0.0.1"))
 
 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
 FIREGEX_PORT = int(os.getenv("PORT","4444"))
 JWT_ALGORITHM: str = "HS256"
 API_VERSION = "{{VERSION_PLACEHOLDER}}" if "{" not in "{{VERSION_PLACEHOLDER}}" else "0.0.0"
@@ -153,4 +160,50 @@ class NFTableManager(Singleton):
     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
+