From 42f3c22c8289b98f66747cdb0ac721fb80bee15e Mon Sep 17 00:00:00 2001 From: FerTV Date: Wed, 23 Jul 2025 13:52:22 +0200 Subject: [PATCH 01/14] refactor: - controller.py renamed to hub.py - database.py renamed to postgress.py - database services in controller.py moved to databases folder --- Makefile | 2 +- app/deployer.py | 4 +- nebula/controller/database.py | 635 ----------------- nebula/controller/{controller.py => hub.py} | 126 ++-- nebula/controller/start_services.sh | 2 +- .../postgress/docker}/Dockerfile | 2 +- .../postgress/docker}/docker-entrypoint.sh | 0 .../postgress/docker}/init-configs.sql | 0 .../database/adapters/postgress/postgress.py | 655 ++++++++++++++++++ 9 files changed, 704 insertions(+), 722 deletions(-) delete mode 100755 nebula/controller/database.py rename nebula/controller/{controller.py => hub.py} (90%) rename nebula/database/{ => adapters/postgress/docker}/Dockerfile (93%) rename nebula/database/{ => adapters/postgress/docker}/docker-entrypoint.sh (100%) rename nebula/database/{ => adapters/postgress/docker}/init-configs.sql (100%) create mode 100755 nebula/database/adapters/postgress/postgress.py diff --git a/Makefile b/Makefile index c1e0cd768..44a6cf390 100644 --- a/Makefile +++ b/Makefile @@ -76,7 +76,7 @@ update-dockers: ## Update docker images fi @echo "🐳 Building nebula-database docker image. Do you want to continue (overrides existing image)? (y/n)" @read ans; if [ "$${ans:-N}" = y ]; then \ - docker build -t nebula-database -f nebula/database/Dockerfile .; \ + docker build -t nebula-database -f nebula/database/docker/Dockerfile .; \ docker build -t nebula-pgweb -f nebula/database/pgweb/Dockerfile .; \ else \ echo "Skipping nebula-database docker build."; \ diff --git a/app/deployer.py b/app/deployer.py index 95bffda7e..7e46fe1a7 100644 --- a/app/deployer.py +++ b/app/deployer.py @@ -17,7 +17,7 @@ from watchdog.observers import Observer from nebula.addons.env import check_environment -from nebula.controller.controller import TermEscapeCodeFormatter +from nebula.controller.hub import TermEscapeCodeFormatter from nebula.controller.scenarios import ScenarioManagement from nebula.utils import DockerUtils, FileUtils, SocketUtils @@ -1033,7 +1033,7 @@ def run_database(self): "POSTGRES_PASSWORD": os.environ.get("POSTGRES_PASSWORD"), "POSTGRES_DB": "nebula", } - host_sql_path = os.path.join(self.root_path, "nebula/database/init-configs.sql") + host_sql_path = os.path.join(self.root_path, "nebula/database/adapters/postgress/docker/init-configs.sql") db_data_path = os.path.join(self.databases_dir, "postgres-data") os.makedirs(db_data_path, exist_ok=True) diff --git a/nebula/controller/database.py b/nebula/controller/database.py deleted file mode 100755 index 407ce3908..000000000 --- a/nebula/controller/database.py +++ /dev/null @@ -1,635 +0,0 @@ -import logging -import os -import datetime -import json -import asyncpg -import asyncio - -from passlib.context import CryptContext - -# --- Configuration --- -# Use environment variables for database credentials from the Docker Compose file -DATABASE_URL = f"postgresql://{os.environ.get('DB_USER')}:{os.environ.get('DB_PASSWORD')}@{os.environ.get('DB_HOST')}:{os.environ.get('DB_PORT')}/nebula" - -# Password hashing context (using Argon2) -pwd_context = CryptContext(schemes=["argon2"], deprecated="auto") - -# Asynchronous lock for node updates -_node_lock = asyncio.Lock() - -# --- Connection Pool Management --- -# Global pool variable, should be initialized at application startup -POOL = None - -async def init_db_pool(): - """ - Initializes the asynchronous PostgreSQL connection pool. - This should be called once when the application starts. - """ - global POOL - if POOL is None: - try: - POOL = await asyncpg.create_pool( - dsn=DATABASE_URL, - min_size=5, # Minimum number of connections in the pool - max_size=20, # Maximum number of connections in the pool - ) - logging.info("Database connection pool successfully created.") - except Exception as e: - logging.critical(f"Failed to create database connection pool: {e}", exc_info=True) - # Exit or handle the failure appropriately - raise - -async def close_db_pool(): - """ - Closes the asynchronous PostgreSQL connection pool. - This should be called once when the application shuts down gracefully. - """ - global POOL - if POOL: - await POOL.close() - logging.info("Database connection pool closed.") - - -# --- User Management Functions --- - -async def insert_default_admin(): - """ - Inserts a default 'ADMIN' user into the database with a hashed password. - The password must be provided via the ADMIN_PASSWORD environment variable. - """ - admin_password = os.environ.get("NEBULA_ADMIN_PASSWORD") - - hashed_password = pwd_context.hash(admin_password) - - query = """ - INSERT INTO users ("user", password, role) - VALUES ($1, $2, $3) - ON CONFLICT ("user") DO NOTHING; - """ - try: - async with POOL.acquire() as conn: - await conn.execute(query, "ADMIN", hashed_password, "admin") - logging.info("Default admin user inserted (or already exists).") - except Exception as e: - logging.error(f"Failed to insert default admin user: {e}", exc_info=True) - -async def list_users(all_info=False): - """ - Retrieves a list of users from the users database. - """ - async with POOL.acquire() as conn: - result = await conn.fetch("SELECT * FROM users") - - if not all_info: - result = [user["user"] for user in result] - - return result - - -async def get_user_info(user): - """ - Fetches detailed information for a specific user from the users database. - """ - async with POOL.acquire() as conn: - return await conn.fetchrow('SELECT * FROM users WHERE "user" = $1', user) - - -async def verify(user, password): - """ - Verifies whether the provided password matches the stored hashed password for a user. - """ - async with POOL.acquire() as conn: - result = await conn.fetchrow('SELECT password FROM users WHERE "user" = $1', user) - if result: - try: - return pwd_context.verify(password, result[0]) - except Exception: - # Catch more general exceptions during verification to be safe - logging.error(f"Error during password verification for user {user}", exc_info=True) - return False - return False - - -async def verify_hash_algorithm(user): - """ - Checks if the stored password hash for a user uses a supported Argon2 algorithm. - """ - user = user.upper() - argon2_prefixes = ("$argon2i$", "$argon2id$") - async with POOL.acquire() as conn: - result = await conn.fetchrow('SELECT password FROM users WHERE "user" = $1', user) - if result: - password_hash = result["password"] - return password_hash.startswith(argon2_prefixes) - return False - - -async def delete_user_from_db(user): - """ - Deletes a user record from the users database. - """ - async with POOL.acquire() as conn: - await conn.execute('DELETE FROM users WHERE "user" = $1', user) - - -async def add_user(user, password, role): - """ - Adds a new user to the users database with a hashed password. - """ - hashed_password = pwd_context.hash(password) - async with POOL.acquire() as conn: - await conn.execute( - 'INSERT INTO users ("user", password, role) VALUES ($1, $2, $3)', - user.upper(), hashed_password, role, - ) - - -async def update_user(user, password, role): - """ - Updates the password and role of an existing user in the users database. - """ - hashed_password = pwd_context.hash(password) - async with POOL.acquire() as conn: - await conn.execute( - 'UPDATE users SET password = $1, role = $2 WHERE "user" = $3', - hashed_password, role, user.upper(), - ) - -# --- Node Management Functions --- - -async def list_nodes(scenario_name=None, sort_by="idx"): - """ - Retrieves a list of nodes from the nodes database, optionally filtered by scenario and sorted. - """ - # Validate sort_by to prevent SQL injection - allowed_sort_fields = ["uid", "idx", "ip", "port", "role", "timestamp", "federation", "round"] - if sort_by not in allowed_sort_fields: - sort_by = "idx" # Default to a safe field - - try: - async with POOL.acquire() as conn: - if scenario_name: - # Using f-string for column names is generally safe if validated as above - command = f"SELECT * FROM nodes WHERE scenario = $1 ORDER BY {sort_by};" - result = await conn.fetch(command, scenario_name) - else: - command = f"SELECT * FROM nodes ORDER BY {sort_by};" - result = await conn.fetch(command) - return result - except asyncpg.PostgresError as e: - logging.error(f"Error occurred while listing nodes: {e}") - return None - - -async def list_nodes_by_scenario_name(scenario_name): - """ - Fetches all nodes associated with a specific scenario, ordered by their index as integers. - """ - try: - async with POOL.acquire() as conn: - command = "SELECT * FROM nodes WHERE scenario = $1 ORDER BY CAST(idx AS INTEGER) ASC;" - result = await conn.fetch(command, scenario_name) - return [dict(record) for record in result] - except Exception as e: - logging.error(f"Error occurred while listing nodes by scenario name: {e}") - return None - - -async def update_node_record( - node_uid, idx, ip, port, role, neighbors, latitude, longitude, - timestamp, federation, federation_round, scenario, run_hash, malicious, -): - """ - Inserts or updates a node record in the database for a given scenario, ensuring thread-safe access. - """ - async with _node_lock: - async with POOL.acquire() as conn: - try: - async with conn.transaction(): - result = await conn.fetchrow( - "SELECT * FROM nodes WHERE uid = $1 AND scenario = $2 FOR UPDATE;", - node_uid, scenario - ) - - if result is None: - # Insert new node - await conn.execute( - """ - INSERT INTO nodes (uid, idx, ip, port, role, neighbors, latitude, longitude, - timestamp, federation, round, scenario, hash, malicious) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14); - """, - node_uid, idx, ip, port, role, neighbors, latitude, longitude, - timestamp, federation, federation_round, scenario, run_hash, malicious, - ) - else: - # Update existing node - await conn.execute( - """ - UPDATE nodes SET idx = $1, ip = $2, port = $3, role = $4, neighbors = $5, - latitude = $6, longitude = $7, timestamp = $8, federation = $9, round = $10, - hash = $11, malicious = $12 - WHERE uid = $13 AND scenario = $14; - """, - idx, ip, port, role, neighbors, latitude, longitude, - timestamp, federation, federation_round, run_hash, malicious, - node_uid, scenario, - ) - - updated_row = await conn.fetchrow("SELECT * from nodes WHERE uid = $1 AND scenario = $2;", node_uid, scenario) - return dict(updated_row) if updated_row else None - except asyncpg.PostgresError as e: - logging.error(f"Database error during node record update: {e}", exc_info=True) - return None - - -async def remove_all_nodes(): - """ - Deletes all node records from the nodes database. - """ - async with POOL.acquire() as conn: - await conn.execute("TRUNCATE nodes CASCADE;") # Use CASCADE if there are foreign key dependencies - - -async def remove_nodes_by_scenario_name(scenario_name): - """ - Deletes all nodes associated with a specific scenario from the database. - """ - async with POOL.acquire() as conn: - await conn.execute("DELETE FROM nodes WHERE scenario = $1;", scenario_name) - -# --- Scenario Management Functions --- - -async def get_all_scenarios(username, role, sort_by="start_time"): - """ - Retrieves all scenarios from the database, accessing fields from the 'config' (JSONB) column - and direct columns. Filters by user role and sorts by the specified field. - """ - allowed_sort_fields = ["start_time", "title", "username", "status", "name"] - if sort_by not in allowed_sort_fields: - sort_by = "start_time" - - # Determine the ORDER BY clause based on sort_by - if sort_by == "start_time": - order_by_clause = """ - ORDER BY - CASE - WHEN start_time IS NULL OR start_time = '' THEN 1 - ELSE 0 - END, - to_timestamp(start_time, 'DD/MM/YYYY HH24:MI:SS') DESC - """ - elif sort_by in ["title", "model", "dataset", "rounds"]: # These are inside config JSONB - order_by_clause = f"ORDER BY config->>'{sort_by}'" - else: # For direct table columns like name, username, status - order_by_clause = f"ORDER BY {sort_by}" - - async with POOL.acquire() as conn: - # Select direct columns and relevant fields from config JSONB - command = """ - SELECT - name, - username, - status, - start_time, - end_time, - config->>'title' AS title, - config->>'model' AS model, - config->>'dataset' AS dataset, - config->>'rounds' AS rounds, - config -- return the full config JSONB - FROM scenarios - """ - params = [] - - if role != "admin": - command += " WHERE username = $1" # username is a direct column now - params.append(username) - - full_command = f"{command} {order_by_clause};" - return await conn.fetch(full_command, *params) - - -async def get_all_scenarios_and_check_completed(username, role, sort_by="start_time"): - """ - Retrieves all scenarios, sorts them, and updates the status if necessary. - Returns a list of dictionaries, where each dictionary represents a scenario. - """ - # Safe list of allowed sorting fields to prevent SQL injection. - allowed_sort_fields = ["start_time", "title", "username", "status", "name"] - if sort_by not in allowed_sort_fields: - sort_by = "start_time" # Safe default value - - # Building the ORDER BY clause - if sort_by == "start_time": - order_by_clause = """ - ORDER BY - CASE - WHEN start_time IS NULL OR start_time = '' THEN 1 - ELSE 0 - END, - to_timestamp(start_time, 'DD/MM/YYYY HH24:MI:SS') DESC - """ - elif sort_by in ["title", "model", "dataset", "rounds"]: # These are inside config JSONB - order_by_clause = f"ORDER BY config->>'{sort_by}'" - else: # For direct table columns like name, username, status - order_by_clause = f"ORDER BY {sort_by}" - - async with POOL.acquire() as conn: - # Base query that extracts fields from the JSONB using the ->> operator - command = f""" - SELECT - name, - username, - status, - start_time, - end_time, - config->>'title' AS title, - config->>'model' AS model, - config->>'dataset' AS dataset, - config->>'rounds' AS rounds, - config -- Return the full config object - FROM scenarios - """ - params = [] - if role != "admin": - command += " WHERE username = $1" # username is a direct column - params.append(username) - - command += f" {order_by_clause};" - - result_dicts = await conn.fetch(command, *params) - - scenarios_to_return = [dict(s) for s in result_dicts] - - re_fetch_required = False - for scenario in scenarios_to_return: - if scenario["status"] == "running": - if await check_scenario_federation_completed(scenario["name"]): - await scenario_set_status_to_completed(scenario["name"]) - re_fetch_required = True - break - - if re_fetch_required: - # Recursively call to get fresh data after status update - return await get_all_scenarios_and_check_completed(username, role, sort_by) - - return scenarios_to_return - - -async def scenario_update_record(name, start_time, end_time, scenario_config, status, username): - """ - Inserts or updates a scenario record using the PostgreSQL "UPSERT" pattern. - All configuration is saved in the 'config' column of type JSONB. - Direct columns (name, start_time, end_time, username, status) are also handled. - """ - # Ensure scenario_config is a dictionary before dumping to JSON - if not isinstance(scenario_config, dict): - try: - scenario_config = json.loads(scenario_config) - except (json.JSONDecodeError, TypeError): - logging.error("scenario_config is not a valid JSON string or dict.") - return - - command = """ - INSERT INTO scenarios (name, start_time, end_time, username, status, config) - VALUES ($1, $2, $3, $4, $5, $6::jsonb) - ON CONFLICT (name) DO UPDATE SET - start_time = EXCLUDED.start_time, - end_time = EXCLUDED.end_time, - username = EXCLUDED.username, - status = EXCLUDED.status, - config = scenarios.config || EXCLUDED.config; -- Merge JSONB - """ - async with POOL.acquire() as conn: - await conn.execute(command, name, start_time, end_time, username, status, json.dumps(scenario_config)) - - -async def scenario_set_all_status_to_finished(): - """ - Sets the status of all 'running' scenarios to 'finished' - and updates their 'end_time' (both in the direct column and within JSONB). - """ - current_time = datetime.datetime.now().strftime('%Y/%m/%d %H:%M:%S') # Consistent format - command = """ - UPDATE scenarios - SET - status = 'finished', - end_time = $1, - config = jsonb_set(config, '{status}', '"finished"') || - jsonb_set(config, '{end_time}', $2::jsonb) - WHERE status = 'running'; - """ - async with POOL.acquire() as conn: - await conn.execute(command, current_time, json.dumps(current_time)) - - -async def scenario_set_status_to_finished(scenario_name): - """ - Sets the status of a specific scenario to 'finished' and updates its 'end_time'. - Updates both the direct columns and the JSONB 'config'. - """ - current_time = datetime.datetime.now().strftime('%Y/%m/%d %H:%M:%S') # Consistent format - command = """ - UPDATE scenarios - SET - status = 'finished', - end_time = $1, - config = jsonb_set( - jsonb_set(config, '{status}', '"finished"'), - '{end_time}', $2::jsonb - ) - WHERE name = $3; - """ - async with POOL.acquire() as conn: - await conn.execute(command, current_time, json.dumps(current_time), scenario_name) - - -async def scenario_set_status_to_completed(scenario_name): - """ - Sets the status of a specific scenario to 'completed'. - Updates both the direct column and the JSONB 'config'. - """ - command = """ - UPDATE scenarios - SET - status = 'completed', - config = jsonb_set(config, '{status}', '"completed"') - WHERE name = $1; - """ - async with POOL.acquire() as conn: - await conn.execute(command, scenario_name) - - -async def get_running_scenario(username=None, get_all=False): - """ - Retrieves scenarios with a 'running' status, optionally filtered by user. - Returns full scenario record (including direct columns and config JSONB). - """ - async with POOL.acquire() as conn: - params = ["running"] - # Select all columns to get both direct and config data - command = "SELECT name, username, status, start_time, end_time, config FROM scenarios WHERE status = $1" - - if username: - command += " AND username = $2" - params.append(username) - - if get_all: - result = [dict(row) for row in await conn.fetch(command, *params)] # Convert records to dicts - else: - result_row = await conn.fetchrow(command, *params) - result = dict(result_row) if result_row else None - return result - - -async def get_completed_scenario(): - """ - Retrieves a single scenario with a 'completed' status. - Returns full scenario record (including direct columns and config JSONB). - """ - async with POOL.acquire() as conn: - command = "SELECT name, username, status, start_time, end_time, config FROM scenarios WHERE status = $1;" - result_row = await conn.fetchrow(command, "completed") - return dict(result_row) if result_row else None - - -async def get_scenario_by_name(scenario_name): - """ - Retrieves the complete record of a scenario by its name. - """ - async with POOL.acquire() as conn: - result_row = await conn.fetchrow("SELECT name, start_time, end_time, username, status, config FROM scenarios WHERE name = $1;", scenario_name) - - result = dict(result_row) if result_row else None - - if result and result.get('config'): - # Assuming 'config' is a JSON string from the DB, so we parse it - # It might already be a dict if asyncpg handles JSONB conversion automatically - config_data = result['config'] - if isinstance(config_data, str): - try: - config_data = json.loads(config_data) - except json.JSONDecodeError: - config_data = {} - - # Extract the 'scenario_title' and add it as a top-level key - result['title'] = config_data.get('scenario_title') - result['description'] = config_data.get('description') - - return result - - -async def get_user_by_scenario_name(scenario_name): - """ - Retrieves the username associated with a scenario (from the direct 'username' column). - """ - async with POOL.acquire() as conn: - return await conn.fetchval("SELECT username FROM scenarios WHERE name = $1;", scenario_name) - - -async def remove_scenario_by_name(scenario_name): - """ - Delete a scenario from the database by its unique name. - """ - try: - async with POOL.acquire() as conn: - await conn.execute("DELETE FROM scenarios WHERE name = $1;", scenario_name) - logging.info(f"Scenario '{scenario_name}' successfully removed.") - except asyncpg.PostgresError as e: - logging.error(f"Error occurred while deleting scenario '{scenario_name}': {e}") - - -async def check_scenario_federation_completed(scenario_name): - """ - Check if all nodes in a given scenario have completed the required federation rounds. - """ - try: - async with POOL.acquire() as conn: - # Retrieve the total rounds for the scenario from the 'config' JSONB column - scenario_rounds_str = await conn.fetchval("SELECT config->>'rounds' AS rounds FROM scenarios WHERE name = $1;", scenario_name) - - if not scenario_rounds_str: - logging.warning(f"Scenario '{scenario_name}' not found or 'rounds' not defined.") - return False - - # Ensure total_rounds is an integer for comparison - try: - total_rounds = int(scenario_rounds_str) - except (ValueError, TypeError): - logging.error(f"Invalid 'rounds' value for scenario '{scenario_name}': {scenario_rounds_str}") - return False - - # Fetch the current round progress of all nodes in that scenario - nodes = await conn.fetch("SELECT round FROM nodes WHERE scenario = $1;", scenario_name) - - if not nodes: - logging.info(f"No nodes found for scenario '{scenario_name}'. Federation not considered completed.") - return False - - # Check if all nodes have completed the total rounds - return all(int(node["round"]) >= total_rounds for node in nodes) - - except asyncpg.PostgresError as e: - logging.error(f"PostgreSQL error during check_scenario_federation_completed for '{scenario_name}': {e}") - return False - except ValueError as e: - logging.error(f"Data error during check_scenario_federation_completed for '{scenario_name}': {e}") - return False - - -async def check_scenario_with_role(role, scenario_name, current_username=None): - """ - Verify if a scenario exists that the user with the given role and username can access. - """ - scenario_info = await get_scenario_by_name(scenario_name) - - if not scenario_info: - return False # Scenario does not exist - - if role == "admin": - return True # Admins can access any existing scenario - - if current_username is None: - logging.info(f"[FER] db username {scenario_info.get('username')} current_username {current_username}") - logging.warning( - "check_scenario_with_role called for non-admin role without current_username." - ) - return False - - return scenario_info.get("username") == current_username - -# --- Notes Management Functions --- - -async def save_notes(scenario, notes): - """ - Save or update notes associated with a specific scenario. - """ - try: - async with POOL.acquire() as conn: - await conn.execute( - """ - INSERT INTO notes (scenario, scenario_notes) VALUES ($1, $2) - ON CONFLICT(scenario) DO UPDATE SET scenario_notes = EXCLUDED.scenario_notes; - """, - scenario, notes, - ) - except asyncpg.PostgresError as e: - logging.error(f"PostgreSQL error during save_notes: {e}") - - -async def get_notes(scenario): - """ - Retrieve notes associated with a specific scenario. - """ - async with POOL.acquire() as conn: - return await conn.fetchrow("SELECT * FROM notes WHERE scenario = $1;", scenario) - - -async def remove_note(scenario): - """ - Delete the note associated with a specific scenario. - """ - async with POOL.acquire() as conn: - await conn.execute("DELETE FROM notes WHERE scenario = $1;", scenario) diff --git a/nebula/controller/controller.py b/nebula/controller/hub.py similarity index 90% rename from nebula/controller/controller.py rename to nebula/controller/hub.py index 7e59ae7ec..628d61322 100755 --- a/nebula/controller/controller.py +++ b/nebula/controller/hub.py @@ -16,17 +16,17 @@ from fastapi import Body, FastAPI, Request, status, HTTPException, Path, File, UploadFile from fastapi.concurrency import asynccontextmanager -from nebula.controller.database import ( - init_db_pool, - close_db_pool, - insert_default_admin, - scenario_set_all_status_to_finished, - scenario_set_status_to_finished, -) +from nebula.database.database_adapter_factory import factory_database_adapter from nebula.controller.http_helpers import remote_get, remote_post_form from nebula.utils import DockerUtils +# Get a database instance +# db = get_database() +#TODO review +db = factory_database_adapter("PostgresDB") + + # Setup controller logger class TermEscapeCodeFormatter(logging.Formatter): """ @@ -120,13 +120,13 @@ async def lifespan(app: FastAPI): configure_logger(controller_log) # Initialize the database connection pool - await init_db_pool() - await insert_default_admin() + await db.init_db_pool() + await db.insert_default_admin() yield # Code to run on shutdown - await close_db_pool() + await db.close_db_pool() # Initialize FastAPI app outside the Controller class @@ -383,9 +383,9 @@ async def stop_scenario( ScenarioManagement.cleanup_scenario_containers() try: if all: - await scenario_set_all_status_to_finished() + await db.scenario_set_all_status_to_finished() else: - await scenario_set_status_to_finished(scenario_name) + await db.scenario_set_status_to_finished(scenario_name) except Exception as e: logging.exception(f"Error setting scenario {scenario_name} to finished: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -404,11 +404,10 @@ async def remove_scenario( Returns: dict: A message indicating successful removal. """ - from nebula.controller.database import remove_scenario_by_name, get_user_by_scenario_name from nebula.controller.scenarios import ScenarioManagement try: - await remove_scenario_by_name(scenario_name) + await db.remove_scenario_by_name(scenario_name) ScenarioManagement.remove_files_by_scenario(scenario_name) except Exception as e: @@ -433,15 +432,13 @@ async def get_scenarios( Returns: dict: A list of scenarios and the currently running scenario. """ - from nebula.controller.database import get_all_scenarios_and_check_completed, get_running_scenario - try: - scenarios = await get_all_scenarios_and_check_completed(username=user, role=role) + scenarios = await db.get_all_scenarios_and_check_completed(username=user, role=role) if role == "admin": - scenario_running = await get_running_scenario() + scenario_running = await db.get_running_scenario() else: - scenario_running = await get_running_scenario(username=user) + scenario_running = await db.get_running_scenario(username=user) except Exception as e: logging.exception(f"Error obtaining scenarios: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -474,10 +471,8 @@ async def update_scenario( Returns: dict: A message confirming the update. """ - from nebula.controller.database import scenario_update_record - try: - await scenario_update_record(scenario_name, start_time, end_time, scenario, status, username) + await db.scenario_update_record(scenario_name, start_time, end_time, scenario, status, username) except Exception as e: logging.exception(f"Error updating scenario {scenario_name}: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -499,13 +494,11 @@ async def set_scenario_status_to_finished( Returns: dict: A message confirming the operation. """ - from nebula.controller.database import scenario_set_all_status_to_finished, scenario_set_status_to_finished - try: if all: - await scenario_set_all_status_to_finished() + await db.scenario_set_all_status_to_finished() else: - await scenario_set_status_to_finished(scenario_name) + await db.scenario_set_status_to_finished(scenario_name) except Exception as e: logging.exception(f"Error setting scenario {scenario_name} to finished: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -514,7 +507,7 @@ async def set_scenario_status_to_finished( @app.get("/scenarios/running") -async def get_running_scenario(get_all: bool = False): +async def get_running_scenario_endpoint(get_all: bool = False): """ Retrieves the currently running scenario(s). @@ -524,18 +517,15 @@ async def get_running_scenario(get_all: bool = False): Returns: dict or list: Running scenario(s) information. """ - from nebula.controller.database import get_running_scenario - try: - return await get_running_scenario(get_all=get_all) + return await db.get_running_scenario(get_all=get_all) except Exception as e: logging.exception(f"Error obtaining running scenario: {e}") raise HTTPException(status_code=500, detail="Internal server error") -@app.get("/scenarios/check/{user}/{role}/{scenario_name}") +@app.get("/scenarios/check/{role}/{scenario_name}") async def check_scenario( - user: Annotated[str, Path(regex="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50, description="Valid username")], role: Annotated[str, Path(regex="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50, description="Valid role")], scenario_name: Annotated[ str, Path(regex="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50, description="Valid scenario name") @@ -551,10 +541,8 @@ async def check_scenario( Returns: dict: Whether the scenario is allowed for the role. """ - from nebula.controller.database import check_scenario_with_role - try: - allowed = await check_scenario_with_role(role, scenario_name, user) + allowed = await db.check_scenario_with_role(role, scenario_name) return {"allowed": allowed} except Exception as e: logging.exception(f"Error checking scenario with role: {e}") @@ -562,7 +550,7 @@ async def check_scenario( @app.get("/scenarios/{scenario_name}") -async def get_scenario_by_name( +async def get_scenario_by_name_endpoint( scenario_name: Annotated[ str, Path(regex="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50, description="Valid scenario name") ], @@ -576,10 +564,8 @@ async def get_scenario_by_name( Returns: dict: The scenario data. """ - from nebula.controller.database import get_scenario_by_name - try: - scenario = await get_scenario_by_name(scenario_name) + scenario = await db.get_scenario_by_name(scenario_name) except Exception as e: logging.exception(f"Error obtaining scenario {scenario_name}: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -588,7 +574,7 @@ async def get_scenario_by_name( @app.get("/nodes/{scenario_name}") -async def list_nodes_by_scenario_name( +async def list_nodes_by_scenario_name_endpoint( scenario_name: Annotated[ str, Path(regex="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50, description="Valid scenario name") ], @@ -602,10 +588,8 @@ async def list_nodes_by_scenario_name( Returns: list: List of nodes. """ - from nebula.controller.database import list_nodes_by_scenario_name - try: - nodes = await list_nodes_by_scenario_name(scenario_name) + nodes = await db.list_nodes_by_scenario_name(scenario_name) except Exception as e: logging.exception(f"Error obtaining nodes: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -631,13 +615,11 @@ async def update_nodes( Returns: dict: Confirmation or response from the frontend. """ - from nebula.controller.database import update_node_record - try: config = await request.json() timestamp = datetime.datetime.now() # Update the node in database - await update_node_record( + await db.update_node_record( str(config["device_args"]["uid"]), str(config["device_args"]["idx"]), str(config["network_args"]["ip"]), @@ -708,7 +690,7 @@ async def node_done( @app.post("/nodes/remove") -async def remove_nodes_by_scenario_name(scenario_name: str = Body(..., embed=True)): +async def remove_nodes_by_scenario_name_endpoint(scenario_name: str = Body(..., embed=True)): """ Endpoint to remove all nodes associated with a scenario. @@ -717,10 +699,8 @@ async def remove_nodes_by_scenario_name(scenario_name: str = Body(..., embed=Tru Returns a success message or an error if something goes wrong. """ - from nebula.controller.database import remove_nodes_by_scenario_name - try: - await remove_nodes_by_scenario_name(scenario_name) + await db.remove_nodes_by_scenario_name(scenario_name) except Exception as e: logging.exception(f"Error removing nodes: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -737,10 +717,8 @@ async def get_notes_by_scenario_name( """ Endpoint to retrieve notes associated with a scenario. """ - from nebula.controller.database import get_notes - try: - notes_record = await get_notes(scenario_name) + notes_record = await db.get_notes(scenario_name) if notes_record is not None: notes_record = dict(notes_record.items()) @@ -763,10 +741,8 @@ async def update_notes_by_scenario_name(scenario_name: str = Body(..., embed=Tru Returns a success message or an error if something goes wrong. """ - from nebula.controller.database import save_notes - try: - await save_notes(scenario_name, notes) + await db.save_notes(scenario_name, notes) except Exception as e: logging.exception(f"Error updating notes: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -775,7 +751,7 @@ async def update_notes_by_scenario_name(scenario_name: str = Body(..., embed=Tru @app.post("/notes/remove") -async def remove_notes_by_scenario_name(scenario_name: str = Body(..., embed=True)): +async def remove_notes_by_scenario_name_endpoint(scenario_name: str = Body(..., embed=True)): """ Endpoint to remove notes associated with a scenario. @@ -784,10 +760,8 @@ async def remove_notes_by_scenario_name(scenario_name: str = Body(..., embed=Tru Returns a success message or an error if something goes wrong. """ - from nebula.controller.database import remove_note - try: - await remove_note(scenario_name) + await db.remove_note(scenario_name) except Exception as e: logging.exception(f"Error removing notes: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -805,10 +779,8 @@ async def list_users_controller(all_info: bool = False): Returns a list of users or raises an HTTPException on error. """ - from nebula.controller.database import list_users - try: - user_list = await list_users(all_info) + user_list = await db.list_users(all_info) if all_info: # Convert each asyncpg.Record to a dictionary so that it is JSON serializable. user_list = [dict(user) for user in user_list] @@ -818,7 +790,7 @@ async def list_users_controller(all_info: bool = False): @app.get("/user/{scenario_name}") -async def get_user_by_scenario_name( +async def get_user_by_scenario_name_endpoint( scenario_name: Annotated[ str, Path(regex="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50, description="Valid scenario name") ], @@ -831,10 +803,8 @@ async def get_user_by_scenario_name( Returns user info or raises an HTTPException on error. """ - from nebula.controller.database import get_user_by_scenario_name - try: - user = await get_user_by_scenario_name(scenario_name) + user = await db.get_user_by_scenario_name(scenario_name) except Exception as e: logging.exception(f"Error obtaining user {user}: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -995,11 +965,11 @@ async def get_physical_scenario_state(scenario_name: str): } """ # 1) Retrieve scenario metadata and node list from the DB - scenario = await get_scenario_by_name(scenario_name) + scenario = await db.get_scenario_by_name(scenario_name) if not scenario: raise HTTPException(status_code=404, detail="Scenario not found") - nodes = await list_nodes_by_scenario_name(scenario_name) + nodes = await db.list_nodes_by_scenario_name(scenario_name) if not nodes: raise HTTPException(status_code=404, detail="No nodes found for scenario") @@ -1036,10 +1006,8 @@ async def add_user_controller(user: str = Body(...), password: str = Body(...), Returns a success message or an error if the user could not be added. """ - from nebula.controller.database import add_user - try: - await add_user(user, password, role) + await db.add_user(user, password, role) return {"detail": "User added successfully"} except Exception as e: logging.exception(f"Error adding user: {e}") @@ -1056,10 +1024,8 @@ async def remove_user_controller(user: str = Body(..., embed=True)): Returns a success message if the user is deleted, or an HTTP error if an exception occurs. """ - from nebula.controller.database import delete_user_from_db - try: - await delete_user_from_db(user) + await db.delete_user_from_db(user) return {"detail": "User deleted successfully"} except Exception as e: logging.exception(f"Error deleting user: {e}") @@ -1078,10 +1044,8 @@ async def update_user_controller(user: str = Body(...), password: str = Body(... Returns a success message if the user is updated, or an HTTP error if an exception occurs. """ - from nebula.controller.database import update_user - try: - await update_user(user, password, role) + await db.update_user(user, password, role) return {"detail": "User updated successfully"} except Exception as e: logging.exception(f"Error updating user: {e}") @@ -1099,12 +1063,10 @@ async def verify_user_controller(user: str = Body(...), password: str = Body(... Returns the user role on success or raises an error on failure. """ - from nebula.controller.database import get_user_info, list_users, verify - try: user_submitted = user.upper() - if (await list_users() and await verify(user_submitted, password)): - user_info = await get_user_info(user_submitted) + if (await db.list_users() and await db.verify(user_submitted, password)): + user_info = await db.get_user_info(user_submitted) return {"user": user_submitted, "role": user_info[2]} else: raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED) diff --git a/nebula/controller/start_services.sh b/nebula/controller/start_services.sh index fa0eaf031..43043e99c 100644 --- a/nebula/controller/start_services.sh +++ b/nebula/controller/start_services.sh @@ -12,6 +12,6 @@ echo "path $(pwd)" NEBULA_SOCK=nebula.sock echo "Starting Gunicorn..." -uvicorn nebula.controller.controller:app --host 0.0.0.0 --port $NEBULA_CONTROLLER_PORT --log-level debug --proxy-headers --forwarded-allow-ips "*" & +uvicorn nebula.controller.hub:app --host 0.0.0.0 --port $NEBULA_CONTROLLER_PORT --log-level debug --proxy-headers --forwarded-allow-ips "*" & tail -f /dev/null diff --git a/nebula/database/Dockerfile b/nebula/database/adapters/postgress/docker/Dockerfile similarity index 93% rename from nebula/database/Dockerfile rename to nebula/database/adapters/postgress/docker/Dockerfile index 04ed953e9..b1b1d26b6 100644 --- a/nebula/database/Dockerfile +++ b/nebula/database/adapters/postgress/docker/Dockerfile @@ -4,7 +4,7 @@ FROM postgres:17.5-alpine3.22 RUN mv /usr/local/bin/docker-entrypoint.sh /usr/local/bin/docker-entrypoint.sh.orig # Copy SQL init file and custom entrypoint script -COPY /nebula/database/docker-entrypoint.sh /usr/local/bin/docker-entrypoint.sh +COPY /nebula/database/docker/docker-entrypoint.sh /usr/local/bin/docker-entrypoint.sh RUN chmod +x /usr/local/bin/docker-entrypoint.sh # Install Python 3.11.7 from source diff --git a/nebula/database/docker-entrypoint.sh b/nebula/database/adapters/postgress/docker/docker-entrypoint.sh similarity index 100% rename from nebula/database/docker-entrypoint.sh rename to nebula/database/adapters/postgress/docker/docker-entrypoint.sh diff --git a/nebula/database/init-configs.sql b/nebula/database/adapters/postgress/docker/init-configs.sql similarity index 100% rename from nebula/database/init-configs.sql rename to nebula/database/adapters/postgress/docker/init-configs.sql diff --git a/nebula/database/adapters/postgress/postgress.py b/nebula/database/adapters/postgress/postgress.py new file mode 100755 index 000000000..3de289780 --- /dev/null +++ b/nebula/database/adapters/postgress/postgress.py @@ -0,0 +1,655 @@ +import logging +import os +import datetime +import json +import asyncpg +import asyncio + +from passlib.context import CryptContext + +from nebula.database.database_adapter_interface import DatabaseAdapter + +# --- Configuration --- +# Use environment variables for database credentials from the Docker Compose file +DATABASE_URL = f"postgresql://{os.environ.get('DB_USER')}:{os.environ.get('DB_PASSWORD')}@{os.environ.get('DB_HOST')}:{os.environ.get('DB_PORT')}/nebula" + +# Password hashing context (using Argon2) +pwd_context = CryptContext(schemes=["argon2"], deprecated="auto") + +# Asynchronous lock for node updates +_node_lock = asyncio.Lock() + + +class PostgresDB(DatabaseAdapter): + """ + PostgreSQL implementation of the Database interface. + """ + def __init__(self): + self.pool = None + + async def init_db_pool(self): + """ + Initializes the asynchronous PostgreSQL connection pool. + This should be called once when the application starts. + Retries connection on failure to handle race conditions during startup. + """ + if self.pool is None: + attempts = 10 + for attempt in range(attempts): + try: + self.pool = await asyncpg.create_pool( + dsn=DATABASE_URL, + min_size=5, # Minimum number of connections in the pool + max_size=20, # Maximum number of connections in the pool + ) + logging.info("Database connection pool successfully created.") + return + except (ConnectionRefusedError, asyncpg.exceptions.CannotConnectNowError) as e: + if attempt < attempts - 1: + logging.warning( + f"Database connection failed. Attempt {attempt + 1}/{attempts}. Retrying in 5 seconds... " + f"Error: {e}" + ) + await asyncio.sleep(5) + else: + logging.critical( + f"Failed to create database connection pool after {attempts} attempts: {e}", exc_info=True + ) + raise + except Exception as e: + logging.critical( + f"An unexpected error occurred while creating database connection pool: {e}", exc_info=True + ) + raise + + async def close_db_pool(self): + """ + Closes the asynchronous PostgreSQL connection pool. + This should be called once when the application shuts down gracefully. + """ + if self.pool: + await self.pool.close() + logging.info("Database connection pool closed.") + + + # --- User Management Functions --- + + async def insert_default_admin(self): + """ + Inserts a default 'ADMIN' user into the database with a hashed password. + The password must be provided via the ADMIN_PASSWORD environment variable. + """ + admin_password = os.environ.get("NEBULA_ADMIN_PASSWORD") + + hashed_password = pwd_context.hash(admin_password) + + query = """ + INSERT INTO users ("user", password, role) + VALUES ($1, $2, $3) + ON CONFLICT ("user") DO NOTHING; + """ + try: + async with self.pool.acquire() as conn: + await conn.execute(query, "ADMIN", hashed_password, "admin") + logging.info("Default admin user inserted (or already exists).") + except Exception as e: + logging.error(f"Failed to insert default admin user: {e}", exc_info=True) + + async def list_users(self, all_info=False): + """ + Retrieves a list of users from the users database. + """ + async with self.pool.acquire() as conn: + result = await conn.fetch("SELECT * FROM users") + + if not all_info: + result = [user["user"] for user in result] + + return result + + + async def get_user_info(self, user): + """ + Fetches detailed information for a specific user from the users database. + """ + async with self.pool.acquire() as conn: + return await conn.fetchrow('SELECT * FROM users WHERE "user" = $1', user) + + + async def verify(self, user, password): + """ + Verifies whether the provided password matches the stored hashed password for a user. + """ + async with self.pool.acquire() as conn: + result = await conn.fetchrow('SELECT password FROM users WHERE "user" = $1', user) + if result: + try: + return pwd_context.verify(password, result[0]) + except Exception: + # Catch more general exceptions during verification to be safe + logging.error(f"Error during password verification for user {user}", exc_info=True) + return False + return False + + + async def verify_hash_algorithm(self, user): + """ + Checks if the stored password hash for a user uses a supported Argon2 algorithm. + """ + user = user.upper() + argon2_prefixes = ("$argon2i$", "$argon2id$") + async with self.pool.acquire() as conn: + result = await conn.fetchrow('SELECT password FROM users WHERE "user" = $1', user) + if result: + password_hash = result["password"] + return password_hash.startswith(argon2_prefixes) + return False + + + async def delete_user_from_db(self, user): + """ + Deletes a user record from the users database. + """ + async with self.pool.acquire() as conn: + await conn.execute('DELETE FROM users WHERE "user" = $1', user) + + + async def add_user(self, user, password, role): + """ + Adds a new user to the users database with a hashed password. + """ + hashed_password = pwd_context.hash(password) + async with self.pool.acquire() as conn: + await conn.execute( + 'INSERT INTO users ("user", password, role) VALUES ($1, $2, $3)', + user.upper(), hashed_password, role, + ) + + + async def update_user(self, user, password, role): + """ + Updates the password and role of an existing user in the users database. + """ + hashed_password = pwd_context.hash(password) + async with self.pool.acquire() as conn: + await conn.execute( + 'UPDATE users SET password = $1, role = $2 WHERE "user" = $3', + hashed_password, role, user.upper(), + ) + + # --- Node Management Functions --- + + async def list_nodes(self, scenario_name=None, sort_by="idx"): + """ + Retrieves a list of nodes from the nodes database, optionally filtered by scenario and sorted. + """ + # Validate sort_by to prevent SQL injection + allowed_sort_fields = ["uid", "idx", "ip", "port", "role", "timestamp", "federation", "round"] + if sort_by not in allowed_sort_fields: + sort_by = "idx" # Default to a safe field + + try: + async with self.pool.acquire() as conn: + if scenario_name: + # Using f-string for column names is generally safe if validated as above + command = f"SELECT * FROM nodes WHERE scenario = $1 ORDER BY {sort_by};" + result = await conn.fetch(command, scenario_name) + else: + command = f"SELECT * FROM nodes ORDER BY {sort_by};" + result = await conn.fetch(command) + return result + except asyncpg.PostgresError as e: + logging.error(f"Error occurred while listing nodes: {e}") + return None + + + async def list_nodes_by_scenario_name(self, scenario_name): + """ + Fetches all nodes associated with a specific scenario, ordered by their index as integers. + """ + try: + async with self.pool.acquire() as conn: + command = "SELECT * FROM nodes WHERE scenario = $1 ORDER BY CAST(idx AS INTEGER) ASC;" + result = await conn.fetch(command, scenario_name) + return [dict(record) for record in result] + except Exception as e: + logging.error(f"Error occurred while listing nodes by scenario name: {e}") + return None + + + async def update_node_record( + self, node_uid, idx, ip, port, role, neighbors, latitude, longitude, + timestamp, federation, federation_round, scenario, run_hash, malicious, + ): + """ + Inserts or updates a node record in the database for a given scenario, ensuring thread-safe access. + """ + async with _node_lock: + async with self.pool.acquire() as conn: + try: + async with conn.transaction(): + result = await conn.fetchrow( + "SELECT * FROM nodes WHERE uid = $1 AND scenario = $2 FOR UPDATE;", + node_uid, scenario + ) + + if result is None: + # Insert new node + await conn.execute( + """ + INSERT INTO nodes (uid, idx, ip, port, role, neighbors, latitude, longitude, + timestamp, federation, round, scenario, hash, malicious) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14); + """, + node_uid, idx, ip, port, role, neighbors, latitude, longitude, + timestamp, federation, federation_round, scenario, run_hash, malicious, + ) + else: + # Update existing node + await conn.execute( + """ + UPDATE nodes SET idx = $1, ip = $2, port = $3, role = $4, neighbors = $5, + latitude = $6, longitude = $7, timestamp = $8, federation = $9, round = $10, + hash = $11, malicious = $12 + WHERE uid = $13 AND scenario = $14; + """, + idx, ip, port, role, neighbors, latitude, longitude, + timestamp, federation, federation_round, run_hash, malicious, + node_uid, scenario, + ) + + updated_row = await conn.fetchrow("SELECT * from nodes WHERE uid = $1 AND scenario = $2;", node_uid, scenario) + return dict(updated_row) if updated_row else None + except asyncpg.PostgresError as e: + logging.error(f"Database error during node record update: {e}", exc_info=True) + return None + + + async def remove_all_nodes(self): + """ + Deletes all node records from the nodes database. + """ + async with self.pool.acquire() as conn: + await conn.execute("TRUNCATE nodes CASCADE;") # Use CASCADE if there are foreign key dependencies + + + async def remove_nodes_by_scenario_name(self, scenario_name): + """ + Deletes all nodes associated with a specific scenario from the database. + """ + async with self.pool.acquire() as conn: + await conn.execute("DELETE FROM nodes WHERE scenario = $1;", scenario_name) + + # --- Scenario Management Functions --- + + async def get_all_scenarios(self, username, role, sort_by="start_time"): + """ + Retrieves all scenarios from the database, accessing fields from the 'config' (JSONB) column + and direct columns. Filters by user role and sorts by the specified field. + """ + allowed_sort_fields = ["start_time", "title", "username", "status", "name"] + if sort_by not in allowed_sort_fields: + sort_by = "start_time" + + # Determine the ORDER BY clause based on sort_by + if sort_by == "start_time": + order_by_clause = """ + ORDER BY + CASE + WHEN start_time IS NULL OR start_time = '' THEN 1 + ELSE 0 + END, + to_timestamp(start_time, 'DD/MM/YYYY HH24:MI:SS') DESC + """ + elif sort_by in ["title", "model", "dataset", "rounds"]: # These are inside config JSONB + order_by_clause = f"ORDER BY config->>'{sort_by}'" + else: # For direct table columns like name, username, status + order_by_clause = f"ORDER BY {sort_by}" + + async with self.pool.acquire() as conn: + # Select direct columns and relevant fields from config JSONB + command = """ + SELECT + name, + username, + status, + start_time, + end_time, + config->>'title' AS title, + config->>'model' AS model, + config->>'dataset' AS dataset, + config->>'rounds' AS rounds, + config -- return the full config JSONB + FROM scenarios + """ + params = [] + + if role != "admin": + command += " WHERE username = $1" # username is a direct column now + params.append(username) + + full_command = f"{command} {order_by_clause};" + return await conn.fetch(full_command, *params) + + + async def get_all_scenarios_and_check_completed(self, username, role, sort_by="start_time"): + """ + Retrieves all scenarios, sorts them, and updates the status if necessary. + Returns a list of dictionaries, where each dictionary represents a scenario. + """ + # Safe list of allowed sorting fields to prevent SQL injection. + allowed_sort_fields = ["start_time", "title", "username", "status", "name"] + if sort_by not in allowed_sort_fields: + sort_by = "start_time" # Safe default value + + # Building the ORDER BY clause + if sort_by == "start_time": + order_by_clause = """ + ORDER BY + CASE + WHEN start_time IS NULL OR start_time = '' THEN 1 + ELSE 0 + END, + to_timestamp(start_time, 'DD/MM/YYYY HH24:MI:SS') DESC + """ + elif sort_by in ["title", "model", "dataset", "rounds"]: # These are inside config JSONB + order_by_clause = f"ORDER BY config->>'{sort_by}'" + else: # For direct table columns like name, username, status + order_by_clause = f"ORDER BY {sort_by}" + + async with self.pool.acquire() as conn: + # Base query that extracts fields from the JSONB using the ->> operator + command = f""" + SELECT + name, + username, + status, + start_time, + end_time, + config->>'title' AS title, + config->>'model' AS model, + config->>'dataset' AS dataset, + config->>'rounds' AS rounds, + config -- Return the full config object + FROM scenarios + """ + params = [] + if role != "admin": + command += " WHERE username = $1" # username is a direct column + params.append(username) + + command += f" {order_by_clause};" + + result_dicts = await conn.fetch(command, *params) + + scenarios_to_return = [dict(s) for s in result_dicts] + + re_fetch_required = False + for scenario in scenarios_to_return: + if scenario["status"] == "running": + if await self.check_scenario_federation_completed(scenario["name"]): + await self.scenario_set_status_to_completed(scenario["name"]) + re_fetch_required = True + break + + if re_fetch_required: + # Recursively call to get fresh data after status update + return await self.get_all_scenarios_and_check_completed(username, role, sort_by) + + return scenarios_to_return + + + async def scenario_update_record(self, name, start_time, end_time, scenario_config, status, username): + """ + Inserts or updates a scenario record using the PostgreSQL "UPSERT" pattern. + All configuration is saved in the 'config' column of type JSONB. + Direct columns (name, start_time, end_time, username, status) are also handled. + """ + # Ensure scenario_config is a dictionary before dumping to JSON + if not isinstance(scenario_config, dict): + try: + scenario_config = json.loads(scenario_config) + except (json.JSONDecodeError, TypeError): + logging.error("scenario_config is not a valid JSON string or dict.") + return + + command = """ + INSERT INTO scenarios (name, start_time, end_time, username, status, config) + VALUES ($1, $2, $3, $4, $5, $6::jsonb) + ON CONFLICT (name) DO UPDATE SET + start_time = EXCLUDED.start_time, + end_time = EXCLUDED.end_time, + username = EXCLUDED.username, + status = EXCLUDED.status, + config = scenarios.config || EXCLUDED.config; -- Merge JSONB + """ + async with self.pool.acquire() as conn: + await conn.execute(command, name, start_time, end_time, username, status, json.dumps(scenario_config)) + + + async def scenario_set_all_status_to_finished(self): + """ + Sets the status of all 'running' scenarios to 'finished' + and updates their 'end_time' (both in the direct column and within JSONB). + """ + current_time = datetime.datetime.now().strftime('%Y/%m/%d %H:%M:%S') # Consistent format + command = """ + UPDATE scenarios + SET + status = 'finished', + end_time = $1, + config = jsonb_set(config, '{status}', '"finished"') || + jsonb_set(config, '{end_time}', $2::jsonb) + WHERE status = 'running'; + """ + async with self.pool.acquire() as conn: + await conn.execute(command, current_time, json.dumps(current_time)) + + + async def scenario_set_status_to_finished(self, scenario_name): + """ + Sets the status of a specific scenario to 'finished' and updates its 'end_time'. + Updates both the direct columns and the JSONB 'config'. + """ + current_time = datetime.datetime.now().strftime('%Y/%m/%d %H:%M:%S') # Consistent format + command = """ + UPDATE scenarios + SET + status = 'finished', + end_time = $1, + config = jsonb_set( + jsonb_set(config, '{status}', '"finished"'), + '{end_time}', $2::jsonb + ) + WHERE name = $3; + """ + async with self.pool.acquire() as conn: + await conn.execute(command, current_time, json.dumps(current_time), scenario_name) + + + async def scenario_set_status_to_completed(self, scenario_name): + """ + Sets the status of a specific scenario to 'completed'. + Updates both the direct column and the JSONB 'config'. + """ + command = """ + UPDATE scenarios + SET + status = 'completed', + config = jsonb_set(config, '{status}', '"completed"') + WHERE name = $1; + """ + async with self.pool.acquire() as conn: + await conn.execute(command, scenario_name) + + + async def get_running_scenario(self, username=None, get_all=False): + """ + Retrieves scenarios with a 'running' status, optionally filtered by user. + Returns full scenario record (including direct columns and config JSONB). + """ + async with self.pool.acquire() as conn: + params = ["running"] + # Select all columns to get both direct and config data + command = "SELECT name, username, status, start_time, end_time, config FROM scenarios WHERE status = $1" + + if username: + command += " AND username = $2" + params.append(username) + + if get_all: + result = [dict(row) for row in await conn.fetch(command, *params)] # Convert records to dicts + else: + result_row = await conn.fetchrow(command, *params) + result = dict(result_row) if result_row else None + return result + + + async def get_completed_scenario(self): + """ + Retrieves a single scenario with a 'completed' status. + Returns full scenario record (including direct columns and config JSONB). + """ + async with self.pool.acquire() as conn: + command = "SELECT name, username, status, start_time, end_time, config FROM scenarios WHERE status = $1;" + result_row = await conn.fetchrow(command, "completed") + return dict(result_row) if result_row else None + + + async def get_scenario_by_name(self, scenario_name): + """ + Retrieves the complete record of a scenario by its name. + """ + async with self.pool.acquire() as conn: + result_row = await conn.fetchrow("SELECT name, start_time, end_time, username, status, config FROM scenarios WHERE name = $1;", scenario_name) + + result = dict(result_row) if result_row else None + + if result and result.get('config'): + # Assuming 'config' is a JSON string from the DB, so we parse it + # It might already be a dict if asyncpg handles JSONB conversion automatically + config_data = result['config'] + if isinstance(config_data, str): + try: + config_data = json.loads(config_data) + except json.JSONDecodeError: + config_data = {} + + # Extract the 'scenario_title' and add it as a top-level key + result['title'] = config_data.get('scenario_title') + result['description'] = config_data.get('description') + + return result + + + async def get_user_by_scenario_name(self, scenario_name): + """ + Retrieves the username associated with a scenario (from the direct 'username' column). + """ + async with self.pool.acquire() as conn: + return await conn.fetchval("SELECT username FROM scenarios WHERE name = $1;", scenario_name) + + + async def remove_scenario_by_name(self, scenario_name): + """ + Delete a scenario from the database by its unique name. + """ + try: + async with self.pool.acquire() as conn: + await conn.execute("DELETE FROM scenarios WHERE name = $1;", scenario_name) + logging.info(f"Scenario '{scenario_name}' successfully removed.") + except asyncpg.PostgresError as e: + logging.error(f"Error occurred while deleting scenario '{scenario_name}': {e}") + + + async def check_scenario_federation_completed(self, scenario_name): + """ + Check if all nodes in a given scenario have completed the required federation rounds. + """ + try: + async with self.pool.acquire() as conn: + # Retrieve the total rounds for the scenario from the 'config' JSONB column + scenario_rounds_str = await conn.fetchval("SELECT config->>'rounds' AS rounds FROM scenarios WHERE name = $1;", scenario_name) + + if not scenario_rounds_str: + logging.warning(f"Scenario '{scenario_name}' not found or 'rounds' not defined.") + return False + + # Ensure total_rounds is an integer for comparison + try: + total_rounds = int(scenario_rounds_str) + except (ValueError, TypeError): + logging.error(f"Invalid 'rounds' value for scenario '{scenario_name}': {scenario_rounds_str}") + return False + + # Fetch the current round progress of all nodes in that scenario + nodes = await conn.fetch("SELECT round FROM nodes WHERE scenario = $1;", scenario_name) + + if not nodes: + logging.info(f"No nodes found for scenario '{scenario_name}'. Federation not considered completed.") + return False + + # Check if all nodes have completed the total rounds + return all(int(node["round"]) >= total_rounds for node in nodes) + + except asyncpg.PostgresError as e: + logging.error(f"PostgreSQL error during check_scenario_federation_completed for '{scenario_name}': {e}") + return False + except ValueError as e: + logging.error(f"Data error during check_scenario_federation_completed for '{scenario_name}': {e}") + return False + + + async def check_scenario_with_role(self, role, scenario_name, current_username=None): + """ + Verify if a scenario exists that the user with the given role and username can access. + """ + scenario_info = await self.get_scenario_by_name(scenario_name) + + if not scenario_info: + return False # Scenario does not exist + + if role == "admin": + return True # Admins can access any existing scenario + + if current_username is None: + logging.warning( + "check_scenario_with_role called for non-admin role without current_username." + ) + return False + + return scenario_info.get("username") == current_username + + # --- Notes Management Functions --- + + async def save_notes(self, scenario, notes): + """ + Save or update notes associated with a specific scenario. + """ + try: + async with self.pool.acquire() as conn: + await conn.execute( + """ + INSERT INTO notes (scenario, scenario_notes) VALUES ($1, $2) + ON CONFLICT(scenario) DO UPDATE SET scenario_notes = EXCLUDED.scenario_notes; + """, + scenario, notes, + ) + except asyncpg.PostgresError as e: + logging.error(f"PostgreSQL error during save_notes: {e}") + + + async def get_notes(self, scenario): + """ + Retrieve notes associated with a specific scenario. + """ + async with self.pool.acquire() as conn: + return await conn.fetchrow("SELECT * FROM notes WHERE scenario = $1;", scenario) + + + async def remove_note(self, scenario): + """ + Delete the note associated with a specific scenario. + """ + async with self.pool.acquire() as conn: + await conn.execute("DELETE FROM notes WHERE scenario = $1;", scenario) From 21c46f0c8a474ffa7ca9a53cea1ef613f3e27f1e Mon Sep 17 00:00:00 2001 From: FerTV Date: Wed, 23 Jul 2025 13:53:12 +0200 Subject: [PATCH 02/14] feature: database adapter created --- nebula/database/database_adapter_factory.py | 17 ++ nebula/database/database_adapter_interface.py | 174 ++++++++++++++++++ 2 files changed, 191 insertions(+) create mode 100644 nebula/database/database_adapter_factory.py create mode 100644 nebula/database/database_adapter_interface.py diff --git a/nebula/database/database_adapter_factory.py b/nebula/database/database_adapter_factory.py new file mode 100644 index 000000000..ddbe443c1 --- /dev/null +++ b/nebula/database/database_adapter_factory.py @@ -0,0 +1,17 @@ +from nebula.database.adapters.postgress.postgress import PostgresDB +from nebula.database.database_adapter_interface import DatabaseAdapter + +class DatabaseAdapterException(Exception): + pass + +def factory_database_adapter(database_adapter: str) -> DatabaseAdapter: + + ADAPTERS = { + "PostgresDB": PostgresDB + } + + db_adapter = ADAPTERS.get(database_adapter, None) + if db_adapter: + return db_adapter() + else: + raise DatabaseAdapterException(f"Database Adapter \"{database_adapter}\" not supported") diff --git a/nebula/database/database_adapter_interface.py b/nebula/database/database_adapter_interface.py new file mode 100644 index 000000000..d527bb467 --- /dev/null +++ b/nebula/database/database_adapter_interface.py @@ -0,0 +1,174 @@ +from abc import ABC, abstractmethod + + +class DatabaseAdapter(ABC): + """ + Abstract base class for database operations. + Defines a common interface for interacting with different database systems. + """ + + @abstractmethod + async def init_db_pool(self): + """Initializes the database connection pool.""" + raise NotImplementedError + + @abstractmethod + async def close_db_pool(self): + """Closes the database connection pool.""" + raise NotImplementedError + + # --- User Management Functions --- + + @abstractmethod + async def insert_default_admin(self): + """Inserts a default admin user.""" + raise NotImplementedError + + @abstractmethod + async def list_users(self, all_info=False): + """Retrieves a list of users.""" + raise NotImplementedError + + @abstractmethod + async def get_user_info(self, user): + """Fetches detailed information for a specific user.""" + raise NotImplementedError + + @abstractmethod + async def verify(self, user, password): + """Verifies user credentials.""" + raise NotImplementedError + + @abstractmethod + async def verify_hash_algorithm(self, user): + """Checks the password hash algorithm for a user.""" + raise NotImplementedError + + @abstractmethod + async def delete_user_from_db(self, user): + """Deletes a user from the database.""" + raise NotImplementedError + + @abstractmethod + async def add_user(self, user, password, role): + """Adds a new user.""" + raise NotImplementedError + + @abstractmethod + async def update_user(self, user, password, role): + """Updates an existing user.""" + raise NotImplementedError + + # --- Node Management Functions --- + + @abstractmethod + async def list_nodes(self, scenario_name=None, sort_by="idx"): + """Retrieves a list of nodes.""" + raise NotImplementedError + + @abstractmethod + async def list_nodes_by_scenario_name(self, scenario_name): + """Fetches all nodes for a specific scenario.""" + raise NotImplementedError + + @abstractmethod + async def update_node_record( + self, node_uid, idx, ip, port, role, neighbors, latitude, longitude, + timestamp, federation, federation_round, scenario, run_hash, malicious, + ): + """Inserts or updates a node record.""" + raise NotImplementedError + + @abstractmethod + async def remove_all_nodes(self): + """Deletes all node records.""" + raise NotImplementedError + + @abstractmethod + async def remove_nodes_by_scenario_name(self, scenario_name): + """Deletes all nodes for a specific scenario.""" + raise NotImplementedError + + # --- Scenario Management Functions --- + + @abstractmethod + async def get_all_scenarios(self, username, role, sort_by="start_time"): + """Retrieves all scenarios.""" + raise NotImplementedError + + @abstractmethod + async def get_all_scenarios_and_check_completed(self, username, role, sort_by="start_time"): + """Retrieves all scenarios and checks for completion.""" + raise NotImplementedError + + @abstractmethod + async def scenario_update_record(self, name, start_time, end_time, scenario_config, status, username): + """Inserts or updates a scenario record.""" + raise NotImplementedError + + @abstractmethod + async def scenario_set_all_status_to_finished(self): + """Sets the status of all running scenarios to 'finished'.""" + raise NotImplementedError + + @abstractmethod + async def scenario_set_status_to_finished(self, scenario_name): + """Sets the status of a specific scenario to 'finished'.""" + raise NotImplementedError + + @abstractmethod + async def scenario_set_status_to_completed(self, scenario_name): + """Sets the status of a specific scenario to 'completed'.""" + raise NotImplementedError + + @abstractmethod + async def get_running_scenario(self, username=None, get_all=False): + """Retrieves running scenarios.""" + raise NotImplementedError + + @abstractmethod + async def get_completed_scenario(self): + """Retrieves a completed scenario.""" + raise NotImplementedError + + @abstractmethod + async def get_scenario_by_name(self, scenario_name): + """Retrieves a scenario by its name.""" + raise NotImplementedError + + @abstractmethod + async def get_user_by_scenario_name(self, scenario_name): + """Retrieves the user associated with a scenario.""" + raise NotImplementedError + + @abstractmethod + async def remove_scenario_by_name(self, scenario_name): + """Deletes a scenario by its name.""" + raise NotImplementedError + + @abstractmethod + async def check_scenario_federation_completed(self, scenario_name): + """Checks if a scenario's federation is complete.""" + raise NotImplementedError + + @abstractmethod + async def check_scenario_with_role(self, role, scenario_name, current_username=None): + """Verifies if a user can access a scenario.""" + raise NotImplementedError + + # --- Notes Management Functions --- + + @abstractmethod + async def save_notes(self, scenario, notes): + """Saves or updates notes for a scenario.""" + raise NotImplementedError + + @abstractmethod + async def get_notes(self, scenario): + """Retrieves notes for a scenario.""" + raise NotImplementedError + + @abstractmethod + async def remove_note(self, scenario): + """Deletes the note for a scenario.""" + raise NotImplementedError From 6a5ebab3912570824ca13240b9021e75fa63b98e Mon Sep 17 00:00:00 2001 From: FerTV Date: Tue, 29 Jul 2025 12:20:27 +0200 Subject: [PATCH 03/14] feature: database api implemented --- Makefile | 2 +- app/deployer.py | 16 +- nebula/controller/hub.py | 250 +++++++------ .../adapters/postgress/docker/Dockerfile | 11 +- .../postgress/docker/docker-entrypoint.sh | 31 +- .../postgress/docker/init-configs.sql | 4 +- .../database/adapters/postgress/postgress.py | 2 +- nebula/database/database_api.py | 337 ++++++++++++++++++ nebula/frontend/app.py | 4 +- pyproject.toml | 9 +- 10 files changed, 526 insertions(+), 140 deletions(-) create mode 100644 nebula/database/database_api.py diff --git a/Makefile b/Makefile index 44a6cf390..6094db00a 100644 --- a/Makefile +++ b/Makefile @@ -76,7 +76,7 @@ update-dockers: ## Update docker images fi @echo "🐳 Building nebula-database docker image. Do you want to continue (overrides existing image)? (y/n)" @read ans; if [ "$${ans:-N}" = y ]; then \ - docker build -t nebula-database -f nebula/database/docker/Dockerfile .; \ + docker build -t nebula-database -f nebula/database/adapters/postgress/docker/Dockerfile .; \ docker build -t nebula-pgweb -f nebula/database/pgweb/Dockerfile .; \ else \ echo "Skipping nebula-database docker build."; \ diff --git a/app/deployer.py b/app/deployer.py index 7e46fe1a7..d04f74654 100644 --- a/app/deployer.py +++ b/app/deployer.py @@ -1032,6 +1032,12 @@ def run_database(self): "POSTGRES_USER": "nebula", "POSTGRES_PASSWORD": os.environ.get("POSTGRES_PASSWORD"), "POSTGRES_DB": "nebula", + "NEBULA_DATABASE_LOG": "/nebula/app/logs/database.log", + "DB_HOST": "localhost", + "DB_PORT": 5432, + "DB_USER": "nebula", + "DB_PASSWORD": os.environ.get("POSTGRES_PASSWORD"), + "NEBULA_ADMIN_PASSWORD": os.environ.get("NEBULA_ADMIN_PASSWORD") } host_sql_path = os.path.join(self.root_path, "nebula/database/adapters/postgress/docker/init-configs.sql") db_data_path = os.path.join(self.databases_dir, "postgres-data") @@ -1039,10 +1045,11 @@ def run_database(self): pg_host_config = client.api.create_host_config( binds=[ + f"{self.root_path}:/nebula", f"{host_sql_path}:/docker-entrypoint-initdb.d/init-configs.sql", f"{db_data_path}:/var/lib/postgresql/data", ], - port_bindings={5432: 5432}, + port_bindings={5432: 5432, 5051: 5051}, ) pg_networking_config = client.api.create_networking_config( {f"{network_name}": client.api.create_endpoint_config(ipv4_address=f"{base}.125")} @@ -1055,6 +1062,7 @@ def run_database(self): environment=pg_environment, host_config=pg_host_config, networking_config=pg_networking_config, + ports=[5432, 5051], ) client.api.start(pg_container) Deployer._add_container_to_metadata(pg_container_name) @@ -1117,11 +1125,7 @@ def run_controller(self): "NEBULA_CONTROLLER_PORT": self.controller_port, "NEBULA_CONTROLLER_HOST": self.controller_host, "NEBULA_FRONTEND_PORT": self.frontend_port, - "DB_HOST": self.get_container_name("nebula-database"), - "DB_PORT": 5432, - "DB_USER": "nebula", - "DB_PASSWORD": os.environ.get("POSTGRES_PASSWORD"), - "NEBULA_ADMIN_PASSWORD": os.environ.get("NEBULA_ADMIN_PASSWORD") + "NEBULA_DATABASE_API_URL": f"http://{self.get_container_name('nebula-database')}:5051" } volumes = ["/nebula", "/var/run/docker.sock"] diff --git a/nebula/controller/hub.py b/nebula/controller/hub.py index 628d61322..c170102bf 100755 --- a/nebula/controller/hub.py +++ b/nebula/controller/hub.py @@ -16,15 +16,11 @@ from fastapi import Body, FastAPI, Request, status, HTTPException, Path, File, UploadFile from fastapi.concurrency import asynccontextmanager -from nebula.database.database_adapter_factory import factory_database_adapter from nebula.controller.http_helpers import remote_get, remote_post_form from nebula.utils import DockerUtils - -# Get a database instance -# db = get_database() -#TODO review -db = factory_database_adapter("PostgresDB") +# URL for the database API +DATABASE_API_URL = os.environ.get("NEBULA_DATABASE_API_URL", "http://nebula-database:5051") # Setup controller logger @@ -111,22 +107,16 @@ def configure_logger(controller_log): async def lifespan(app: FastAPI): """ Application lifespan context manager. - - Initializes the database connection pool on startup. - - Configures logging. - - Cleans up resources like the database pool on shutdown. + - Configures logging on startup. """ # Code to run on startup controller_log: str = os.environ.get("NEBULA_CONTROLLER_LOG") configure_logger(controller_log) - # Initialize the database connection pool - await db.init_db_pool() - await db.insert_default_admin() - yield # Code to run on shutdown - await db.close_db_pool() + pass # Initialize FastAPI app outside the Controller class @@ -321,8 +311,6 @@ async def run_scenario( validate_physical_fields(scenario_data) - db_scenario = copy.deepcopy(scenario_data) - # Manager for the actual scenario scenarioManagement = ScenarioManagement(scenario_data, user) @@ -332,7 +320,6 @@ async def run_scenario( end_time="", scenario=scenario_data, status="running", - role=role, username=user, ) @@ -382,10 +369,12 @@ async def stop_scenario( ScenarioManagement.cleanup_scenario_containers() try: - if all: - await db.scenario_set_all_status_to_finished() - else: - await db.scenario_set_status_to_finished(scenario_name) + async with aiohttp.ClientSession() as session: + async with session.post( + f"{DATABASE_API_URL}/scenarios/stop", json={"scenario_name": scenario_name, "all": all} + ) as response: + if response.status != 200: + raise HTTPException(status_code=response.status, detail=await response.text()) except Exception as e: logging.exception(f"Error setting scenario {scenario_name} to finished: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -407,7 +396,12 @@ async def remove_scenario( from nebula.controller.scenarios import ScenarioManagement try: - await db.remove_scenario_by_name(scenario_name) + async with aiohttp.ClientSession() as session: + async with session.post( + f"{DATABASE_API_URL}/scenarios/remove", json={"scenario_name": scenario_name} + ) as response: + if response.status != 200: + raise HTTPException(status_code=response.status, detail=await response.text()) ScenarioManagement.remove_files_by_scenario(scenario_name) except Exception as e: @@ -433,18 +427,16 @@ async def get_scenarios( dict: A list of scenarios and the currently running scenario. """ try: - scenarios = await db.get_all_scenarios_and_check_completed(username=user, role=role) - - if role == "admin": - scenario_running = await db.get_running_scenario() - else: - scenario_running = await db.get_running_scenario(username=user) + async with aiohttp.ClientSession() as session: + async with session.get(f"{DATABASE_API_URL}/scenarios/{user}/{role}") as response: + if response.status == 200: + return await response.json() + else: + raise HTTPException(status_code=response.status, detail=await response.text()) except Exception as e: logging.exception(f"Error obtaining scenarios: {e}") raise HTTPException(status_code=500, detail="Internal server error") - return {"scenarios": scenarios, "scenario_running": scenario_running} - @app.post("/scenarios/update") async def update_scenario( @@ -453,7 +445,6 @@ async def update_scenario( end_time: str = Body(..., embed=True), scenario: dict = Body(..., embed=True), status: str = Body(..., embed=True), - role: str = Body(..., embed=True), username: str = Body(..., embed=True), ): """ @@ -465,20 +456,30 @@ async def update_scenario( end_time (str): End time of the scenario. scenario (dict): Scenario configuration. status (str): New status of the scenario (e.g., "running", "finished"). - role (str): Role associated with the scenario. username (str): User performing the update. Returns: dict: A message confirming the update. """ try: - await db.scenario_update_record(scenario_name, start_time, end_time, scenario, status, username) + payload = { + "scenario_name": scenario_name, + "start_time": start_time, + "end_time": end_time, + "scenario": scenario, + "status": status, + "username": username, + } + async with aiohttp.ClientSession() as session: + async with session.post(f"{DATABASE_API_URL}/scenarios/update", json=payload) as response: + if response.status == 200: + return await response.json() + else: + raise HTTPException(status_code=response.status, detail=await response.text()) except Exception as e: logging.exception(f"Error updating scenario {scenario_name}: {e}") raise HTTPException(status_code=500, detail="Internal server error") - return {"message": f"Scenario {scenario_name} updated successfully"} - @app.post("/scenarios/set_status_to_finished") async def set_scenario_status_to_finished( @@ -495,16 +496,17 @@ async def set_scenario_status_to_finished( dict: A message confirming the operation. """ try: - if all: - await db.scenario_set_all_status_to_finished() - else: - await db.scenario_set_status_to_finished(scenario_name) + payload = {"scenario_name": scenario_name, "all": all} + async with aiohttp.ClientSession() as session: + async with session.post(f"{DATABASE_API_URL}/scenarios/set_status_to_finished", json=payload) as response: + if response.status == 200: + return await response.json() + else: + raise HTTPException(status_code=response.status, detail=await response.text()) except Exception as e: logging.exception(f"Error setting scenario {scenario_name} to finished: {e}") raise HTTPException(status_code=500, detail="Internal server error") - return {"message": f"Scenario {scenario_name} status set to finished successfully"} - @app.get("/scenarios/running") async def get_running_scenario_endpoint(get_all: bool = False): @@ -518,7 +520,12 @@ async def get_running_scenario_endpoint(get_all: bool = False): dict or list: Running scenario(s) information. """ try: - return await db.get_running_scenario(get_all=get_all) + async with aiohttp.ClientSession() as session: + async with session.get(f"{DATABASE_API_URL}/scenarios/running", params={"get_all": str(get_all)}) as response: + if response.status == 200: + return await response.json() + else: + raise HTTPException(status_code=response.status, detail=await response.text()) except Exception as e: logging.exception(f"Error obtaining running scenario: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -542,8 +549,12 @@ async def check_scenario( dict: Whether the scenario is allowed for the role. """ try: - allowed = await db.check_scenario_with_role(role, scenario_name) - return {"allowed": allowed} + async with aiohttp.ClientSession() as session: + async with session.get(f"{DATABASE_API_URL}/scenarios/check/{role}/{scenario_name}") as response: + if response.status == 200: + return await response.json() + else: + raise HTTPException(status_code=response.status, detail=await response.text()) except Exception as e: logging.exception(f"Error checking scenario with role: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -565,13 +576,16 @@ async def get_scenario_by_name_endpoint( dict: The scenario data. """ try: - scenario = await db.get_scenario_by_name(scenario_name) + async with aiohttp.ClientSession() as session: + async with session.get(f"{DATABASE_API_URL}/scenarios/{scenario_name}") as response: + if response.status == 200: + return await response.json() + else: + raise HTTPException(status_code=response.status, detail=await response.text()) except Exception as e: logging.exception(f"Error obtaining scenario {scenario_name}: {e}") raise HTTPException(status_code=500, detail="Internal server error") - return scenario - @app.get("/nodes/{scenario_name}") async def list_nodes_by_scenario_name_endpoint( @@ -589,13 +603,16 @@ async def list_nodes_by_scenario_name_endpoint( list: List of nodes. """ try: - nodes = await db.list_nodes_by_scenario_name(scenario_name) + async with aiohttp.ClientSession() as session: + async with session.get(f"{DATABASE_API_URL}/nodes/{scenario_name}") as response: + if response.status == 200: + return await response.json() + else: + raise HTTPException(status_code=response.status, detail=await response.text()) except Exception as e: logging.exception(f"Error obtaining nodes: {e}") raise HTTPException(status_code=500, detail="Internal server error") - return nodes - @app.post("/nodes/{scenario_name}/update") async def update_nodes( @@ -618,23 +635,14 @@ async def update_nodes( try: config = await request.json() timestamp = datetime.datetime.now() + config["timestamp"] = str(timestamp) + # Update the node in database - await db.update_node_record( - str(config["device_args"]["uid"]), - str(config["device_args"]["idx"]), - str(config["network_args"]["ip"]), - str(config["network_args"]["port"]), - str(config["device_args"]["role"]), - config["network_args"]["neighbors"], - str(config["mobility_args"]["latitude"]), - str(config["mobility_args"]["longitude"]), - str(timestamp), - str(config["scenario_args"]["federation"]), - str(config["federation_args"]["round"]), - str(config["scenario_args"]["name"]), - str(config["tracking_args"]["run_hash"]), - str(config["device_args"]["malicious"]), - ) + async with aiohttp.ClientSession() as session: + async with session.post(f"{DATABASE_API_URL}/nodes/update", json=config) as response: + if response.status != 200: + raise HTTPException(status_code=response.status, detail=await response.text()) + except Exception as e: logging.exception(f"Error updating nodes: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -643,8 +651,6 @@ async def update_nodes( f"http://{os.environ['NEBULA_ENV_TAG']}_{os.environ['NEBULA_PREFIX_TAG']}_{os.environ['NEBULA_USER_TAG']}_nebula-frontend/platform/dashboard/{scenario_name}/node/update" ) - config["timestamp"] = str(timestamp) - async with aiohttp.ClientSession() as session: async with session.post(url, json=config) as response: if response.status == 200: @@ -700,7 +706,10 @@ async def remove_nodes_by_scenario_name_endpoint(scenario_name: str = Body(..., Returns a success message or an error if something goes wrong. """ try: - await db.remove_nodes_by_scenario_name(scenario_name) + async with aiohttp.ClientSession() as session: + async with session.post(f"{DATABASE_API_URL}/nodes/remove", json={"scenario_name": scenario_name}) as response: + if response.status != 200: + raise HTTPException(status_code=response.status, detail=await response.text()) except Exception as e: logging.exception(f"Error removing nodes: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -718,13 +727,12 @@ async def get_notes_by_scenario_name( Endpoint to retrieve notes associated with a scenario. """ try: - notes_record = await db.get_notes(scenario_name) - - if notes_record is not None: - notes_record = dict(notes_record.items()) - - return notes_record - + async with aiohttp.ClientSession() as session: + async with session.get(f"{DATABASE_API_URL}/notes/{scenario_name}") as response: + if response.status == 200: + return await response.json() + else: + raise HTTPException(status_code=response.status, detail=await response.text()) except Exception as e: logging.exception(f"Error obtaining notes for scenario {scenario_name}: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -742,13 +750,17 @@ async def update_notes_by_scenario_name(scenario_name: str = Body(..., embed=Tru Returns a success message or an error if something goes wrong. """ try: - await db.save_notes(scenario_name, notes) + payload = {"scenario_name": scenario_name, "notes": notes} + async with aiohttp.ClientSession() as session: + async with session.post(f"{DATABASE_API_URL}/notes/update", json=payload) as response: + if response.status == 200: + return await response.json() + else: + raise HTTPException(status_code=response.status, detail=await response.text()) except Exception as e: logging.exception(f"Error updating notes: {e}") raise HTTPException(status_code=500, detail="Internal server error") - return {"message": f"Notes for scenario {scenario_name} updated successfully"} - @app.post("/notes/remove") async def remove_notes_by_scenario_name_endpoint(scenario_name: str = Body(..., embed=True)): @@ -761,7 +773,10 @@ async def remove_notes_by_scenario_name_endpoint(scenario_name: str = Body(..., Returns a success message or an error if something goes wrong. """ try: - await db.remove_note(scenario_name) + async with aiohttp.ClientSession() as session: + async with session.post(f"{DATABASE_API_URL}/notes/remove", json={"scenario_name": scenario_name}) as response: + if response.status != 200: + raise HTTPException(status_code=response.status, detail=await response.text()) except Exception as e: logging.exception(f"Error removing notes: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -780,12 +795,16 @@ async def list_users_controller(all_info: bool = False): Returns a list of users or raises an HTTPException on error. """ try: - user_list = await db.list_users(all_info) - if all_info: - # Convert each asyncpg.Record to a dictionary so that it is JSON serializable. - user_list = [dict(user) for user in user_list] - return {"users": user_list} + async with aiohttp.ClientSession() as session: + async with session.get( + f"{DATABASE_API_URL}/user/list", params={"all_info": str(all_info)} + ) as response: + if response.status == 200: + return await response.json() + else: + raise HTTPException(status_code=response.status, detail=await response.text()) except Exception as e: + logging.exception(f"Error retrieving users: {e}") raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error retrieving users: {e}") @@ -804,13 +823,16 @@ async def get_user_by_scenario_name_endpoint( Returns user info or raises an HTTPException on error. """ try: - user = await db.get_user_by_scenario_name(scenario_name) + async with aiohttp.ClientSession() as session: + async with session.get(f"{DATABASE_API_URL}/user/{scenario_name}") as response: + if response.status == 200: + return await response.json() + else: + raise HTTPException(status_code=response.status, detail=await response.text()) except Exception as e: - logging.exception(f"Error obtaining user {user}: {e}") + logging.exception(f"Error obtaining user for scenario {scenario_name}: {e}") raise HTTPException(status_code=500, detail="Internal server error") - return user - @app.get("/discover-vpn") async def discover_vpn(): @@ -965,11 +987,11 @@ async def get_physical_scenario_state(scenario_name: str): } """ # 1) Retrieve scenario metadata and node list from the DB - scenario = await db.get_scenario_by_name(scenario_name) + scenario = await get_scenario_by_name_endpoint(scenario_name) if not scenario: raise HTTPException(status_code=404, detail="Scenario not found") - nodes = await db.list_nodes_by_scenario_name(scenario_name) + nodes = await list_nodes_by_scenario_name_endpoint(scenario_name) if not nodes: raise HTTPException(status_code=404, detail="No nodes found for scenario") @@ -1007,8 +1029,13 @@ async def add_user_controller(user: str = Body(...), password: str = Body(...), Returns a success message or an error if the user could not be added. """ try: - await db.add_user(user, password, role) - return {"detail": "User added successfully"} + payload = {"user": user, "password": password, "role": role} + async with aiohttp.ClientSession() as session: + async with session.post(f"{DATABASE_API_URL}/user/add", json=payload) as response: + if response.status == 200: + return await response.json() + else: + raise HTTPException(status_code=response.status, detail=await response.text()) except Exception as e: logging.exception(f"Error adding user: {e}") raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error adding user: {e}") @@ -1025,8 +1052,12 @@ async def remove_user_controller(user: str = Body(..., embed=True)): Returns a success message if the user is deleted, or an HTTP error if an exception occurs. """ try: - await db.delete_user_from_db(user) - return {"detail": "User deleted successfully"} + async with aiohttp.ClientSession() as session: + async with session.post(f"{DATABASE_API_URL}/user/delete", json={"user": user}) as response: + if response.status == 200: + return await response.json() + else: + raise HTTPException(status_code=response.status, detail=await response.text()) except Exception as e: logging.exception(f"Error deleting user: {e}") raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error deleting user: {e}") @@ -1045,8 +1076,13 @@ async def update_user_controller(user: str = Body(...), password: str = Body(... Returns a success message if the user is updated, or an HTTP error if an exception occurs. """ try: - await db.update_user(user, password, role) - return {"detail": "User updated successfully"} + payload = {"user": user, "password": password, "role": role} + async with aiohttp.ClientSession() as session: + async with session.post(f"{DATABASE_API_URL}/user/update", json=payload) as response: + if response.status == 200: + return await response.json() + else: + raise HTTPException(status_code=response.status, detail=await response.text()) except Exception as e: logging.exception(f"Error updating user: {e}") raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error updating user: {e}") @@ -1064,12 +1100,20 @@ async def verify_user_controller(user: str = Body(...), password: str = Body(... Returns the user role on success or raises an error on failure. """ try: - user_submitted = user.upper() - if (await db.list_users() and await db.verify(user_submitted, password)): - user_info = await db.get_user_info(user_submitted) - return {"user": user_submitted, "role": user_info[2]} - else: - raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED) + payload = {"user": user, "password": password} + async with aiohttp.ClientSession() as session: + async with session.post(f"{DATABASE_API_URL}/user/verify", json=payload) as response: + if response.status == 200: + return await response.json() + elif response.status == 401: + raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED) + else: + raise HTTPException(status_code=response.status, detail=await response.text()) + except HTTPException as e: + if e.status_code == 401: + raise e + logging.exception(f"Error verifying user: {e}") + raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error verifying user: {e}") except Exception as e: logging.exception(f"Error verifying user: {e}") raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error verifying user: {e}") diff --git a/nebula/database/adapters/postgress/docker/Dockerfile b/nebula/database/adapters/postgress/docker/Dockerfile index b1b1d26b6..2859de8da 100644 --- a/nebula/database/adapters/postgress/docker/Dockerfile +++ b/nebula/database/adapters/postgress/docker/Dockerfile @@ -4,7 +4,7 @@ FROM postgres:17.5-alpine3.22 RUN mv /usr/local/bin/docker-entrypoint.sh /usr/local/bin/docker-entrypoint.sh.orig # Copy SQL init file and custom entrypoint script -COPY /nebula/database/docker/docker-entrypoint.sh /usr/local/bin/docker-entrypoint.sh +COPY ./nebula/database/adapters/postgress/docker/docker-entrypoint.sh /usr/local/bin/docker-entrypoint.sh RUN chmod +x /usr/local/bin/docker-entrypoint.sh # Install Python 3.11.7 from source @@ -41,12 +41,13 @@ ADD https://astral.sh/uv/install.sh /uv-installer.sh RUN sh /uv-installer.sh && rm /uv-installer.sh ENV PATH="/root/.local/bin/:$PATH" -# Install Python dependencies using uv COPY pyproject.toml . + +# Install Python dependencies using uv RUN uv python pin 3.11.7 RUN uv sync --group database -ENV PATH=".venv/bin:$PATH" +ENV PATH="/.venv/bin:$PATH" -# ENTRYPOINT ["/usr/local/bin/docker-entrypoint.sh"] -# CMD ["postgres"] +ENTRYPOINT ["/bin/bash", "/usr/local/bin/docker-entrypoint.sh"] +CMD ["postgres"] diff --git a/nebula/database/adapters/postgress/docker/docker-entrypoint.sh b/nebula/database/adapters/postgress/docker/docker-entrypoint.sh index a298ca23b..911961294 100644 --- a/nebula/database/adapters/postgress/docker/docker-entrypoint.sh +++ b/nebula/database/adapters/postgress/docker/docker-entrypoint.sh @@ -1,18 +1,21 @@ #!/bin/sh -set -e +set -x -# 1) Run the original entrypoint and wait for it to finish initialization -/usr/local/bin/docker-entrypoint.sh.orig "$@" +# Start the python API in the background +echo "🐍 Starting Nebula Database API in the background..." +( + # Wait for postgres to be ready + until pg_isready -U "$POSTGRES_USER" -d "$POSTGRES_DB" -h localhost >/dev/null 2>&1; do + sleep 1 + done + echo "✅ PostgreSQL is ready, starting API." -# 2) Wait until PostgreSQL accepts connections to the configured database -echo "⏳ Waiting for PostgreSQL to be ready..." -until pg_isready -U "$POSTGRES_USER" -d "$POSTGRES_DB" >/dev/null 2>&1; do - sleep 1 -done + cd nebula + NEBULA_SOCK=nebula.sock -# 3) Apply the SQL initialization script -echo "🚀 Applying init-configs.sql..." -psql -v ON_ERROR_STOP=1 \ - -U "$POSTGRES_USER" \ - -d "$POSTGRES_DB" \ - -f /docker-entrypoint-initdb.d/init-configs.sql + uvicorn nebula.database.database_api:app --host 0.0.0.0 --port 5051 --log-level debug --proxy-headers --forwarded-allow-ips "*" +) & + +# Run the original postgres entrypoint in the foreground +# This will become the main process of the container +exec /usr/local/bin/docker-entrypoint.sh.orig "$@" diff --git a/nebula/database/adapters/postgress/docker/init-configs.sql b/nebula/database/adapters/postgress/docker/init-configs.sql index a34b17841..9a452af93 100644 --- a/nebula/database/adapters/postgress/docker/init-configs.sql +++ b/nebula/database/adapters/postgress/docker/init-configs.sql @@ -13,7 +13,7 @@ CREATE TABLE IF NOT EXISTS users ( role TEXT ); --- 2) Nodes como JSONB +-- 2) Nodes CREATE TABLE IF NOT EXISTS nodes ( uid TEXT PRIMARY KEY, idx TEXT, @@ -31,7 +31,7 @@ CREATE TABLE IF NOT EXISTS nodes ( malicious TEXT ); --- 3) Configs como JSONB +-- 3) Configs as JSONB DROP INDEX IF EXISTS idx_configs_config_gin; DROP TABLE IF EXISTS configs; CREATE TABLE configs ( diff --git a/nebula/database/adapters/postgress/postgress.py b/nebula/database/adapters/postgress/postgress.py index 3de289780..57019164f 100755 --- a/nebula/database/adapters/postgress/postgress.py +++ b/nebula/database/adapters/postgress/postgress.py @@ -95,7 +95,7 @@ async def insert_default_admin(self): except Exception as e: logging.error(f"Failed to insert default admin user: {e}", exc_info=True) - async def list_users(self, all_info=False): + async def list_users(self, all_info: bool = False): """ Retrieves a list of users from the users database. """ diff --git a/nebula/database/database_api.py b/nebula/database/database_api.py new file mode 100644 index 000000000..aac0a981c --- /dev/null +++ b/nebula/database/database_api.py @@ -0,0 +1,337 @@ + +import argparse +import logging +import os +import sys +from typing import Annotated + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../'))) + +from fastapi import Body, FastAPI, HTTPException, Path, status +from fastapi.concurrency import asynccontextmanager + +from nebula.database.database_adapter_factory import factory_database_adapter + +# Get a database instance +db = factory_database_adapter("PostgresDB") + + +# Setup logger +def configure_logger(log_file): + """ + Configures the logging system for the database API. + """ + log_console_format = "[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s" + console_handler = logging.StreamHandler() + console_handler.setLevel(logging.INFO) + console_handler.setFormatter(logging.Formatter(log_console_format)) + file_handler = logging.FileHandler(log_file, mode="w") + file_handler.setLevel(logging.INFO) + file_handler.setFormatter(logging.Formatter("[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s")) + logging.basicConfig( + level=logging.DEBUG, + handlers=[ + console_handler, + file_handler, + ], + ) + uvicorn_loggers = ["uvicorn", "uvicorn.error", "uvicorn.access"] + for logger_name in uvicorn_loggers: + logger = logging.getLogger(logger_name) + logger.handlers = [] + logger.propagate = False + handler = logging.FileHandler(log_file, mode="a") + handler.setFormatter(logging.Formatter("[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s")) + logger.addHandler(handler) + + +@asynccontextmanager +async def lifespan(app: FastAPI): + """ + Application lifespan context manager for the database API. + """ + # Code to run on startup + db_log = os.environ.get("NEBULA_DATABASE_LOG", "database.log") + configure_logger(db_log) + + # Initialize the database connection pool + await db.init_db_pool() + await db.insert_default_admin() + + yield + + # Code to run on shutdown + await db.close_db_pool() + + +app = FastAPI(lifespan=lifespan) + + +@app.get("/") +async def read_root(): + return {"message": "Welcome to the NEBULA Database API"} + + +# Scenarios +@app.post("/scenarios/update") +async def update_scenario( + scenario_name: str = Body(..., embed=True), + start_time: str = Body(..., embed=True), + end_time: str = Body(..., embed=True), + scenario: dict = Body(..., embed=True), + status: str = Body(..., embed=True), + username: str = Body(..., embed=True), +): + try: + await db.scenario_update_record(scenario_name, start_time, end_time, scenario, status, username) + return {"message": f"Scenario {scenario_name} updated successfully"} + except Exception as e: + logging.exception(f"Error updating scenario {scenario_name}: {e}") + raise HTTPException(status_code=500, detail="Internal server error") + + +@app.post("/scenarios/stop") +async def stop_scenario( + scenario_name: str = Body(..., embed=True), + all: bool = Body(False, embed=True), +): + try: + if all: + await db.scenario_set_all_status_to_finished() + else: + await db.scenario_set_status_to_finished(scenario_name) + return {"message": "Scenario status updated successfully"} + except Exception as e: + logging.exception(f"Error stopping scenario {scenario_name}: {e}") + raise HTTPException(status_code=500, detail="Internal server error") + + +@app.post("/scenarios/remove") +async def remove_scenario( + scenario_name: str = Body(..., embed=True), +): + try: + await db.remove_scenario_by_name(scenario_name) + return {"message": f"Scenario {scenario_name} removed successfully"} + except Exception as e: + logging.exception(f"Error removing scenario {scenario_name}: {e}") + raise HTTPException(status_code=500, detail="Internal server error") + + +@app.get("/scenarios/{user}/{role}") +async def get_scenarios( + user: Annotated[str, Path(pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50)], + role: Annotated[str, Path(pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50)], +): + try: + scenarios = await db.get_all_scenarios_and_check_completed(username=user, role=role) + if role == "admin": + scenario_running = await db.get_running_scenario() + else: + scenario_running = await db.get_running_scenario(username=user) + return {"scenarios": scenarios, "scenario_running": scenario_running} + except Exception as e: + logging.exception(f"Error obtaining scenarios: {e}") + raise HTTPException(status_code=500, detail="Internal server error") + + +@app.post("/scenarios/set_status_to_finished") +async def set_scenario_status_to_finished( + scenario_name: str = Body(..., embed=True), all: bool = Body(False, embed=True) +): + try: + if all: + await db.scenario_set_all_status_to_finished() + else: + await db.scenario_set_status_to_finished(scenario_name) + return {"message": f"Scenario {scenario_name} status set to finished successfully"} + except Exception as e: + logging.exception(f"Error setting scenario {scenario_name} to finished: {e}") + raise HTTPException(status_code=500, detail="Internal server error") + + +@app.get("/scenarios/running") +async def get_running_scenario_endpoint(get_all: bool = False): + try: + return await db.get_running_scenario(get_all=get_all) + except Exception as e: + logging.exception(f"Error obtaining running scenario: {e}") + raise HTTPException(status_code=500, detail="Internal server error") + + +@app.get("/scenarios/check/{role}/{scenario_name}") +async def check_scenario( + role: Annotated[str, Path(pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50)], + scenario_name: Annotated[str, Path(pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50)], +): + try: + allowed = await db.check_scenario_with_role(role, scenario_name) + return {"allowed": allowed} + except Exception as e: + logging.exception(f"Error checking scenario with role: {e}") + raise HTTPException(status_code=500, detail="Internal server error") + + +@app.get("/scenarios/{scenario_name}") +async def get_scenario_by_name_endpoint( + scenario_name: Annotated[str, Path(pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50)], +): + try: + scenario = await db.get_scenario_by_name(scenario_name) + return scenario + except Exception as e: + logging.exception(f"Error obtaining scenario {scenario_name}: {e}") + raise HTTPException(status_code=500, detail="Internal server error") + + +# Nodes +@app.get("/nodes/{scenario_name}") +async def list_nodes_by_scenario_name_endpoint( + scenario_name: Annotated[str, Path(pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50)], +): + try: + nodes = await db.list_nodes_by_scenario_name(scenario_name) + return nodes + except Exception as e: + logging.exception(f"Error obtaining nodes: {e}") + raise HTTPException(status_code=500, detail="Internal server error") + + +@app.post("/nodes/update") +async def update_node_record(data: dict): + try: + await db.update_node_record( + str(data["device_args"]["uid"]), + str(data["device_args"]["idx"]), + str(data["network_args"]["ip"]), + str(data["network_args"]["port"]), + str(data["device_args"]["role"]), + data["network_args"]["neighbors"], + str(data["mobility_args"]["latitude"]), + str(data["mobility_args"]["longitude"]), + str(data["timestamp"]), + str(data["scenario_args"]["federation"]), + str(data["federation_args"]["round"]), + str(data["scenario_args"]["name"]), + str(data["tracking_args"]["run_hash"]), + str(data["device_args"]["malicious"]), + ) + return {"message": "Node updated successfully"} + except Exception as e: + logging.exception(f"Error updating node: {e}") + raise HTTPException(status_code=500, detail="Internal server error") + + +@app.post("/nodes/remove") +async def remove_nodes_by_scenario_name_endpoint(scenario_name: str = Body(..., embed=True)): + try: + await db.remove_nodes_by_scenario_name(scenario_name) + return {"message": f"Nodes for scenario {scenario_name} removed successfully"} + except Exception as e: + logging.exception(f"Error removing nodes: {e}") + raise HTTPException(status_code=500, detail="Internal server error") + + +# Notes +@app.get("/notes/{scenario_name}") +async def get_notes_by_scenario_name( + scenario_name: Annotated[str, Path(pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50)], +): + try: + notes_record = await db.get_notes(scenario_name) + if notes_record is not None: + notes_record = dict(notes_record.items()) + return notes_record + except Exception as e: + logging.exception(f"Error obtaining notes for scenario {scenario_name}: {e}") + raise HTTPException(status_code=500, detail="Internal server error") + + +@app.post("/notes/update") +async def update_notes_by_scenario_name(scenario_name: str = Body(..., embed=True), notes: str = Body(..., embed=True)): + try: + await db.save_notes(scenario_name, notes) + return {"message": f"Notes for scenario {scenario_name} updated successfully"} + except Exception as e: + logging.exception(f"Error updating notes: {e}") + raise HTTPException(status_code=500, detail="Internal server error") + + +@app.post("/notes/remove") +async def remove_notes_by_scenario_name_endpoint(scenario_name: str = Body(..., embed=True)): + try: + await db.remove_note(scenario_name) + return {"message": f"Notes for scenario {scenario_name} removed successfully"} + except Exception as e: + logging.exception(f"Error removing notes: {e}") + raise HTTPException(status_code=500, detail="Internal server error") + + +# Users +@app.get("/user/list") +async def list_users_controller(all_info: bool = False): + try: + user_list = await db.list_users(all_info) + if all_info: + user_list = [dict(user) for user in user_list] + return {"users": user_list} + except Exception as e: + logging.exception(f"Error retrieving users: {e}") + raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error retrieving users: {e}") + + +@app.get("/user/{scenario_name}") +async def get_user_by_scenario_name_endpoint( + scenario_name: Annotated[str, Path(pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50)], +): + try: + user = await db.get_user_by_scenario_name(scenario_name) + return user + except Exception as e: + logging.exception(f"Error obtaining user for scenario {scenario_name}: {e}") + raise HTTPException(status_code=500, detail="Internal server error") + + +@app.post("/user/add") +async def add_user_controller(user: str = Body(...), password: str = Body(...), role: str = Body(...)): + try: + await db.add_user(user, password, role) + return {"detail": "User added successfully"} + except Exception as e: + logging.exception(f"Error adding user: {e}") + raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error adding user: {e}") + + +@app.post("/user/delete") +async def remove_user_controller(user: str = Body(..., embed=True)): + try: + await db.delete_user_from_db(user) + return {"detail": "User deleted successfully"} + except Exception as e: + logging.exception(f"Error deleting user: {e}") + raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error deleting user: {e}") + + +@app.post("/user/update") +async def update_user_controller(user: str = Body(...), password: str = Body(...), role: str = Body(...)): + try: + await db.update_user(user, password, role) + return {"detail": "User updated successfully"} + except Exception as e: + logging.exception(f"Error updating user: {e}") + raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error updating user: {e}") + + +@app.post("/user/verify") +async def verify_user_controller(user: str = Body(...), password: str = Body(...)): + try: + user_submitted = user.upper() + users = await db.list_users() + if users and await db.verify(user_submitted, password): + user_info = await db.get_user_info(user_submitted) + return {"user": user_submitted, "role": user_info[2]} + else: + raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED) + except Exception as e: + logging.exception(f"Error verifying user: {e}") + raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error verifying user: {e}") diff --git a/nebula/frontend/app.py b/nebula/frontend/app.py index 42f5b3a68..fb065e308 100755 --- a/nebula/frontend/app.py +++ b/nebula/frontend/app.py @@ -627,7 +627,7 @@ async def get_scenarios(user, role): return await controller_get(url) -async def scenario_update_record(scenario_name, start_time, end_time, scenario, status, role, username): +async def scenario_update_record(scenario_name, start_time, end_time, scenario, status, username): """ Update the record of a scenario's execution status on the controller. @@ -637,7 +637,6 @@ async def scenario_update_record(scenario_name, start_time, end_time, scenario, end_time (str): ISO-formatted end timestamp. scenario (Any): Scenario payload or identifier. status (str): New status value (e.g., 'running', 'finished'). - role (str): Role associated with the scenario. username (str): User who ran or updated the scenario. Raises: @@ -650,7 +649,6 @@ async def scenario_update_record(scenario_name, start_time, end_time, scenario, "end_time": end_time, "scenario": scenario, "status": status, - "role": role, "username": username, } await controller_post(url, data) diff --git a/pyproject.toml b/pyproject.toml index 7c7e666b5..d1bf059c1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,12 +62,8 @@ docs = [ "mkdocstrings[python]<1.0.0,>=0.26.2", ] controller = [ - "psycopg2-binary==2.9.10", - "asyncpg==0.30.0", - "passlib==1.7.4", "aiohttp==3.10.5", "aiosqlite==0.20.0", - "argon2-cffi==23.1.0", "docker==7.1.0", "fastapi[all]==0.114.0", "gunicorn==23.0.0", @@ -84,8 +80,11 @@ controller = [ "scikit-learn==1.5.1", ] database = [ + "argon2-cffi==23.1.0", "asyncpg==0.30.0", - "psycopg2-binary==2.9.10" + "psycopg2-binary==2.9.10", + "passlib==1.7.4", + "fastapi[all]==0.114.0", ] core = [ "aiohttp==3.10.5", From c70cec84db407a8a8b2bd0e1ba82cf901ada6179 Mon Sep 17 00:00:00 2001 From: FerTV Date: Tue, 29 Jul 2025 14:35:54 +0200 Subject: [PATCH 04/14] refactor: using apiutils for hub endpoints --- nebula/controller/hub.py | 163 ++++++--------------------------------- nebula/utils.py | 155 +++++++++++++++++++++++++++++++++++++ 2 files changed, 179 insertions(+), 139 deletions(-) diff --git a/nebula/controller/hub.py b/nebula/controller/hub.py index c170102bf..713316362 100755 --- a/nebula/controller/hub.py +++ b/nebula/controller/hub.py @@ -17,7 +17,7 @@ from fastapi.concurrency import asynccontextmanager from nebula.controller.http_helpers import remote_get, remote_post_form -from nebula.utils import DockerUtils +from nebula.utils import APIUtils, DockerUtils # URL for the database API DATABASE_API_URL = os.environ.get("NEBULA_DATABASE_API_URL", "http://nebula-database:5051") @@ -369,12 +369,7 @@ async def stop_scenario( ScenarioManagement.cleanup_scenario_containers() try: - async with aiohttp.ClientSession() as session: - async with session.post( - f"{DATABASE_API_URL}/scenarios/stop", json={"scenario_name": scenario_name, "all": all} - ) as response: - if response.status != 200: - raise HTTPException(status_code=response.status, detail=await response.text()) + await APIUtils.post(f"{DATABASE_API_URL}/scenarios/stop", data={"scenario_name": scenario_name, "all": all}) except Exception as e: logging.exception(f"Error setting scenario {scenario_name} to finished: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -396,12 +391,7 @@ async def remove_scenario( from nebula.controller.scenarios import ScenarioManagement try: - async with aiohttp.ClientSession() as session: - async with session.post( - f"{DATABASE_API_URL}/scenarios/remove", json={"scenario_name": scenario_name} - ) as response: - if response.status != 200: - raise HTTPException(status_code=response.status, detail=await response.text()) + await APIUtils.post(f"{DATABASE_API_URL}/scenarios/remove", data={"scenario_name": scenario_name}) ScenarioManagement.remove_files_by_scenario(scenario_name) except Exception as e: @@ -427,12 +417,7 @@ async def get_scenarios( dict: A list of scenarios and the currently running scenario. """ try: - async with aiohttp.ClientSession() as session: - async with session.get(f"{DATABASE_API_URL}/scenarios/{user}/{role}") as response: - if response.status == 200: - return await response.json() - else: - raise HTTPException(status_code=response.status, detail=await response.text()) + return await APIUtils.get(f"{DATABASE_API_URL}/scenarios/{user}/{role}") except Exception as e: logging.exception(f"Error obtaining scenarios: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -470,12 +455,7 @@ async def update_scenario( "status": status, "username": username, } - async with aiohttp.ClientSession() as session: - async with session.post(f"{DATABASE_API_URL}/scenarios/update", json=payload) as response: - if response.status == 200: - return await response.json() - else: - raise HTTPException(status_code=response.status, detail=await response.text()) + return await APIUtils.post(f"{DATABASE_API_URL}/scenarios/update", data=payload) except Exception as e: logging.exception(f"Error updating scenario {scenario_name}: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -497,12 +477,7 @@ async def set_scenario_status_to_finished( """ try: payload = {"scenario_name": scenario_name, "all": all} - async with aiohttp.ClientSession() as session: - async with session.post(f"{DATABASE_API_URL}/scenarios/set_status_to_finished", json=payload) as response: - if response.status == 200: - return await response.json() - else: - raise HTTPException(status_code=response.status, detail=await response.text()) + return await APIUtils.post(f"{DATABASE_API_URL}/scenarios/set_status_to_finished", data=payload) except Exception as e: logging.exception(f"Error setting scenario {scenario_name} to finished: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -520,12 +495,7 @@ async def get_running_scenario_endpoint(get_all: bool = False): dict or list: Running scenario(s) information. """ try: - async with aiohttp.ClientSession() as session: - async with session.get(f"{DATABASE_API_URL}/scenarios/running", params={"get_all": str(get_all)}) as response: - if response.status == 200: - return await response.json() - else: - raise HTTPException(status_code=response.status, detail=await response.text()) + return await APIUtils.get(f"{DATABASE_API_URL}/scenarios/running", params={"get_all": str(get_all)}) except Exception as e: logging.exception(f"Error obtaining running scenario: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -549,12 +519,7 @@ async def check_scenario( dict: Whether the scenario is allowed for the role. """ try: - async with aiohttp.ClientSession() as session: - async with session.get(f"{DATABASE_API_URL}/scenarios/check/{role}/{scenario_name}") as response: - if response.status == 200: - return await response.json() - else: - raise HTTPException(status_code=response.status, detail=await response.text()) + return await APIUtils.get(f"{DATABASE_API_URL}/scenarios/check/{role}/{scenario_name}") except Exception as e: logging.exception(f"Error checking scenario with role: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -576,12 +541,7 @@ async def get_scenario_by_name_endpoint( dict: The scenario data. """ try: - async with aiohttp.ClientSession() as session: - async with session.get(f"{DATABASE_API_URL}/scenarios/{scenario_name}") as response: - if response.status == 200: - return await response.json() - else: - raise HTTPException(status_code=response.status, detail=await response.text()) + return await APIUtils.get(f"{DATABASE_API_URL}/scenarios/{scenario_name}") except Exception as e: logging.exception(f"Error obtaining scenario {scenario_name}: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -603,12 +563,7 @@ async def list_nodes_by_scenario_name_endpoint( list: List of nodes. """ try: - async with aiohttp.ClientSession() as session: - async with session.get(f"{DATABASE_API_URL}/nodes/{scenario_name}") as response: - if response.status == 200: - return await response.json() - else: - raise HTTPException(status_code=response.status, detail=await response.text()) + return await APIUtils.get(f"{DATABASE_API_URL}/nodes/{scenario_name}") except Exception as e: logging.exception(f"Error obtaining nodes: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -638,10 +593,7 @@ async def update_nodes( config["timestamp"] = str(timestamp) # Update the node in database - async with aiohttp.ClientSession() as session: - async with session.post(f"{DATABASE_API_URL}/nodes/update", json=config) as response: - if response.status != 200: - raise HTTPException(status_code=response.status, detail=await response.text()) + await APIUtils.post(f"{DATABASE_API_URL}/nodes/update", data=config) except Exception as e: logging.exception(f"Error updating nodes: {e}") @@ -651,14 +603,7 @@ async def update_nodes( f"http://{os.environ['NEBULA_ENV_TAG']}_{os.environ['NEBULA_PREFIX_TAG']}_{os.environ['NEBULA_USER_TAG']}_nebula-frontend/platform/dashboard/{scenario_name}/node/update" ) - async with aiohttp.ClientSession() as session: - async with session.post(url, json=config) as response: - if response.status == 200: - return await response.json() - else: - raise HTTPException(status_code=response.status, detail="Error posting data") - - return {"message": "Nodes updated successfully in the database"} + return await APIUtils.post(url, data=config) @app.post("/nodes/{scenario_name}/done") @@ -685,14 +630,7 @@ async def node_done( data = await request.json() - async with aiohttp.ClientSession() as session: - async with session.post(url, json=data) as response: - if response.status == 200: - return await response.json() - else: - raise HTTPException(status_code=response.status, detail="Error posting data") - - return {"message": "Nodes done"} + return await APIUtils.post(url, data=data) @app.post("/nodes/remove") @@ -706,10 +644,7 @@ async def remove_nodes_by_scenario_name_endpoint(scenario_name: str = Body(..., Returns a success message or an error if something goes wrong. """ try: - async with aiohttp.ClientSession() as session: - async with session.post(f"{DATABASE_API_URL}/nodes/remove", json={"scenario_name": scenario_name}) as response: - if response.status != 200: - raise HTTPException(status_code=response.status, detail=await response.text()) + await APIUtils.post(f"{DATABASE_API_URL}/nodes/remove", data={"scenario_name": scenario_name}) except Exception as e: logging.exception(f"Error removing nodes: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -727,12 +662,7 @@ async def get_notes_by_scenario_name( Endpoint to retrieve notes associated with a scenario. """ try: - async with aiohttp.ClientSession() as session: - async with session.get(f"{DATABASE_API_URL}/notes/{scenario_name}") as response: - if response.status == 200: - return await response.json() - else: - raise HTTPException(status_code=response.status, detail=await response.text()) + return await APIUtils.get(f"{DATABASE_API_URL}/notes/{scenario_name}") except Exception as e: logging.exception(f"Error obtaining notes for scenario {scenario_name}: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -751,12 +681,7 @@ async def update_notes_by_scenario_name(scenario_name: str = Body(..., embed=Tru """ try: payload = {"scenario_name": scenario_name, "notes": notes} - async with aiohttp.ClientSession() as session: - async with session.post(f"{DATABASE_API_URL}/notes/update", json=payload) as response: - if response.status == 200: - return await response.json() - else: - raise HTTPException(status_code=response.status, detail=await response.text()) + return await APIUtils.post(f"{DATABASE_API_URL}/notes/update", data=payload) except Exception as e: logging.exception(f"Error updating notes: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -773,10 +698,7 @@ async def remove_notes_by_scenario_name_endpoint(scenario_name: str = Body(..., Returns a success message or an error if something goes wrong. """ try: - async with aiohttp.ClientSession() as session: - async with session.post(f"{DATABASE_API_URL}/notes/remove", json={"scenario_name": scenario_name}) as response: - if response.status != 200: - raise HTTPException(status_code=response.status, detail=await response.text()) + await APIUtils.post(f"{DATABASE_API_URL}/notes/remove", data={"scenario_name": scenario_name}) except Exception as e: logging.exception(f"Error removing notes: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -795,14 +717,7 @@ async def list_users_controller(all_info: bool = False): Returns a list of users or raises an HTTPException on error. """ try: - async with aiohttp.ClientSession() as session: - async with session.get( - f"{DATABASE_API_URL}/user/list", params={"all_info": str(all_info)} - ) as response: - if response.status == 200: - return await response.json() - else: - raise HTTPException(status_code=response.status, detail=await response.text()) + return await APIUtils.get(f"{DATABASE_API_URL}/user/list", params={"all_info": str(all_info)}) except Exception as e: logging.exception(f"Error retrieving users: {e}") raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error retrieving users: {e}") @@ -823,12 +738,7 @@ async def get_user_by_scenario_name_endpoint( Returns user info or raises an HTTPException on error. """ try: - async with aiohttp.ClientSession() as session: - async with session.get(f"{DATABASE_API_URL}/user/{scenario_name}") as response: - if response.status == 200: - return await response.json() - else: - raise HTTPException(status_code=response.status, detail=await response.text()) + return await APIUtils.get(f"{DATABASE_API_URL}/user/{scenario_name}") except Exception as e: logging.exception(f"Error obtaining user for scenario {scenario_name}: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -1030,12 +940,7 @@ async def add_user_controller(user: str = Body(...), password: str = Body(...), """ try: payload = {"user": user, "password": password, "role": role} - async with aiohttp.ClientSession() as session: - async with session.post(f"{DATABASE_API_URL}/user/add", json=payload) as response: - if response.status == 200: - return await response.json() - else: - raise HTTPException(status_code=response.status, detail=await response.text()) + return await APIUtils.post(f"{DATABASE_API_URL}/user/add", data=payload) except Exception as e: logging.exception(f"Error adding user: {e}") raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error adding user: {e}") @@ -1052,12 +957,7 @@ async def remove_user_controller(user: str = Body(..., embed=True)): Returns a success message if the user is deleted, or an HTTP error if an exception occurs. """ try: - async with aiohttp.ClientSession() as session: - async with session.post(f"{DATABASE_API_URL}/user/delete", json={"user": user}) as response: - if response.status == 200: - return await response.json() - else: - raise HTTPException(status_code=response.status, detail=await response.text()) + return await APIUtils.post(f"{DATABASE_API_URL}/user/delete", data={"user": user}) except Exception as e: logging.exception(f"Error deleting user: {e}") raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error deleting user: {e}") @@ -1077,12 +977,7 @@ async def update_user_controller(user: str = Body(...), password: str = Body(... """ try: payload = {"user": user, "password": password, "role": role} - async with aiohttp.ClientSession() as session: - async with session.post(f"{DATABASE_API_URL}/user/update", json=payload) as response: - if response.status == 200: - return await response.json() - else: - raise HTTPException(status_code=response.status, detail=await response.text()) + return await APIUtils.post(f"{DATABASE_API_URL}/user/update", data=payload) except Exception as e: logging.exception(f"Error updating user: {e}") raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error updating user: {e}") @@ -1101,20 +996,10 @@ async def verify_user_controller(user: str = Body(...), password: str = Body(... """ try: payload = {"user": user, "password": password} - async with aiohttp.ClientSession() as session: - async with session.post(f"{DATABASE_API_URL}/user/verify", json=payload) as response: - if response.status == 200: - return await response.json() - elif response.status == 401: - raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED) - else: - raise HTTPException(status_code=response.status, detail=await response.text()) + return await APIUtils.post(f"{DATABASE_API_URL}/user/verify", data=payload) except HTTPException as e: if e.status_code == 401: - raise e - logging.exception(f"Error verifying user: {e}") - raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error verifying user: {e}") - except Exception as e: + raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED) from e logging.exception(f"Error verifying user: {e}") raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error verifying user: {e}") diff --git a/nebula/utils.py b/nebula/utils.py index 60819ed1a..ec8aa114a 100644 --- a/nebula/utils.py +++ b/nebula/utils.py @@ -2,8 +2,16 @@ import os import socket +import aiohttp import docker +import re +from typing import Optional + +from fastapi import HTTPException +from aiohttp import ClientConnectorError +from aiohttp.client_exceptions import ClientError +import asyncio class FileUtils: """ @@ -202,3 +210,150 @@ def check_docker_by_prefix(cls, prefix): logging.exception("Error interacting with Docker") except Exception: logging.exception("Unexpected error") + + +class LoggerUtils: + + @staticmethod + def configure_logger( + name: Optional[str] = None, + log_file: Optional[str] = None, + level: int = logging.INFO, + console: bool = True, + strip_ansi: bool = True, + file_mode: str = "w", + log_format: str = "[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s", + date_format: str = "%Y-%m-%d %H:%M:%S", + ) -> logging.Logger: + """ + Configure and return a logger with optional console and file output. + + Args: + name (str): Logger name. If None, the root logger is used. + log_file (str): Path to the log file. + level (int): Logging level (DEBUG, INFO, etc). + console (bool): If True, output is also printed to the console. + strip_ansi (bool): Placeholder for future ANSI stripping support. + file_mode (str): File mode for the log file ('a' for append, 'w' for overwrite). + log_format (str): Format for log messages. + date_format (str): Format for timestamps. + + Returns: + logging.Logger: Configured logger instance. + """ + logger = logging.getLogger(name) + logger.setLevel(level) + + # Prevent duplicate handler setup + if getattr(logger, "_is_configured", False): + return logger + + formatter = logging.Formatter(fmt=log_format, datefmt=date_format) + + if log_file: + os.makedirs(os.path.dirname(log_file), exist_ok=True) + fh = logging.FileHandler(log_file, mode=file_mode) + fh.setLevel(level) + fh.setFormatter(formatter) + logger.addHandler(fh) + + if console: + ch = logging.StreamHandler() + ch.setLevel(level) + ch.setFormatter(formatter) + logger.addHandler(ch) + + # Mark this logger as configured to avoid re-adding handlers + logger._is_configured = True + logger.propagate = False + + return logger + +class APIUtils(): + + @staticmethod + async def retry_with_backoff(func, *args, max_retries=5, initial_delay=1): + """ + Retry a function with exponential backoff. + + Args: + func: The async function to retry + *args: Arguments to pass to the function + max_retries: Maximum number of retry attempts + initial_delay: Initial delay between retries in seconds + + Returns: + The result of the function if successful + + Raises: + The last exception if all retries fail + """ + delay = initial_delay + last_exception = None + + for attempt in range(max_retries): + try: + return await func(*args) + except (ClientConnectorError, ClientError) as e: + last_exception = e + if attempt < max_retries - 1: + logging.warning(f"Connection attempt {attempt + 1} failed: {str(e)}. Retrying in {delay} seconds...") + await asyncio.sleep(delay) + delay *= 2 # Exponential backoff + else: + logging.error(f"All {max_retries} connection attempts failed") + raise last_exception + + @staticmethod + async def get(url, params=None): + """ + Fetch JSON data from a remote controller endpoint via asynchronous HTTP GET. + + Parameters: + url (str): The full URL of the controller API endpoint. + params (dict, optional): A dictionary of query parameters to be sent with the request. + + Returns: + Any: Parsed JSON response when the HTTP status code is 200. + + Raises: + HTTPException: If the response status is not 200, raises with the response status code and an error detail. + """ + + async def _get(): + async with aiohttp.ClientSession() as session: + async with session.get(url, params=params) as response: + if response.status == 200: + return await response.json() + else: + detail = await response.text() + raise HTTPException(status_code=response.status, detail=detail) + + return await APIUtils.retry_with_backoff(_get) + + @staticmethod + async def post(url, data=None): + """ + Asynchronously send a JSON payload via HTTP POST to a controller endpoint and parse the response. + + Parameters: + url (str): The full URL of the controller API endpoint. + data (Any, optional): JSON-serializable payload to include in the POST request (default: None). + + Returns: + Any: Parsed JSON response when the HTTP status code is 200. + + Raises: + HTTPException: If the response status is not 200, with the status code and an error detail. + """ + + async def _post(): + async with aiohttp.ClientSession() as session: + async with session.post(url, json=data) as response: + if response.status == 200: + return await response.json() + else: + detail = await response.text() + raise HTTPException(status_code=response.status, detail=detail) + + return await APIUtils.retry_with_backoff(_post) From 54ba89dd7cdc9a22a4b76ef9fabb4849369dc6e7 Mon Sep 17 00:00:00 2001 From: FerTV Date: Wed, 17 Sep 2025 16:54:01 +0200 Subject: [PATCH 05/14] refactor: database and hub endpoints --- nebula/controller/hub.py | 195 +++++++++++++++++----------- nebula/database/database_api.py | 43 +++--- nebula/database/utils_requests.py | 209 ++++++++++++++++++++++++++++++ 3 files changed, 353 insertions(+), 94 deletions(-) create mode 100644 nebula/database/utils_requests.py diff --git a/nebula/controller/hub.py b/nebula/controller/hub.py index 713316362..9bb3c8cee 100755 --- a/nebula/controller/hub.py +++ b/nebula/controller/hub.py @@ -18,6 +18,23 @@ from nebula.controller.http_helpers import remote_get, remote_post_form from nebula.utils import APIUtils, DockerUtils +from nebula.database.utils_requests import ( + NodesUpdateRequest, + factory_requests_path, + ScenarioUpdateRequest, + ScenarioStopRequest, + ScenarioRemoveRequest, + ScenarioFinishRequest, + NotesUpdateRequest, + NotesRemoveRequest, + NodesRemoveRequest, + UserAddRequest, + UserDeleteRequest, + UserUpdateRequest, + UserVerifyRequest, + Routes, + RunScenarioRequest, +) # URL for the database API DATABASE_API_URL = os.environ.get("NEBULA_DATABASE_API_URL", "http://nebula-database:5051") @@ -124,7 +141,7 @@ async def lifespan(app: FastAPI): # Define endpoints outside the Controller class -@app.get("/") +@app.get(Routes.INIT) async def read_root(): """ Root endpoint of the NEBULA Controller API. @@ -135,7 +152,7 @@ async def read_root(): return {"message": "Welcome to the NEBULA Controller API"} -@app.get("/status") +@app.get(Routes.STATUS) async def get_status(): """ Check the status of the NEBULA Controller API. @@ -146,7 +163,7 @@ async def get_status(): return {"status": "NEBULA Controller API is running"} -@app.get("/resources") +@app.get(Routes.RESOURCES) async def get_resources(): """ Get system resource usage including RAM and GPU memory usage. @@ -188,7 +205,7 @@ async def get_resources(): } -@app.get("/least_memory_gpu") +@app.get(Routes.LEAST_MEMORY_GPU) async def get_least_memory_gpu(): """ Identify the GPU with the highest memory usage above a threshold (50%). @@ -230,7 +247,7 @@ async def get_least_memory_gpu(): } -@app.get("/available_gpus/") +@app.get(Routes.AVAILABLE_GPUS) async def get_available_gpu(): """ Get the list of GPUs with memory usage below 5%. @@ -289,10 +306,8 @@ def validate_physical_fields(data: dict): raise HTTPException(status_code=400, detail=str(e)) -@app.post("/scenarios/run") -async def run_scenario( - scenario_data: dict = Body(..., embed=True), role: str = Body(..., embed=True), user: str = Body(..., embed=True) -): +@app.post(Routes.RUN) +async def run_scenario(run_scenario_request: RunScenarioRequest): """ Launches a new scenario based on the provided configuration. @@ -309,6 +324,10 @@ async def run_scenario( from nebula.controller.scenarios import ScenarioManagement + # Unpack request data (role is intentionally ignored for now) + scenario_data = run_scenario_request.scenario_data + user = run_scenario_request.user + validate_physical_fields(scenario_data) # Manager for the actual scenario @@ -340,7 +359,7 @@ async def run_scenario( return scenarioManagement.scenario_name -@app.post("/scenarios/stop") +@app.post(Routes.STOP) async def stop_scenario( scenario_name: str = Body(..., embed=True), all: bool = Body(False, embed=True), @@ -369,13 +388,15 @@ async def stop_scenario( ScenarioManagement.cleanup_scenario_containers() try: - await APIUtils.post(f"{DATABASE_API_URL}/scenarios/stop", data={"scenario_name": scenario_name, "all": all}) + path = factory_requests_path("stop") + payload = ScenarioStopRequest(scenario_name=scenario_name, all=all).dict() + await APIUtils.post(f"{DATABASE_API_URL}{path}", data=payload) except Exception as e: logging.exception(f"Error setting scenario {scenario_name} to finished: {e}") raise HTTPException(status_code=500, detail="Internal server error") -@app.post("/scenarios/remove") +@app.post(Routes.REMOVE) async def remove_scenario( scenario_name: str = Body(..., embed=True), ): @@ -391,7 +412,9 @@ async def remove_scenario( from nebula.controller.scenarios import ScenarioManagement try: - await APIUtils.post(f"{DATABASE_API_URL}/scenarios/remove", data={"scenario_name": scenario_name}) + path = factory_requests_path("remove") + payload = ScenarioRemoveRequest(scenario_name=scenario_name).dict() + await APIUtils.post(f"{DATABASE_API_URL}{path}", data=payload) ScenarioManagement.remove_files_by_scenario(scenario_name) except Exception as e: @@ -401,7 +424,7 @@ async def remove_scenario( return {"message": f"Scenario {scenario_name} removed successfully"} -@app.get("/scenarios/{user}/{role}") +@app.get(Routes.GET_SCENARIOS_BY_USER) async def get_scenarios( user: Annotated[str, Path(regex="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50, description="Valid username")], role: Annotated[str, Path(regex="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50, description="Valid role")], @@ -417,13 +440,14 @@ async def get_scenarios( dict: A list of scenarios and the currently running scenario. """ try: - return await APIUtils.get(f"{DATABASE_API_URL}/scenarios/{user}/{role}") + path = factory_requests_path("get_scenarios_by_user", user=user, role=role) + return await APIUtils.get(f"{DATABASE_API_URL}{path}") except Exception as e: logging.exception(f"Error obtaining scenarios: {e}") raise HTTPException(status_code=500, detail="Internal server error") -@app.post("/scenarios/update") +@app.post(Routes.UPDATE) async def update_scenario( scenario_name: str = Body(..., embed=True), start_time: str = Body(..., embed=True), @@ -447,21 +471,22 @@ async def update_scenario( dict: A message confirming the update. """ try: - payload = { - "scenario_name": scenario_name, - "start_time": start_time, - "end_time": end_time, - "scenario": scenario, - "status": status, - "username": username, - } - return await APIUtils.post(f"{DATABASE_API_URL}/scenarios/update", data=payload) + payload = ScenarioUpdateRequest( + scenario_name=scenario_name, + start_time=start_time, + end_time=end_time, + scenario=scenario, + status=status, + username=username, + ).dict() + path = factory_requests_path("update") + return await APIUtils.post(f"{DATABASE_API_URL}{path}", data=payload) except Exception as e: logging.exception(f"Error updating scenario {scenario_name}: {e}") raise HTTPException(status_code=500, detail="Internal server error") -@app.post("/scenarios/set_status_to_finished") +@app.post(Routes.FINISH) async def set_scenario_status_to_finished( scenario_name: str = Body(..., embed=True), all: bool = Body(False, embed=True) ): @@ -476,14 +501,15 @@ async def set_scenario_status_to_finished( dict: A message confirming the operation. """ try: - payload = {"scenario_name": scenario_name, "all": all} - return await APIUtils.post(f"{DATABASE_API_URL}/scenarios/set_status_to_finished", data=payload) + payload = ScenarioFinishRequest(scenario_name=scenario_name, all=all).dict() + path = factory_requests_path("finish") + return await APIUtils.post(f"{DATABASE_API_URL}{path}", data=payload) except Exception as e: logging.exception(f"Error setting scenario {scenario_name} to finished: {e}") raise HTTPException(status_code=500, detail="Internal server error") -@app.get("/scenarios/running") +@app.get(Routes.RUNNING) async def get_running_scenario_endpoint(get_all: bool = False): """ Retrieves the currently running scenario(s). @@ -495,13 +521,14 @@ async def get_running_scenario_endpoint(get_all: bool = False): dict or list: Running scenario(s) information. """ try: - return await APIUtils.get(f"{DATABASE_API_URL}/scenarios/running", params={"get_all": str(get_all)}) + path = factory_requests_path("running") + return await APIUtils.get(f"{DATABASE_API_URL}{path}", params={"get_all": str(get_all)}) except Exception as e: logging.exception(f"Error obtaining running scenario: {e}") raise HTTPException(status_code=500, detail="Internal server error") -@app.get("/scenarios/check/{role}/{scenario_name}") +@app.get(Routes.CHECK_SCENARIO) async def check_scenario( role: Annotated[str, Path(regex="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50, description="Valid role")], scenario_name: Annotated[ @@ -519,13 +546,14 @@ async def check_scenario( dict: Whether the scenario is allowed for the role. """ try: - return await APIUtils.get(f"{DATABASE_API_URL}/scenarios/check/{role}/{scenario_name}") + path = factory_requests_path("check_scenario", role=role, scenario_name=scenario_name) + return await APIUtils.get(f"{DATABASE_API_URL}{path}") except Exception as e: logging.exception(f"Error checking scenario with role: {e}") raise HTTPException(status_code=500, detail="Internal server error") -@app.get("/scenarios/{scenario_name}") +@app.get(Routes.GET_SCENARIOS_BY_SCENARIO_NAME) async def get_scenario_by_name_endpoint( scenario_name: Annotated[ str, Path(regex="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50, description="Valid scenario name") @@ -541,13 +569,14 @@ async def get_scenario_by_name_endpoint( dict: The scenario data. """ try: - return await APIUtils.get(f"{DATABASE_API_URL}/scenarios/{scenario_name}") + path = factory_requests_path("get_scenarios_by_scenario_name", scenario_name=scenario_name) + return await APIUtils.get(f"{DATABASE_API_URL}{path}") except Exception as e: logging.exception(f"Error obtaining scenario {scenario_name}: {e}") raise HTTPException(status_code=500, detail="Internal server error") -@app.get("/nodes/{scenario_name}") +@app.get(Routes.NODES_BY_SCENARIO_NAME) async def list_nodes_by_scenario_name_endpoint( scenario_name: Annotated[ str, Path(regex="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50, description="Valid scenario name") @@ -563,13 +592,14 @@ async def list_nodes_by_scenario_name_endpoint( list: List of nodes. """ try: - return await APIUtils.get(f"{DATABASE_API_URL}/nodes/{scenario_name}") + path = factory_requests_path("get_nodes_by_scenario_name", scenario_name=scenario_name) + return await APIUtils.get(f"{DATABASE_API_URL}{path}") except Exception as e: logging.exception(f"Error obtaining nodes: {e}") raise HTTPException(status_code=500, detail="Internal server error") -@app.post("/nodes/{scenario_name}/update") +@app.post(Routes.NODES_UPDATE_BY_SCENARIO) async def update_nodes( scenario_name: Annotated[ str, @@ -588,12 +618,18 @@ async def update_nodes( dict: Confirmation or response from the frontend. """ try: - config = await request.json() - timestamp = datetime.datetime.now() - config["timestamp"] = str(timestamp) + config:dict = await request.json() + config["timestamp"] = str(datetime.datetime.now()) + + mobility_args = config.get("mobility_args", None) + if not mobility_args: + config["mobility_args"] = {"38.0235", "-1.1744"} + # Validate and normalize payload + validated = NodesUpdateRequest(**config) - # Update the node in database - await APIUtils.post(f"{DATABASE_API_URL}/nodes/update", data=config) + # Update the node in database with validated data + path = factory_requests_path("update_nodes") + await APIUtils.post(f"{DATABASE_API_URL}{path}", data=validated.dict()) except Exception as e: logging.exception(f"Error updating nodes: {e}") @@ -606,7 +642,7 @@ async def update_nodes( return await APIUtils.post(url, data=config) -@app.post("/nodes/{scenario_name}/done") +@app.post(Routes.NODES_DONE_BY_SCENARIO) async def node_done( scenario_name: Annotated[ str, @@ -633,7 +669,7 @@ async def node_done( return await APIUtils.post(url, data=data) -@app.post("/nodes/remove") +@app.post(Routes.NODES_REMOVE) async def remove_nodes_by_scenario_name_endpoint(scenario_name: str = Body(..., embed=True)): """ Endpoint to remove all nodes associated with a scenario. @@ -644,7 +680,9 @@ async def remove_nodes_by_scenario_name_endpoint(scenario_name: str = Body(..., Returns a success message or an error if something goes wrong. """ try: - await APIUtils.post(f"{DATABASE_API_URL}/nodes/remove", data={"scenario_name": scenario_name}) + path = factory_requests_path("remove_nodes") + payload = NodesRemoveRequest(scenario_name=scenario_name).dict() + await APIUtils.post(f"{DATABASE_API_URL}{path}", data=payload) except Exception as e: logging.exception(f"Error removing nodes: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -652,7 +690,7 @@ async def remove_nodes_by_scenario_name_endpoint(scenario_name: str = Body(..., return {"message": f"Nodes for scenario {scenario_name} removed successfully"} -@app.get("/notes/{scenario_name}") +@app.get(Routes.NOTES_BY_SCENARIO_NAME) async def get_notes_by_scenario_name( scenario_name: Annotated[ str, Path(regex="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50, description="Valid scenario name") @@ -662,13 +700,14 @@ async def get_notes_by_scenario_name( Endpoint to retrieve notes associated with a scenario. """ try: - return await APIUtils.get(f"{DATABASE_API_URL}/notes/{scenario_name}") + path = factory_requests_path("get_notes_by_scenario_name", scenario_name=scenario_name) + return await APIUtils.get(f"{DATABASE_API_URL}{path}") except Exception as e: logging.exception(f"Error obtaining notes for scenario {scenario_name}: {e}") raise HTTPException(status_code=500, detail="Internal server error") -@app.post("/notes/update") +@app.post(Routes.NOTES_UPDATE) async def update_notes_by_scenario_name(scenario_name: str = Body(..., embed=True), notes: str = Body(..., embed=True)): """ Endpoint to update notes for a given scenario. @@ -680,14 +719,15 @@ async def update_notes_by_scenario_name(scenario_name: str = Body(..., embed=Tru Returns a success message or an error if something goes wrong. """ try: - payload = {"scenario_name": scenario_name, "notes": notes} - return await APIUtils.post(f"{DATABASE_API_URL}/notes/update", data=payload) + payload = NotesUpdateRequest(scenario_name=scenario_name, notes=notes).dict() + path = factory_requests_path("update_notes") + return await APIUtils.post(f"{DATABASE_API_URL}{path}", data=payload) except Exception as e: logging.exception(f"Error updating notes: {e}") raise HTTPException(status_code=500, detail="Internal server error") -@app.post("/notes/remove") +@app.post(Routes.NOTES_REMOVE) async def remove_notes_by_scenario_name_endpoint(scenario_name: str = Body(..., embed=True)): """ Endpoint to remove notes associated with a scenario. @@ -698,7 +738,9 @@ async def remove_notes_by_scenario_name_endpoint(scenario_name: str = Body(..., Returns a success message or an error if something goes wrong. """ try: - await APIUtils.post(f"{DATABASE_API_URL}/notes/remove", data={"scenario_name": scenario_name}) + path = factory_requests_path("remove_notes") + payload = NotesRemoveRequest(scenario_name=scenario_name).dict() + await APIUtils.post(f"{DATABASE_API_URL}{path}", data=payload) except Exception as e: logging.exception(f"Error removing notes: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -706,7 +748,7 @@ async def remove_notes_by_scenario_name_endpoint(scenario_name: str = Body(..., return {"message": f"Notes for scenario {scenario_name} removed successfully"} -@app.get("/user/list") +@app.get(Routes.USER_LIST) async def list_users_controller(all_info: bool = False): """ Endpoint to list all users in the database. @@ -717,13 +759,14 @@ async def list_users_controller(all_info: bool = False): Returns a list of users or raises an HTTPException on error. """ try: - return await APIUtils.get(f"{DATABASE_API_URL}/user/list", params={"all_info": str(all_info)}) + path = factory_requests_path("list_users") + return await APIUtils.get(f"{DATABASE_API_URL}{path}", params={"all_info": str(all_info)}) except Exception as e: logging.exception(f"Error retrieving users: {e}") raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error retrieving users: {e}") -@app.get("/user/{scenario_name}") +@app.get(Routes.USER_BY_SCENARIO_NAME) async def get_user_by_scenario_name_endpoint( scenario_name: Annotated[ str, Path(regex="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50, description="Valid scenario name") @@ -738,13 +781,14 @@ async def get_user_by_scenario_name_endpoint( Returns user info or raises an HTTPException on error. """ try: - return await APIUtils.get(f"{DATABASE_API_URL}/user/{scenario_name}") + path = factory_requests_path("get_user_by_scenario_name", scenario_name=scenario_name) + return await APIUtils.get(f"{DATABASE_API_URL}{path}") except Exception as e: logging.exception(f"Error obtaining user for scenario {scenario_name}: {e}") raise HTTPException(status_code=500, detail="Internal server error") -@app.get("/discover-vpn") +@app.get(Routes.DISCOVER_VPN) async def discover_vpn(): """ Calls the Tailscale CLI to fetch the current status in JSON format, @@ -785,7 +829,7 @@ async def discover_vpn(): raise HTTPException(status_code=500, detail="No devices discovered") -@app.get("/physical/run/{ip}", tags=["physical"]) +@app.get(Routes.PHYSICAL_RUN, tags=["physical"]) async def physical_run(ip: str): status, data = await remote_get(ip, "/run/") @@ -796,7 +840,7 @@ async def physical_run(ip: str): raise HTTPException(status_code=status, detail=data) -@app.get("/physical/stop/{ip}", tags=["physical"]) +@app.get(Routes.PHYSICAL_STOP, tags=["physical"]) async def physical_stop(ip: str): status, data = await remote_get(ip, "/stop/") if status == 200: @@ -806,7 +850,7 @@ async def physical_stop(ip: str): raise HTTPException(status_code=status, detail=data) -@app.put("/physical/setup/{ip}", tags=["physical"], +@app.put(Routes.PHYSICAL_SETUP, tags=["physical"], status_code=status.HTTP_201_CREATED) async def physical_setup( ip: str, @@ -839,7 +883,7 @@ async def physical_setup( # ────────────────────────────────────────────────────────────── # Physical · single-node state # ────────────────────────────────────────────────────────────── -@app.get("/physical/state/{ip}", tags=["physical"]) +@app.get(Routes.PHYSICAL_STATE, tags=["physical"]) async def get_physical_node_state(ip: str): """ Query a single Raspberry Pi (or other node) for its training state. @@ -876,7 +920,7 @@ async def get_physical_node_state(ip: str): # ────────────────────────────────────────────────────────────── # Physical · aggregate state for an entire scenario # ────────────────────────────────────────────────────────────── -@app.get("/physical/scenario-state/{scenario_name}", tags=["physical"]) +@app.get(Routes.PHYSICAL_SCENARIO_STATE, tags=["physical"]) async def get_physical_scenario_state(scenario_name: str): """ Check the training state of *every* physical node assigned to a scenario. @@ -926,7 +970,7 @@ async def get_physical_scenario_state(scenario_name: str): } -@app.post("/user/add") +@app.post(Routes.USER_ADD) async def add_user_controller(user: str = Body(...), password: str = Body(...), role: str = Body(...)): """ Endpoint to add a new user to the database. @@ -939,14 +983,15 @@ async def add_user_controller(user: str = Body(...), password: str = Body(...), Returns a success message or an error if the user could not be added. """ try: - payload = {"user": user, "password": password, "role": role} - return await APIUtils.post(f"{DATABASE_API_URL}/user/add", data=payload) + payload = UserAddRequest(user=user, password=password, role=role).dict() + path = factory_requests_path("add_user") + return await APIUtils.post(f"{DATABASE_API_URL}{path}", data=payload) except Exception as e: logging.exception(f"Error adding user: {e}") raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error adding user: {e}") -@app.post("/user/delete") +@app.post(Routes.USER_DELETE) async def remove_user_controller(user: str = Body(..., embed=True)): """ Controller endpoint that inserts a new user into the database. @@ -957,13 +1002,15 @@ async def remove_user_controller(user: str = Body(..., embed=True)): Returns a success message if the user is deleted, or an HTTP error if an exception occurs. """ try: - return await APIUtils.post(f"{DATABASE_API_URL}/user/delete", data={"user": user}) + path = factory_requests_path("delete_user") + payload = UserDeleteRequest(user=user).dict() + return await APIUtils.post(f"{DATABASE_API_URL}{path}", data=payload) except Exception as e: logging.exception(f"Error deleting user: {e}") raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error deleting user: {e}") -@app.post("/user/update") +@app.post(Routes.USER_UPDATE) async def update_user_controller(user: str = Body(...), password: str = Body(...), role: str = Body(...)): """ Controller endpoint that modifies a user of the database. @@ -976,14 +1023,15 @@ async def update_user_controller(user: str = Body(...), password: str = Body(... Returns a success message if the user is updated, or an HTTP error if an exception occurs. """ try: - payload = {"user": user, "password": password, "role": role} - return await APIUtils.post(f"{DATABASE_API_URL}/user/update", data=payload) + payload = UserUpdateRequest(user=user, password=password, role=role).dict() + path = factory_requests_path("update_user") + return await APIUtils.post(f"{DATABASE_API_URL}{path}", data=payload) except Exception as e: logging.exception(f"Error updating user: {e}") raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error updating user: {e}") -@app.post("/user/verify") +@app.post(Routes.USER_VERIFY) async def verify_user_controller(user: str = Body(...), password: str = Body(...)): """ Endpoint to verify user credentials. @@ -995,8 +1043,9 @@ async def verify_user_controller(user: str = Body(...), password: str = Body(... Returns the user role on success or raises an error on failure. """ try: - payload = {"user": user, "password": password} - return await APIUtils.post(f"{DATABASE_API_URL}/user/verify", data=payload) + payload = UserVerifyRequest(user=user, password=password).dict() + path = factory_requests_path("verify_user") + return await APIUtils.post(f"{DATABASE_API_URL}{path}", data=payload) except HTTPException as e: if e.status_code == 401: raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED) from e diff --git a/nebula/database/database_api.py b/nebula/database/database_api.py index aac0a981c..f92dfefa1 100644 --- a/nebula/database/database_api.py +++ b/nebula/database/database_api.py @@ -11,6 +11,7 @@ from fastapi.concurrency import asynccontextmanager from nebula.database.database_adapter_factory import factory_database_adapter +from nebula.database.utils_requests import Routes # Get a database instance db = factory_database_adapter("PostgresDB") @@ -67,13 +68,13 @@ async def lifespan(app: FastAPI): app = FastAPI(lifespan=lifespan) -@app.get("/") +@app.get(Routes.INIT) async def read_root(): return {"message": "Welcome to the NEBULA Database API"} # Scenarios -@app.post("/scenarios/update") +@app.post(Routes.UPDATE) async def update_scenario( scenario_name: str = Body(..., embed=True), start_time: str = Body(..., embed=True), @@ -90,7 +91,7 @@ async def update_scenario( raise HTTPException(status_code=500, detail="Internal server error") -@app.post("/scenarios/stop") +@app.post(Routes.STOP) async def stop_scenario( scenario_name: str = Body(..., embed=True), all: bool = Body(False, embed=True), @@ -106,7 +107,7 @@ async def stop_scenario( raise HTTPException(status_code=500, detail="Internal server error") -@app.post("/scenarios/remove") +@app.post(Routes.REMOVE) async def remove_scenario( scenario_name: str = Body(..., embed=True), ): @@ -118,7 +119,7 @@ async def remove_scenario( raise HTTPException(status_code=500, detail="Internal server error") -@app.get("/scenarios/{user}/{role}") +@app.get(Routes.GET_SCENARIOS_BY_USER) async def get_scenarios( user: Annotated[str, Path(pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50)], role: Annotated[str, Path(pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50)], @@ -135,7 +136,7 @@ async def get_scenarios( raise HTTPException(status_code=500, detail="Internal server error") -@app.post("/scenarios/set_status_to_finished") +@app.post(Routes.FINISH) async def set_scenario_status_to_finished( scenario_name: str = Body(..., embed=True), all: bool = Body(False, embed=True) ): @@ -150,7 +151,7 @@ async def set_scenario_status_to_finished( raise HTTPException(status_code=500, detail="Internal server error") -@app.get("/scenarios/running") +@app.get(Routes.RUNNING) async def get_running_scenario_endpoint(get_all: bool = False): try: return await db.get_running_scenario(get_all=get_all) @@ -159,7 +160,7 @@ async def get_running_scenario_endpoint(get_all: bool = False): raise HTTPException(status_code=500, detail="Internal server error") -@app.get("/scenarios/check/{role}/{scenario_name}") +@app.get(Routes.CHECK_SCENARIO) async def check_scenario( role: Annotated[str, Path(pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50)], scenario_name: Annotated[str, Path(pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50)], @@ -172,7 +173,7 @@ async def check_scenario( raise HTTPException(status_code=500, detail="Internal server error") -@app.get("/scenarios/{scenario_name}") +@app.get(Routes.GET_SCENARIOS_BY_SCENARIO_NAME) async def get_scenario_by_name_endpoint( scenario_name: Annotated[str, Path(pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50)], ): @@ -185,7 +186,7 @@ async def get_scenario_by_name_endpoint( # Nodes -@app.get("/nodes/{scenario_name}") +@app.get(Routes.NODES_BY_SCENARIO_NAME) async def list_nodes_by_scenario_name_endpoint( scenario_name: Annotated[str, Path(pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50)], ): @@ -197,7 +198,7 @@ async def list_nodes_by_scenario_name_endpoint( raise HTTPException(status_code=500, detail="Internal server error") -@app.post("/nodes/update") +@app.post(Routes.NODES_UPDATE) async def update_node_record(data: dict): try: await db.update_node_record( @@ -222,7 +223,7 @@ async def update_node_record(data: dict): raise HTTPException(status_code=500, detail="Internal server error") -@app.post("/nodes/remove") +@app.post(Routes.NODES_REMOVE) async def remove_nodes_by_scenario_name_endpoint(scenario_name: str = Body(..., embed=True)): try: await db.remove_nodes_by_scenario_name(scenario_name) @@ -233,7 +234,7 @@ async def remove_nodes_by_scenario_name_endpoint(scenario_name: str = Body(..., # Notes -@app.get("/notes/{scenario_name}") +@app.get(Routes.NOTES_BY_SCENARIO_NAME) async def get_notes_by_scenario_name( scenario_name: Annotated[str, Path(pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50)], ): @@ -247,7 +248,7 @@ async def get_notes_by_scenario_name( raise HTTPException(status_code=500, detail="Internal server error") -@app.post("/notes/update") +@app.post(Routes.NOTES_UPDATE) async def update_notes_by_scenario_name(scenario_name: str = Body(..., embed=True), notes: str = Body(..., embed=True)): try: await db.save_notes(scenario_name, notes) @@ -257,7 +258,7 @@ async def update_notes_by_scenario_name(scenario_name: str = Body(..., embed=Tru raise HTTPException(status_code=500, detail="Internal server error") -@app.post("/notes/remove") +@app.post(Routes.NOTES_REMOVE) async def remove_notes_by_scenario_name_endpoint(scenario_name: str = Body(..., embed=True)): try: await db.remove_note(scenario_name) @@ -268,7 +269,7 @@ async def remove_notes_by_scenario_name_endpoint(scenario_name: str = Body(..., # Users -@app.get("/user/list") +@app.get(Routes.USER_LIST) async def list_users_controller(all_info: bool = False): try: user_list = await db.list_users(all_info) @@ -280,7 +281,7 @@ async def list_users_controller(all_info: bool = False): raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error retrieving users: {e}") -@app.get("/user/{scenario_name}") +@app.get(Routes.USER_BY_SCENARIO_NAME) async def get_user_by_scenario_name_endpoint( scenario_name: Annotated[str, Path(pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50)], ): @@ -292,7 +293,7 @@ async def get_user_by_scenario_name_endpoint( raise HTTPException(status_code=500, detail="Internal server error") -@app.post("/user/add") +@app.post(Routes.USER_ADD) async def add_user_controller(user: str = Body(...), password: str = Body(...), role: str = Body(...)): try: await db.add_user(user, password, role) @@ -302,7 +303,7 @@ async def add_user_controller(user: str = Body(...), password: str = Body(...), raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error adding user: {e}") -@app.post("/user/delete") +@app.post(Routes.USER_DELETE) async def remove_user_controller(user: str = Body(..., embed=True)): try: await db.delete_user_from_db(user) @@ -312,7 +313,7 @@ async def remove_user_controller(user: str = Body(..., embed=True)): raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error deleting user: {e}") -@app.post("/user/update") +@app.post(Routes.USER_UPDATE) async def update_user_controller(user: str = Body(...), password: str = Body(...), role: str = Body(...)): try: await db.update_user(user, password, role) @@ -322,7 +323,7 @@ async def update_user_controller(user: str = Body(...), password: str = Body(... raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error updating user: {e}") -@app.post("/user/verify") +@app.post(Routes.USER_VERIFY) async def verify_user_controller(user: str = Body(...), password: str = Body(...)): try: user_submitted = user.upper() diff --git a/nebula/database/utils_requests.py b/nebula/database/utils_requests.py new file mode 100644 index 000000000..25060322f --- /dev/null +++ b/nebula/database/utils_requests.py @@ -0,0 +1,209 @@ +from typing import Any, Dict, List + +from pydantic import BaseModel, conint, confloat + + +class Routes: + # Scenarios + INIT = "/" + STATUS = "/status" + RESOURCES = "/resources" + LEAST_MEMORY_GPU = "/least_memory_gpu" + AVAILABLE_GPUS = "/available_gpus/" + DISCOVER_VPN = "/discover-vpn" + + RUN = "/scenarios/run" + UPDATE = "/scenarios/update" + STOP = "/scenarios/stop" + REMOVE = "/scenarios/remove" + FINISH = "/scenarios/set_status_to_finished" + RUNNING = "/scenarios/running" + CHECK_SCENARIO = "/scenarios/check/{role}/{scenario_name}" + GET_SCENARIOS_BY_USER = "/scenarios/{user}/{role}" + GET_SCENARIOS_BY_SCENARIO_NAME = "/scenarios/{scenario_name}" + + # Nodes + NODES_BY_SCENARIO_NAME = "/nodes/{scenario_name}" + NODES_UPDATE = "/nodes/update" + NODES_REMOVE = "/nodes/remove" + NODES_UPDATE_BY_SCENARIO = "/nodes/{scenario_name}/update" + NODES_DONE_BY_SCENARIO = "/nodes/{scenario_name}/done" + + # Notes + NOTES_BY_SCENARIO_NAME = "/notes/{scenario_name}" + NOTES_UPDATE = "/notes/update" + NOTES_REMOVE = "/notes/remove" + + # Users + USER_LIST = "/user/list" + USER_BY_SCENARIO_NAME = "/user/{scenario_name}" + USER_ADD = "/user/add" + USER_DELETE = "/user/delete" + USER_UPDATE = "/user/update" + USER_VERIFY = "/user/verify" + + # Physical + PHYSICAL_RUN = "/physical/run/{ip}" + PHYSICAL_STOP = "/physical/stop/{ip}" + PHYSICAL_SETUP = "/physical/setup/{ip}" + PHYSICAL_STATE = "/physical/state/{ip}" + PHYSICAL_SCENARIO_STATE = "/physical/scenario-state/{scenario_name}" + + +class RunScenarioRequest(BaseModel): + """Request model to trigger a scenario run on the controller. + + - Only requires scenario_data and user. + - Extra fields (e.g., role, federation_id) are ignored. + """ + scenario_data: Dict[str, Any] + user: str + + +class ScenarioUpdateRequest(BaseModel): + scenario_name: str + start_time: str + end_time: str + scenario: Dict[str, Any] + status: str + username: str + + +class ScenarioStopRequest(BaseModel): + scenario_name: str + all: bool = False + + +class ScenarioRemoveRequest(BaseModel): + scenario_name: str + + +class ScenarioFinishRequest(BaseModel): + scenario_name: str + all: bool = False + + +class NotesUpdateRequest(BaseModel): + scenario_name: str + notes: str + + +class NotesRemoveRequest(BaseModel): + scenario_name: str + + +class NodesRemoveRequest(BaseModel): + scenario_name: str + + +class UserAddRequest(BaseModel): + user: str + password: str + role: str + + +class UserDeleteRequest(BaseModel): + user: str + + +class UserUpdateRequest(BaseModel): + user: str + password: str + role: str + + +class UserVerifyRequest(BaseModel): + user: str + password: str + + +# Nodes update payload +class DeviceArgs(BaseModel): + uid: str + idx: int + role: str + malicious: bool + + +class NetworkArgs(BaseModel): + ip: str + port: conint(ge=1, le=65535) # type: ignore[valid-type] + neighbors: List[Any] + + +class MobilityArgs(BaseModel): + latitude: confloat(ge=-90, le=90) # type: ignore[valid-type] + longitude: confloat(ge=-180, le=180) # type: ignore[valid-type] + + +class TrackingArgs(BaseModel): + run_hash: str + + +class FederationArgs(BaseModel): + round: int + + +class ScenarioArgs(BaseModel): + federation: str + name: str + + +class NodesUpdateRequest(BaseModel): + device_args: DeviceArgs + network_args: NetworkArgs + mobility_args: MobilityArgs + tracking_args: TrackingArgs + federation_args: FederationArgs + scenario_args: ScenarioArgs + timestamp: str + + +def factory_requests_path(resource: str, user: str = "", role: str = "", scenario_name: str = "") -> str: + if resource == "init": + return Routes.INIT + elif resource == "update": + return Routes.UPDATE + elif resource == "stop": + return Routes.STOP + elif resource == "remove": + return Routes.REMOVE + elif resource == "finish": + return Routes.FINISH + elif resource == "running": + return Routes.RUNNING + elif resource == "check_scenario": + return Routes.CHECK_SCENARIO.format(role=role, scenario_name=scenario_name) + elif resource == "get_scenarios_by_user": + return Routes.GET_SCENARIOS_BY_USER.format(user=user, role=role) + elif resource == "get_scenarios_by_scenario_name": + return Routes.GET_SCENARIOS_BY_SCENARIO_NAME.format(scenario_name=scenario_name) + # Nodes + elif resource == "get_nodes_by_scenario_name": + return Routes.NODES_BY_SCENARIO_NAME.format(scenario_name=scenario_name) + elif resource == "update_nodes": + return Routes.NODES_UPDATE + elif resource == "remove_nodes": + return Routes.NODES_REMOVE + # Notes + elif resource == "get_notes_by_scenario_name": + return Routes.NOTES_BY_SCENARIO_NAME.format(scenario_name=scenario_name) + elif resource == "update_notes": + return Routes.NOTES_UPDATE + elif resource == "remove_notes": + return Routes.NOTES_REMOVE + # Users + elif resource == "list_users": + return Routes.USER_LIST + elif resource == "get_user_by_scenario_name": + return Routes.USER_BY_SCENARIO_NAME.format(scenario_name=scenario_name) + elif resource == "add_user": + return Routes.USER_ADD + elif resource == "delete_user": + return Routes.USER_DELETE + elif resource == "update_user": + return Routes.USER_UPDATE + elif resource == "verify_user": + return Routes.USER_VERIFY + else: + raise Exception(f"resource not found: {resource}") From e41ba4c54567435301972b83ec6c34b15d01e732 Mon Sep 17 00:00:00 2001 From: FerTV Date: Thu, 18 Sep 2025 10:37:11 +0200 Subject: [PATCH 06/14] refactor: extras column added to nodes table --- nebula/controller/hub.py | 33 ++++--- .../postgress/docker/init-configs.sql | 13 ++- .../database/adapters/postgress/postgress.py | 87 ++++++++++++++++--- nebula/database/database_adapter_interface.py | 18 +++- nebula/database/database_api.py | 6 +- 5 files changed, 122 insertions(+), 35 deletions(-) diff --git a/nebula/controller/hub.py b/nebula/controller/hub.py index 9bb3c8cee..cc8cbdb36 100755 --- a/nebula/controller/hub.py +++ b/nebula/controller/hub.py @@ -389,7 +389,7 @@ async def stop_scenario( ScenarioManagement.cleanup_scenario_containers() try: path = factory_requests_path("stop") - payload = ScenarioStopRequest(scenario_name=scenario_name, all=all).dict() + payload = ScenarioStopRequest(scenario_name=scenario_name, all=all).model_dump() await APIUtils.post(f"{DATABASE_API_URL}{path}", data=payload) except Exception as e: logging.exception(f"Error setting scenario {scenario_name} to finished: {e}") @@ -413,7 +413,7 @@ async def remove_scenario( try: path = factory_requests_path("remove") - payload = ScenarioRemoveRequest(scenario_name=scenario_name).dict() + payload = ScenarioRemoveRequest(scenario_name=scenario_name).model_dump() await APIUtils.post(f"{DATABASE_API_URL}{path}", data=payload) ScenarioManagement.remove_files_by_scenario(scenario_name) @@ -478,7 +478,7 @@ async def update_scenario( scenario=scenario, status=status, username=username, - ).dict() + ).model_dump() path = factory_requests_path("update") return await APIUtils.post(f"{DATABASE_API_URL}{path}", data=payload) except Exception as e: @@ -501,7 +501,7 @@ async def set_scenario_status_to_finished( dict: A message confirming the operation. """ try: - payload = ScenarioFinishRequest(scenario_name=scenario_name, all=all).dict() + payload = ScenarioFinishRequest(scenario_name=scenario_name, all=all).model_dump() path = factory_requests_path("finish") return await APIUtils.post(f"{DATABASE_API_URL}{path}", data=payload) except Exception as e: @@ -623,13 +623,18 @@ async def update_nodes( mobility_args = config.get("mobility_args", None) if not mobility_args: - config["mobility_args"] = {"38.0235", "-1.1744"} + # default Murcia coordinates if none provided + config["mobility_args"] = {"latitude": "38.0235", "longitude": "-1.1744"} # Validate and normalize payload validated = NodesUpdateRequest(**config) - # Update the node in database with validated data + # Build payload and include extras with mobility data + payload = validated.model_dump() + payload["extras"] = payload.get("mobility_args", {}) + + # Update the node in database with validated data and extras path = factory_requests_path("update_nodes") - await APIUtils.post(f"{DATABASE_API_URL}{path}", data=validated.dict()) + await APIUtils.post(f"{DATABASE_API_URL}{path}", data=payload) except Exception as e: logging.exception(f"Error updating nodes: {e}") @@ -681,7 +686,7 @@ async def remove_nodes_by_scenario_name_endpoint(scenario_name: str = Body(..., """ try: path = factory_requests_path("remove_nodes") - payload = NodesRemoveRequest(scenario_name=scenario_name).dict() + payload = NodesRemoveRequest(scenario_name=scenario_name).model_dump() await APIUtils.post(f"{DATABASE_API_URL}{path}", data=payload) except Exception as e: logging.exception(f"Error removing nodes: {e}") @@ -719,7 +724,7 @@ async def update_notes_by_scenario_name(scenario_name: str = Body(..., embed=Tru Returns a success message or an error if something goes wrong. """ try: - payload = NotesUpdateRequest(scenario_name=scenario_name, notes=notes).dict() + payload = NotesUpdateRequest(scenario_name=scenario_name, notes=notes).model_dump() path = factory_requests_path("update_notes") return await APIUtils.post(f"{DATABASE_API_URL}{path}", data=payload) except Exception as e: @@ -739,7 +744,7 @@ async def remove_notes_by_scenario_name_endpoint(scenario_name: str = Body(..., """ try: path = factory_requests_path("remove_notes") - payload = NotesRemoveRequest(scenario_name=scenario_name).dict() + payload = NotesRemoveRequest(scenario_name=scenario_name).model_dump() await APIUtils.post(f"{DATABASE_API_URL}{path}", data=payload) except Exception as e: logging.exception(f"Error removing notes: {e}") @@ -983,7 +988,7 @@ async def add_user_controller(user: str = Body(...), password: str = Body(...), Returns a success message or an error if the user could not be added. """ try: - payload = UserAddRequest(user=user, password=password, role=role).dict() + payload = UserAddRequest(user=user, password=password, role=role).model_dump() path = factory_requests_path("add_user") return await APIUtils.post(f"{DATABASE_API_URL}{path}", data=payload) except Exception as e: @@ -1003,7 +1008,7 @@ async def remove_user_controller(user: str = Body(..., embed=True)): """ try: path = factory_requests_path("delete_user") - payload = UserDeleteRequest(user=user).dict() + payload = UserDeleteRequest(user=user).model_dump() return await APIUtils.post(f"{DATABASE_API_URL}{path}", data=payload) except Exception as e: logging.exception(f"Error deleting user: {e}") @@ -1023,7 +1028,7 @@ async def update_user_controller(user: str = Body(...), password: str = Body(... Returns a success message if the user is updated, or an HTTP error if an exception occurs. """ try: - payload = UserUpdateRequest(user=user, password=password, role=role).dict() + payload = UserUpdateRequest(user=user, password=password, role=role).model_dump() path = factory_requests_path("update_user") return await APIUtils.post(f"{DATABASE_API_URL}{path}", data=payload) except Exception as e: @@ -1043,7 +1048,7 @@ async def verify_user_controller(user: str = Body(...), password: str = Body(... Returns the user role on success or raises an error on failure. """ try: - payload = UserVerifyRequest(user=user, password=password).dict() + payload = UserVerifyRequest(user=user, password=password).model_dump() path = factory_requests_path("verify_user") return await APIUtils.post(f"{DATABASE_API_URL}{path}", data=payload) except HTTPException as e: diff --git a/nebula/database/adapters/postgress/docker/init-configs.sql b/nebula/database/adapters/postgress/docker/init-configs.sql index 9a452af93..9370a31d7 100644 --- a/nebula/database/adapters/postgress/docker/init-configs.sql +++ b/nebula/database/adapters/postgress/docker/init-configs.sql @@ -21,16 +21,25 @@ CREATE TABLE IF NOT EXISTS nodes ( port TEXT, role TEXT, neighbors TEXT[], - latitude TEXT, - longitude TEXT, timestamp TEXT, federation TEXT, round TEXT, scenario TEXT, hash TEXT, + extras JSONB, malicious TEXT ); +-- Ensure column exists for pre-existing installations +ALTER TABLE IF EXISTS nodes + ADD COLUMN IF NOT EXISTS extras JSONB; + +-- Drop legacy columns for latitude/longitude if present +ALTER TABLE IF EXISTS nodes + DROP COLUMN IF EXISTS latitude; +ALTER TABLE IF EXISTS nodes + DROP COLUMN IF EXISTS longitude; + -- 3) Configs as JSONB DROP INDEX IF EXISTS idx_configs_config_gin; DROP TABLE IF EXISTS configs; diff --git a/nebula/database/adapters/postgress/postgress.py b/nebula/database/adapters/postgress/postgress.py index 57019164f..65958ef7b 100755 --- a/nebula/database/adapters/postgress/postgress.py +++ b/nebula/database/adapters/postgress/postgress.py @@ -197,7 +197,24 @@ async def list_nodes(self, scenario_name=None, sort_by="idx"): else: command = f"SELECT * FROM nodes ORDER BY {sort_by};" result = await conn.fetch(command) - return result + + # Convert to list of dicts and expose latitude/longitude from extras for compatibility + rows = [] + for record in result: + row = dict(record) + extras = row.get("extras") + if isinstance(extras, str): + try: + extras = json.loads(extras) + except json.JSONDecodeError: + extras = None + if isinstance(extras, dict): + if "latitude" in extras and "latitude" not in row: + row["latitude"] = extras.get("latitude") + if "longitude" in extras and "longitude" not in row: + row["longitude"] = extras.get("longitude") + rows.append(row) + return rows except asyncpg.PostgresError as e: logging.error(f"Error occurred while listing nodes: {e}") return None @@ -211,15 +228,42 @@ async def list_nodes_by_scenario_name(self, scenario_name): async with self.pool.acquire() as conn: command = "SELECT * FROM nodes WHERE scenario = $1 ORDER BY CAST(idx AS INTEGER) ASC;" result = await conn.fetch(command, scenario_name) - return [dict(record) for record in result] + rows = [] + for record in result: + row = dict(record) + extras = row.get("extras") + if isinstance(extras, str): + try: + extras = json.loads(extras) + except json.JSONDecodeError: + extras = None + if isinstance(extras, dict): + if "latitude" in extras and "latitude" not in row: + row["latitude"] = extras.get("latitude") + if "longitude" in extras and "longitude" not in row: + row["longitude"] = extras.get("longitude") + rows.append(row) + return rows except Exception as e: logging.error(f"Error occurred while listing nodes by scenario name: {e}") return None async def update_node_record( - self, node_uid, idx, ip, port, role, neighbors, latitude, longitude, - timestamp, federation, federation_round, scenario, run_hash, malicious, + self, + node_uid, + idx, + ip, + port, + role, + neighbors, + extras, + timestamp, + federation, + federation_round, + scenario, + run_hash, + malicious, ): """ Inserts or updates a node record in the database for a given scenario, ensuring thread-safe access. @@ -227,6 +271,19 @@ async def update_node_record( async with _node_lock: async with self.pool.acquire() as conn: try: + # Ensure `extras` is a JSON string when provided + extras_payload = None + if extras is not None: + if isinstance(extras, str): + extras_payload = extras + else: + try: + extras_payload = json.dumps(extras) + except (TypeError, ValueError): + # Fallback to empty JSON object on serialization issues + logging.warning("Unable to serialize extras to JSON, storing as empty object.") + extras_payload = json.dumps({}) + async with conn.transaction(): result = await conn.fetchrow( "SELECT * FROM nodes WHERE uid = $1 AND scenario = $2 FOR UPDATE;", @@ -237,24 +294,26 @@ async def update_node_record( # Insert new node await conn.execute( """ - INSERT INTO nodes (uid, idx, ip, port, role, neighbors, latitude, longitude, - timestamp, federation, round, scenario, hash, malicious) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14); + INSERT INTO nodes (uid, idx, ip, port, role, neighbors, + timestamp, federation, round, scenario, hash, extras, malicious) + VALUES ($1, $2, $3, $4, $5, $6, + $7, $8, $9, $10, $11, $12::jsonb, $13); """, - node_uid, idx, ip, port, role, neighbors, latitude, longitude, - timestamp, federation, federation_round, scenario, run_hash, malicious, + node_uid, idx, ip, port, role, neighbors, + timestamp, federation, federation_round, scenario, run_hash, extras_payload, malicious, ) else: # Update existing node await conn.execute( """ UPDATE nodes SET idx = $1, ip = $2, port = $3, role = $4, neighbors = $5, - latitude = $6, longitude = $7, timestamp = $8, federation = $9, round = $10, - hash = $11, malicious = $12 - WHERE uid = $13 AND scenario = $14; + timestamp = $6, federation = $7, round = $8, + hash = $9, extras = $10::jsonb, malicious = $11 + WHERE uid = $12 AND scenario = $13; """, - idx, ip, port, role, neighbors, latitude, longitude, - timestamp, federation, federation_round, run_hash, malicious, + idx, ip, port, role, neighbors, + timestamp, federation, federation_round, + run_hash, extras_payload, malicious, node_uid, scenario, ) diff --git a/nebula/database/database_adapter_interface.py b/nebula/database/database_adapter_interface.py index d527bb467..9c0286ef1 100644 --- a/nebula/database/database_adapter_interface.py +++ b/nebula/database/database_adapter_interface.py @@ -73,10 +73,22 @@ async def list_nodes_by_scenario_name(self, scenario_name): @abstractmethod async def update_node_record( - self, node_uid, idx, ip, port, role, neighbors, latitude, longitude, - timestamp, federation, federation_round, scenario, run_hash, malicious, + self, + node_uid, + idx, + ip, + port, + role, + neighbors, + extras, + timestamp, + federation, + federation_round, + scenario, + run_hash, + malicious, ): - """Inserts or updates a node record.""" + """Inserts or updates a node record. Latitude/longitude must be included in `extras` (JSON).""" raise NotImplementedError @abstractmethod diff --git a/nebula/database/database_api.py b/nebula/database/database_api.py index f92dfefa1..02df20898 100644 --- a/nebula/database/database_api.py +++ b/nebula/database/database_api.py @@ -201,6 +201,9 @@ async def list_nodes_by_scenario_name_endpoint( @app.post(Routes.NODES_UPDATE) async def update_node_record(data: dict): try: + # Store latitude/longitude (and any mobility details) inside extras JSONB + # Prefer explicit `extras` sent by controller; fallback to `mobility_args` for backward compatibility + extras = data.get("extras") or data.get("mobility_args", {}) await db.update_node_record( str(data["device_args"]["uid"]), str(data["device_args"]["idx"]), @@ -208,8 +211,7 @@ async def update_node_record(data: dict): str(data["network_args"]["port"]), str(data["device_args"]["role"]), data["network_args"]["neighbors"], - str(data["mobility_args"]["latitude"]), - str(data["mobility_args"]["longitude"]), + extras, str(data["timestamp"]), str(data["scenario_args"]["federation"]), str(data["federation_args"]["round"]), From 0572c0cf71006f30780af3355786cee107df6cf8 Mon Sep 17 00:00:00 2001 From: FerTV Date: Fri, 19 Sep 2025 13:25:33 +0200 Subject: [PATCH 07/14] refactor: database_api endpoints --- nebula/controller/hub.py | 5 +- nebula/controller/utils_requests.py | 214 +++++++++++++++++ .../database/adapters/postgress/postgress.py | 156 ++++++++----- nebula/database/database_adapter_interface.py | 74 +++--- nebula/database/database_api.py | 220 +++++++++--------- nebula/database/utils_requests.py | 66 +++--- 6 files changed, 505 insertions(+), 230 deletions(-) create mode 100644 nebula/controller/utils_requests.py diff --git a/nebula/controller/hub.py b/nebula/controller/hub.py index cc8cbdb36..23b954b3e 100755 --- a/nebula/controller/hub.py +++ b/nebula/controller/hub.py @@ -18,7 +18,7 @@ from nebula.controller.http_helpers import remote_get, remote_post_form from nebula.utils import APIUtils, DockerUtils -from nebula.database.utils_requests import ( +from nebula.controller.utils_requests import ( NodesUpdateRequest, factory_requests_path, ScenarioUpdateRequest, @@ -530,6 +530,7 @@ async def get_running_scenario_endpoint(get_all: bool = False): @app.get(Routes.CHECK_SCENARIO) async def check_scenario( + user: Annotated[str, Path(regex="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50, description="Valid username")], role: Annotated[str, Path(regex="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50, description="Valid role")], scenario_name: Annotated[ str, Path(regex="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50, description="Valid scenario name") @@ -546,7 +547,7 @@ async def check_scenario( dict: Whether the scenario is allowed for the role. """ try: - path = factory_requests_path("check_scenario", role=role, scenario_name=scenario_name) + path = factory_requests_path("check_scenario", user=user, role=role, scenario_name=scenario_name) return await APIUtils.get(f"{DATABASE_API_URL}{path}") except Exception as e: logging.exception(f"Error checking scenario with role: {e}") diff --git a/nebula/controller/utils_requests.py b/nebula/controller/utils_requests.py new file mode 100644 index 000000000..9a66a4c6e --- /dev/null +++ b/nebula/controller/utils_requests.py @@ -0,0 +1,214 @@ +from typing import Any, Dict, List + +from pydantic import BaseModel, conint, confloat + + +class Routes: + # General + INIT = "/" + STATUS = "/status" + RESOURCES = "/resources" + LEAST_MEMORY_GPU = "/least_memory_gpu" + AVAILABLE_GPUS = "/available_gpus/" + + # Scenarios (Controller + DB API routing) + RUN = "/scenarios/run" + UPDATE = "/scenarios/update" + STOP = "/scenarios/stop" + REMOVE = "/scenarios/remove" + FINISH = "/scenarios/set_status_to_finished" + RUNNING = "/scenarios/running" + CHECK_SCENARIO = "/scenarios/check/{user}/{role}/{scenario_name}" + GET_SCENARIOS_BY_USER = "/scenarios/{user}/{role}" + GET_SCENARIOS_BY_SCENARIO_NAME = "/scenarios/{scenario_name}" + + # Nodes + NODES_BY_SCENARIO_NAME = "/nodes/{scenario_name}" + NODES_UPDATE = "/nodes/update" + NODES_UPDATE_BY_SCENARIO = "/nodes/{scenario_name}/update" + NODES_DONE_BY_SCENARIO = "/nodes/{scenario_name}/done" + NODES_REMOVE = "/nodes/remove" + + # Notes + NOTES_BY_SCENARIO_NAME = "/notes/{scenario_name}" + NOTES_UPDATE = "/notes/update" + NOTES_REMOVE = "/notes/remove" + + # Users + USER_LIST = "/user/list" + USER_BY_SCENARIO_NAME = "/user/{scenario_name}" + USER_ADD = "/user/add" + USER_DELETE = "/user/delete" + USER_UPDATE = "/user/update" + USER_VERIFY = "/user/verify" + + # Discovery / Physical management + DISCOVER_VPN = "/discover-vpn" + PHYSICAL_RUN = "/physical/run" + PHYSICAL_STOP = "/physical/stop" + PHYSICAL_SETUP = "/physical/setup" + PHYSICAL_STATE = "/physical/state" + PHYSICAL_SCENARIO_STATE = "/physical/{scenario_name}/state" + + +class RunScenarioRequest(BaseModel): + """Request model to trigger a scenario run on the controller. + + - Only requires scenario_data and user. + - Extra fields (e.g., role, federation_id) are ignored. + """ + scenario_data: Dict[str, Any] + user: str + + +class ScenarioUpdateRequest(BaseModel): + scenario_name: str + start_time: str + end_time: str + scenario: Dict[str, Any] + status: str + username: str + + +class ScenarioStopRequest(BaseModel): + scenario_name: str + all: bool = False + + +class ScenarioRemoveRequest(BaseModel): + scenario_name: str + + +class ScenarioFinishRequest(BaseModel): + scenario_name: str + all: bool = False + + +class NotesUpdateRequest(BaseModel): + scenario_name: str + notes: str + + +class NotesRemoveRequest(BaseModel): + scenario_name: str + + +class NodesRemoveRequest(BaseModel): + scenario_name: str + + +class UserAddRequest(BaseModel): + user: str + password: str + role: str + + +class UserDeleteRequest(BaseModel): + user: str + + +class UserUpdateRequest(BaseModel): + user: str + password: str + role: str + + +class UserVerifyRequest(BaseModel): + user: str + password: str + + +# Nodes update payload +class DeviceArgs(BaseModel): + uid: str + idx: int + role: str + malicious: bool + + +class NetworkArgs(BaseModel): + ip: str + port: conint(ge=1, le=65535) # type: ignore[valid-type] + neighbors: List[Any] + + +class MobilityArgs(BaseModel): + latitude: confloat(ge=-90, le=90) # type: ignore[valid-type] + longitude: confloat(ge=-180, le=180) # type: ignore[valid-type] + + +class TrackingArgs(BaseModel): + run_hash: str + + +class FederationArgs(BaseModel): + round: int + + +class ScenarioArgs(BaseModel): + federation: str + name: str + + +class NodesUpdateRequest(BaseModel): + device_args: DeviceArgs + network_args: NetworkArgs + mobility_args: MobilityArgs + tracking_args: TrackingArgs + federation_args: FederationArgs + scenario_args: ScenarioArgs + timestamp: str + + +def factory_requests_path(resource: str, user: str = "", role: str = "", scenario_name: str = "") -> str: + """Build paths for requests to the Database API from the Controller. + + This factory only maps DB API resources; controller endpoints do not require mapping here. + """ + if resource == "init": + return Routes.INIT + elif resource == "update": + return Routes.UPDATE + elif resource == "stop": + return Routes.STOP + elif resource == "remove": + return Routes.REMOVE + elif resource == "finish": + return Routes.FINISH + elif resource == "running": + return Routes.RUNNING + elif resource == "check_scenario": + return Routes.CHECK_SCENARIO.format(user=user, role=role, scenario_name=scenario_name) + elif resource == "get_scenarios_by_user": + return Routes.GET_SCENARIOS_BY_USER.format(user=user, role=role) + elif resource == "get_scenarios_by_scenario_name": + return Routes.GET_SCENARIOS_BY_SCENARIO_NAME.format(scenario_name=scenario_name) + # Nodes + elif resource == "get_nodes_by_scenario_name": + return Routes.NODES_BY_SCENARIO_NAME.format(scenario_name=scenario_name) + elif resource == "update_nodes": + return Routes.NODES_UPDATE + elif resource == "remove_nodes": + return Routes.NODES_REMOVE + # Notes + elif resource == "get_notes_by_scenario_name": + return Routes.NOTES_BY_SCENARIO_NAME.format(scenario_name=scenario_name) + elif resource == "update_notes": + return Routes.NOTES_UPDATE + elif resource == "remove_notes": + return Routes.NOTES_REMOVE + # Users + elif resource == "list_users": + return Routes.USER_LIST + elif resource == "get_user_by_scenario_name": + return Routes.USER_BY_SCENARIO_NAME.format(scenario_name=scenario_name) + elif resource == "add_user": + return Routes.USER_ADD + elif resource == "delete_user": + return Routes.USER_DELETE + elif resource == "update_user": + return Routes.USER_UPDATE + elif resource == "verify_user": + return Routes.USER_VERIFY + else: + raise Exception(f"resource not found: {resource}") diff --git a/nebula/database/adapters/postgress/postgress.py b/nebula/database/adapters/postgress/postgress.py index 65958ef7b..14423f94d 100755 --- a/nebula/database/adapters/postgress/postgress.py +++ b/nebula/database/adapters/postgress/postgress.py @@ -27,7 +27,7 @@ class PostgresDB(DatabaseAdapter): def __init__(self): self.pool = None - async def init_db_pool(self): + async def _init_db_pool(self): """ Initializes the asynchronous PostgreSQL connection pool. This should be called once when the application starts. @@ -62,7 +62,7 @@ async def init_db_pool(self): ) raise - async def close_db_pool(self): + async def _close_db_pool(self): """ Closes the asynchronous PostgreSQL connection pool. This should be called once when the application shuts down gracefully. @@ -71,10 +71,9 @@ async def close_db_pool(self): await self.pool.close() logging.info("Database connection pool closed.") - # --- User Management Functions --- - async def insert_default_admin(self): + async def _insert_default_admin(self): """ Inserts a default 'ADMIN' user into the database with a hashed password. The password must be provided via the ADMIN_PASSWORD environment variable. @@ -95,20 +94,22 @@ async def insert_default_admin(self): except Exception as e: logging.error(f"Failed to insert default admin user: {e}", exc_info=True) - async def list_users(self, all_info: bool = False): + async def _list_users(self, all_info: bool = False): """ Retrieves a list of users from the users database. """ async with self.pool.acquire() as conn: result = await conn.fetch("SELECT * FROM users") - if not all_info: - result = [user["user"] for user in result] - - return result + if all_info: + # Return JSON-serializable dicts with full info + return [dict(row) for row in result] + else: + # Return just the list of usernames (strings) + return [row["user"] for row in result] - async def get_user_info(self, user): + async def _get_user_info(self, user: str): """ Fetches detailed information for a specific user from the users database. """ @@ -116,23 +117,29 @@ async def get_user_info(self, user): return await conn.fetchrow('SELECT * FROM users WHERE "user" = $1', user) - async def verify(self, user, password): + async def _verify(self, user: str, password: str): """ - Verifies whether the provided password matches the stored hashed password for a user. + Verifies credentials and returns user info when valid. + + Returns + ------- + dict | None + {"user": USER, "role": ROLE} if valid, otherwise None. """ + user_up = user.upper() async with self.pool.acquire() as conn: - result = await conn.fetchrow('SELECT password FROM users WHERE "user" = $1', user) - if result: - try: - return pwd_context.verify(password, result[0]) - except Exception: - # Catch more general exceptions during verification to be safe - logging.error(f"Error during password verification for user {user}", exc_info=True) - return False - return False + row = await conn.fetchrow('SELECT password, role FROM users WHERE "user" = $1', user_up) + if not row: + return None + try: + if pwd_context.verify(password, row["password"]): + return {"user": user_up, "role": row["role"]} + except Exception: + logging.error(f"Error during password verification for user {user_up}", exc_info=True) + return None - async def verify_hash_algorithm(self, user): + async def _verify_hash_algorithm(self, user: str): """ Checks if the stored password hash for a user uses a supported Argon2 algorithm. """ @@ -146,7 +153,7 @@ async def verify_hash_algorithm(self, user): return False - async def delete_user_from_db(self, user): + async def _delete_user_from_db(self, user: str): """ Deletes a user record from the users database. """ @@ -154,7 +161,7 @@ async def delete_user_from_db(self, user): await conn.execute('DELETE FROM users WHERE "user" = $1', user) - async def add_user(self, user, password, role): + async def _add_user(self, user:str, password:str, role:str): """ Adds a new user to the users database with a hashed password. """ @@ -166,7 +173,7 @@ async def add_user(self, user, password, role): ) - async def update_user(self, user, password, role): + async def _update_user(self, user:str, password:str, role:str): """ Updates the password and role of an existing user in the users database. """ @@ -179,7 +186,7 @@ async def update_user(self, user, password, role): # --- Node Management Functions --- - async def list_nodes(self, scenario_name=None, sort_by="idx"): + async def _list_nodes(self, scenario_name:str=None, sort_by:str="idx"): """ Retrieves a list of nodes from the nodes database, optionally filtered by scenario and sorted. """ @@ -220,7 +227,7 @@ async def list_nodes(self, scenario_name=None, sort_by="idx"): return None - async def list_nodes_by_scenario_name(self, scenario_name): + async def _list_nodes_by_scenario_name(self, scenario_name:str): """ Fetches all nodes associated with a specific scenario, ordered by their index as integers. """ @@ -249,7 +256,7 @@ async def list_nodes_by_scenario_name(self, scenario_name): return None - async def update_node_record( + async def _update_node_record( self, node_uid, idx, @@ -284,6 +291,9 @@ async def update_node_record( logging.warning("Unable to serialize extras to JSON, storing as empty object.") extras_payload = json.dumps({}) + # Ensure malicious is stored as text if the column expects text + malicious_payload = malicious if isinstance(malicious, str) else str(malicious) + async with conn.transaction(): result = await conn.fetchrow( "SELECT * FROM nodes WHERE uid = $1 AND scenario = $2 FOR UPDATE;", @@ -300,7 +310,7 @@ async def update_node_record( $7, $8, $9, $10, $11, $12::jsonb, $13); """, node_uid, idx, ip, port, role, neighbors, - timestamp, federation, federation_round, scenario, run_hash, extras_payload, malicious, + timestamp, federation, federation_round, scenario, run_hash, extras_payload, malicious_payload, ) else: # Update existing node @@ -313,7 +323,7 @@ async def update_node_record( """, idx, ip, port, role, neighbors, timestamp, federation, federation_round, - run_hash, extras_payload, malicious, + run_hash, extras_payload, malicious_payload, node_uid, scenario, ) @@ -324,7 +334,7 @@ async def update_node_record( return None - async def remove_all_nodes(self): + async def _remove_all_nodes(self): """ Deletes all node records from the nodes database. """ @@ -332,7 +342,7 @@ async def remove_all_nodes(self): await conn.execute("TRUNCATE nodes CASCADE;") # Use CASCADE if there are foreign key dependencies - async def remove_nodes_by_scenario_name(self, scenario_name): + async def _remove_nodes_by_scenario_name(self, scenario_name:str): """ Deletes all nodes associated with a specific scenario from the database. """ @@ -341,7 +351,7 @@ async def remove_nodes_by_scenario_name(self, scenario_name): # --- Scenario Management Functions --- - async def get_all_scenarios(self, username, role, sort_by="start_time"): + async def _get_all_scenarios(self, username:str, role:str, sort_by:str="start_time"): """ Retrieves all scenarios from the database, accessing fields from the 'config' (JSONB) column and direct columns. Filters by user role and sorts by the specified field. @@ -391,7 +401,7 @@ async def get_all_scenarios(self, username, role, sort_by="start_time"): return await conn.fetch(full_command, *params) - async def get_all_scenarios_and_check_completed(self, username, role, sort_by="start_time"): + async def _get_all_scenarios_and_check_completed(self, user:str, role:str, sort_by:str="start_time"): """ Retrieves all scenarios, sorts them, and updates the status if necessary. Returns a list of dictionaries, where each dictionary represents a scenario. @@ -435,7 +445,7 @@ async def get_all_scenarios_and_check_completed(self, username, role, sort_by="s params = [] if role != "admin": command += " WHERE username = $1" # username is a direct column - params.append(username) + params.append(user) command += f" {order_by_clause};" @@ -446,30 +456,30 @@ async def get_all_scenarios_and_check_completed(self, username, role, sort_by="s re_fetch_required = False for scenario in scenarios_to_return: if scenario["status"] == "running": - if await self.check_scenario_federation_completed(scenario["name"]): - await self.scenario_set_status_to_completed(scenario["name"]) + if await self._check_scenario_federation_completed(scenario["name"]): + await self._scenario_set_status_to_completed(scenario["name"]) re_fetch_required = True break if re_fetch_required: # Recursively call to get fresh data after status update - return await self.get_all_scenarios_and_check_completed(username, role, sort_by) + return await self._get_all_scenarios_and_check_completed(user, role, sort_by) return scenarios_to_return - async def scenario_update_record(self, name, start_time, end_time, scenario_config, status, username): + async def _scenario_update_record(self, scenario_name:str, start_time:datetime, end_time:datetime, scenario:dict, status:str, username:str): """ Inserts or updates a scenario record using the PostgreSQL "UPSERT" pattern. All configuration is saved in the 'config' column of type JSONB. Direct columns (name, start_time, end_time, username, status) are also handled. """ - # Ensure scenario_config is a dictionary before dumping to JSON - if not isinstance(scenario_config, dict): + # Ensure scenario is a dictionary before dumping to JSON + if not isinstance(scenario, dict): try: - scenario_config = json.loads(scenario_config) + scenario = json.loads(scenario) except (json.JSONDecodeError, TypeError): - logging.error("scenario_config is not a valid JSON string or dict.") + logging.error("scenario is not a valid JSON string or dict.") return command = """ @@ -483,10 +493,10 @@ async def scenario_update_record(self, name, start_time, end_time, scenario_conf config = scenarios.config || EXCLUDED.config; -- Merge JSONB """ async with self.pool.acquire() as conn: - await conn.execute(command, name, start_time, end_time, username, status, json.dumps(scenario_config)) + await conn.execute(command, scenario_name, start_time, end_time, username, status, json.dumps(scenario)) - async def scenario_set_all_status_to_finished(self): + async def _scenario_set_all_status_to_finished(self): """ Sets the status of all 'running' scenarios to 'finished' and updates their 'end_time' (both in the direct column and within JSONB). @@ -505,7 +515,7 @@ async def scenario_set_all_status_to_finished(self): await conn.execute(command, current_time, json.dumps(current_time)) - async def scenario_set_status_to_finished(self, scenario_name): + async def _scenario_set_status_to_finished(self, scenario_name:str): """ Sets the status of a specific scenario to 'finished' and updates its 'end_time'. Updates both the direct columns and the JSONB 'config'. @@ -526,7 +536,7 @@ async def scenario_set_status_to_finished(self, scenario_name): await conn.execute(command, current_time, json.dumps(current_time), scenario_name) - async def scenario_set_status_to_completed(self, scenario_name): + async def _scenario_set_status_to_completed(self, scenario_name:str): """ Sets the status of a specific scenario to 'completed'. Updates both the direct column and the JSONB 'config'. @@ -542,7 +552,17 @@ async def scenario_set_status_to_completed(self, scenario_name): await conn.execute(command, scenario_name) - async def get_running_scenario(self, username=None, get_all=False): + async def _finish_scenario(self, scenario_name: str, all: bool = False): + """ + Consolidated method to set scenarios to finished. + """ + if all: + await self._scenario_set_all_status_to_finished() + else: + await self._scenario_set_status_to_finished(scenario_name) + + + async def _get_running_scenario(self, username:str=None, get_all:bool=False): """ Retrieves scenarios with a 'running' status, optionally filtered by user. Returns full scenario record (including direct columns and config JSONB). @@ -564,7 +584,7 @@ async def get_running_scenario(self, username=None, get_all=False): return result - async def get_completed_scenario(self): + async def _get_completed_scenario(self): """ Retrieves a single scenario with a 'completed' status. Returns full scenario record (including direct columns and config JSONB). @@ -574,8 +594,16 @@ async def get_completed_scenario(self): result_row = await conn.fetchrow(command, "completed") return dict(result_row) if result_row else None + async def _get_scenarios(self, user: str, role: str): + """ + Compose scenarios list and running scenario respecting role. + """ + scenarios = await self._get_all_scenarios_and_check_completed(username=user, role=role) + scenario_running = await self._get_running_scenario(None if role == "admin" else user) + return {"scenarios": scenarios, "scenario_running": scenario_running} + - async def get_scenario_by_name(self, scenario_name): + async def _get_scenario_by_name(self, scenario_name:str): """ Retrieves the complete record of a scenario by its name. """ @@ -601,7 +629,7 @@ async def get_scenario_by_name(self, scenario_name): return result - async def get_user_by_scenario_name(self, scenario_name): + async def _get_user_by_scenario_name(self, scenario_name:str): """ Retrieves the username associated with a scenario (from the direct 'username' column). """ @@ -609,7 +637,7 @@ async def get_user_by_scenario_name(self, scenario_name): return await conn.fetchval("SELECT username FROM scenarios WHERE name = $1;", scenario_name) - async def remove_scenario_by_name(self, scenario_name): + async def _remove_scenario_by_name(self, scenario_name:str): """ Delete a scenario from the database by its unique name. """ @@ -621,7 +649,7 @@ async def remove_scenario_by_name(self, scenario_name): logging.error(f"Error occurred while deleting scenario '{scenario_name}': {e}") - async def check_scenario_federation_completed(self, scenario_name): + async def _check_scenario_federation_completed(self, scenario_name:str): """ Check if all nodes in a given scenario have completed the required federation rounds. """ @@ -659,11 +687,11 @@ async def check_scenario_federation_completed(self, scenario_name): return False - async def check_scenario_with_role(self, role, scenario_name, current_username=None): + async def _check_scenario_with_role(self, role:str, scenario_name:str, user:str=None): """ Verify if a scenario exists that the user with the given role and username can access. """ - scenario_info = await self.get_scenario_by_name(scenario_name) + scenario_info = await self._get_scenario_by_name(scenario_name) if not scenario_info: return False # Scenario does not exist @@ -671,17 +699,17 @@ async def check_scenario_with_role(self, role, scenario_name, current_username=N if role == "admin": return True # Admins can access any existing scenario - if current_username is None: + if user is None: logging.warning( - "check_scenario_with_role called for non-admin role without current_username." + "check_scenario_with_role called for non-admin role without user." ) return False - return scenario_info.get("username") == current_username + return scenario_info.get("username") == user # --- Notes Management Functions --- - async def save_notes(self, scenario, notes): + async def _save_notes(self, scenario: str, notes: str): """ Save or update notes associated with a specific scenario. """ @@ -698,15 +726,19 @@ async def save_notes(self, scenario, notes): logging.error(f"PostgreSQL error during save_notes: {e}") - async def get_notes(self, scenario): + async def _get_notes(self, scenario: str): """ Retrieve notes associated with a specific scenario. """ async with self.pool.acquire() as conn: - return await conn.fetchrow("SELECT * FROM notes WHERE scenario = $1;", scenario) + row = await conn.fetchrow("SELECT * FROM notes WHERE scenario = $1;", scenario) + if row is None: + # No notes stored for this scenario yet + return None + return dict(row) - async def remove_note(self, scenario): + async def _remove_note(self, scenario: str): """ Delete the note associated with a specific scenario. """ diff --git a/nebula/database/database_adapter_interface.py b/nebula/database/database_adapter_interface.py index 9c0286ef1..250464fb1 100644 --- a/nebula/database/database_adapter_interface.py +++ b/nebula/database/database_adapter_interface.py @@ -8,71 +8,71 @@ class DatabaseAdapter(ABC): """ @abstractmethod - async def init_db_pool(self): + async def _init_db_pool(self): """Initializes the database connection pool.""" raise NotImplementedError @abstractmethod - async def close_db_pool(self): + async def _close_db_pool(self): """Closes the database connection pool.""" raise NotImplementedError # --- User Management Functions --- @abstractmethod - async def insert_default_admin(self): + async def _insert_default_admin(self): """Inserts a default admin user.""" raise NotImplementedError @abstractmethod - async def list_users(self, all_info=False): + async def _list_users(self, all_info=False): """Retrieves a list of users.""" raise NotImplementedError @abstractmethod - async def get_user_info(self, user): + async def _get_user_info(self, user): """Fetches detailed information for a specific user.""" raise NotImplementedError @abstractmethod - async def verify(self, user, password): + async def _verify(self, user, password): """Verifies user credentials.""" raise NotImplementedError @abstractmethod - async def verify_hash_algorithm(self, user): + async def _verify_hash_algorithm(self, user): """Checks the password hash algorithm for a user.""" raise NotImplementedError @abstractmethod - async def delete_user_from_db(self, user): + async def _delete_user_from_db(self, user): """Deletes a user from the database.""" raise NotImplementedError @abstractmethod - async def add_user(self, user, password, role): + async def _add_user(self, user, password, role): """Adds a new user.""" raise NotImplementedError @abstractmethod - async def update_user(self, user, password, role): + async def _update_user(self, user, password, role): """Updates an existing user.""" raise NotImplementedError # --- Node Management Functions --- @abstractmethod - async def list_nodes(self, scenario_name=None, sort_by="idx"): + async def _list_nodes(self, scenario_name=None, sort_by="idx"): """Retrieves a list of nodes.""" raise NotImplementedError @abstractmethod - async def list_nodes_by_scenario_name(self, scenario_name): + async def _list_nodes_by_scenario_name(self, scenario_name): """Fetches all nodes for a specific scenario.""" raise NotImplementedError @abstractmethod - async def update_node_record( + async def _update_node_record( self, node_uid, idx, @@ -92,95 +92,107 @@ async def update_node_record( raise NotImplementedError @abstractmethod - async def remove_all_nodes(self): + async def _remove_all_nodes(self): """Deletes all node records.""" raise NotImplementedError @abstractmethod - async def remove_nodes_by_scenario_name(self, scenario_name): + async def _remove_nodes_by_scenario_name(self, scenario_name): """Deletes all nodes for a specific scenario.""" raise NotImplementedError # --- Scenario Management Functions --- @abstractmethod - async def get_all_scenarios(self, username, role, sort_by="start_time"): + async def _get_all_scenarios(self, username, role, sort_by="start_time"): """Retrieves all scenarios.""" raise NotImplementedError @abstractmethod - async def get_all_scenarios_and_check_completed(self, username, role, sort_by="start_time"): + async def _get_all_scenarios_and_check_completed(self, username, role, sort_by="start_time"): """Retrieves all scenarios and checks for completion.""" raise NotImplementedError @abstractmethod - async def scenario_update_record(self, name, start_time, end_time, scenario_config, status, username): + async def _scenario_update_record(self, name, start_time, end_time, scenario_config, status, username): """Inserts or updates a scenario record.""" raise NotImplementedError @abstractmethod - async def scenario_set_all_status_to_finished(self): + async def _scenario_set_all_status_to_finished(self): """Sets the status of all running scenarios to 'finished'.""" raise NotImplementedError @abstractmethod - async def scenario_set_status_to_finished(self, scenario_name): + async def _scenario_set_status_to_finished(self, scenario_name): """Sets the status of a specific scenario to 'finished'.""" raise NotImplementedError @abstractmethod - async def scenario_set_status_to_completed(self, scenario_name): + async def _scenario_set_status_to_completed(self, scenario_name): """Sets the status of a specific scenario to 'completed'.""" raise NotImplementedError @abstractmethod - async def get_running_scenario(self, username=None, get_all=False): + async def _get_running_scenario(self, username=None, get_all=False): """Retrieves running scenarios.""" raise NotImplementedError @abstractmethod - async def get_completed_scenario(self): + async def _get_completed_scenario(self): """Retrieves a completed scenario.""" raise NotImplementedError @abstractmethod - async def get_scenario_by_name(self, scenario_name): + async def _get_scenario_by_name(self, scenario_name): """Retrieves a scenario by its name.""" raise NotImplementedError @abstractmethod - async def get_user_by_scenario_name(self, scenario_name): + async def _get_user_by_scenario_name(self, scenario_name): """Retrieves the user associated with a scenario.""" raise NotImplementedError @abstractmethod - async def remove_scenario_by_name(self, scenario_name): + async def _remove_scenario_by_name(self, scenario_name): """Deletes a scenario by its name.""" raise NotImplementedError @abstractmethod - async def check_scenario_federation_completed(self, scenario_name): + async def _check_scenario_federation_completed(self, scenario_name): """Checks if a scenario's federation is complete.""" raise NotImplementedError @abstractmethod - async def check_scenario_with_role(self, role, scenario_name, current_username=None): + async def _check_scenario_with_role(self, role, scenario_name, current_username=None): """Verifies if a user can access a scenario.""" raise NotImplementedError # --- Notes Management Functions --- @abstractmethod - async def save_notes(self, scenario, notes): + async def _save_notes(self, scenario, notes): """Saves or updates notes for a scenario.""" raise NotImplementedError @abstractmethod - async def get_notes(self, scenario): + async def _get_notes(self, scenario): """Retrieves notes for a scenario.""" raise NotImplementedError @abstractmethod - async def remove_note(self, scenario): + async def _remove_note(self, scenario): """Deletes the note for a scenario.""" raise NotImplementedError + + # --- Scenario Finish (no API logic) --- + + @abstractmethod + async def _finish_scenario(self, scenario_name, all: bool = False): + """Sets status to finished for one scenario or all running scenarios.""" + raise NotImplementedError + + @abstractmethod + async def _get_scenarios(self, user: str, role: str): + """Return scenarios list and running scenario, given user and role.""" + raise NotImplementedError diff --git a/nebula/database/database_api.py b/nebula/database/database_api.py index 02df20898..f716117ae 100644 --- a/nebula/database/database_api.py +++ b/nebula/database/database_api.py @@ -1,17 +1,37 @@ -import argparse import logging import os import sys -from typing import Annotated sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../'))) -from fastapi import Body, FastAPI, HTTPException, Path, status +from fastapi import FastAPI, HTTPException, status, Depends from fastapi.concurrency import asynccontextmanager from nebula.database.database_adapter_factory import factory_database_adapter -from nebula.database.utils_requests import Routes +from nebula.database.utils_requests import ( + Routes, + ScenarioUpdateRequest, + ScenarioStopRequest, + ScenarioRemoveRequest, + ScenarioFinishRequest, + NotesUpdateRequest, + NotesRemoveRequest, + NodesRemoveRequest, + UserAddRequest, + UserDeleteRequest, + UserUpdateRequest, + UserVerifyRequest, + NodesUpdateRequest, + GetScenariosRequest, + GetRunningScenarioRequest, + CheckScenarioRequest, + GetScenarioByNameRequest, + ListNodesByScenarioNameRequest, + GetNotesByScenarioNameRequest, + ListUsersRequest, + GetUserByScenarioNameRequest, +) # Get a database instance db = factory_database_adapter("PostgresDB") @@ -56,13 +76,13 @@ async def lifespan(app: FastAPI): configure_logger(db_log) # Initialize the database connection pool - await db.init_db_pool() - await db.insert_default_admin() + await db._init_db_pool() + await db._insert_default_admin() yield # Code to run on shutdown - await db.close_db_pool() + await db._close_db_pool() app = FastAPI(lifespan=lifespan) @@ -76,61 +96,54 @@ async def read_root(): # Scenarios @app.post(Routes.UPDATE) async def update_scenario( - scenario_name: str = Body(..., embed=True), - start_time: str = Body(..., embed=True), - end_time: str = Body(..., embed=True), - scenario: dict = Body(..., embed=True), - status: str = Body(..., embed=True), - username: str = Body(..., embed=True), + payload: ScenarioUpdateRequest, ): try: - await db.scenario_update_record(scenario_name, start_time, end_time, scenario, status, username) - return {"message": f"Scenario {scenario_name} updated successfully"} + await db._scenario_update_record( + **payload.model_dump() + ) + return {"message": f"Scenario {payload.scenario_name} updated successfully"} except Exception as e: - logging.exception(f"Error updating scenario {scenario_name}: {e}") + logging.exception( + f"Error updating scenario {payload.scenario_name}: {e}" + ) raise HTTPException(status_code=500, detail="Internal server error") @app.post(Routes.STOP) async def stop_scenario( - scenario_name: str = Body(..., embed=True), - all: bool = Body(False, embed=True), + payload: ScenarioStopRequest, ): try: - if all: - await db.scenario_set_all_status_to_finished() - else: - await db.scenario_set_status_to_finished(scenario_name) - return {"message": "Scenario status updated successfully"} + await db._finish_scenario(payload.scenario_name, payload.all) + return {"message": "Finished status set successfully"} except Exception as e: - logging.exception(f"Error stopping scenario {scenario_name}: {e}") + logging.exception( + f"Error stopping scenario {payload.scenario_name}: {e}" + ) raise HTTPException(status_code=500, detail="Internal server error") @app.post(Routes.REMOVE) async def remove_scenario( - scenario_name: str = Body(..., embed=True), + payload: ScenarioRemoveRequest, ): try: - await db.remove_scenario_by_name(scenario_name) - return {"message": f"Scenario {scenario_name} removed successfully"} + await db._remove_scenario_by_name(payload.scenario_name) + return {"message": f"Scenario {payload.scenario_name} removed successfully"} except Exception as e: - logging.exception(f"Error removing scenario {scenario_name}: {e}") + logging.exception( + f"Error removing scenario {payload.scenario_name}: {e}" + ) raise HTTPException(status_code=500, detail="Internal server error") @app.get(Routes.GET_SCENARIOS_BY_USER) async def get_scenarios( - user: Annotated[str, Path(pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50)], - role: Annotated[str, Path(pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50)], + payload: GetScenariosRequest = Depends() ): try: - scenarios = await db.get_all_scenarios_and_check_completed(username=user, role=role) - if role == "admin": - scenario_running = await db.get_running_scenario() - else: - scenario_running = await db.get_running_scenario(username=user) - return {"scenarios": scenarios, "scenario_running": scenario_running} + return await db._get_scenarios(payload.user, payload.role) except Exception as e: logging.exception(f"Error obtaining scenarios: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -138,23 +151,24 @@ async def get_scenarios( @app.post(Routes.FINISH) async def set_scenario_status_to_finished( - scenario_name: str = Body(..., embed=True), all: bool = Body(False, embed=True) + payload: ScenarioFinishRequest, ): try: - if all: - await db.scenario_set_all_status_to_finished() - else: - await db.scenario_set_status_to_finished(scenario_name) - return {"message": f"Scenario {scenario_name} status set to finished successfully"} + await db._finish_scenario( + payload.scenario_name, payload.all + ) + return {"message": "Finished status set successfully"} except Exception as e: - logging.exception(f"Error setting scenario {scenario_name} to finished: {e}") + logging.exception( + f"Error setting scenario {payload.scenario_name} to finished: {e}" + ) raise HTTPException(status_code=500, detail="Internal server error") @app.get(Routes.RUNNING) -async def get_running_scenario_endpoint(get_all: bool = False): +async def get_running_scenario_endpoint(payload: GetRunningScenarioRequest = Depends()): try: - return await db.get_running_scenario(get_all=get_all) + return await db._get_running_scenario(get_all=payload.get_all) except Exception as e: logging.exception(f"Error obtaining running scenario: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -162,11 +176,11 @@ async def get_running_scenario_endpoint(get_all: bool = False): @app.get(Routes.CHECK_SCENARIO) async def check_scenario( - role: Annotated[str, Path(pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50)], - scenario_name: Annotated[str, Path(pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50)], + payload: CheckScenarioRequest = Depends() ): try: - allowed = await db.check_scenario_with_role(role, scenario_name) + params = CheckScenarioRequest(**payload.model_dump()) + allowed = await db._check_scenario_with_role(params.role, params.scenario_name, params.user) return {"allowed": allowed} except Exception as e: logging.exception(f"Error checking scenario with role: {e}") @@ -175,23 +189,23 @@ async def check_scenario( @app.get(Routes.GET_SCENARIOS_BY_SCENARIO_NAME) async def get_scenario_by_name_endpoint( - scenario_name: Annotated[str, Path(pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50)], + payload: GetScenarioByNameRequest = Depends(), ): try: - scenario = await db.get_scenario_by_name(scenario_name) + scenario = await db._get_scenario_by_name(payload.scenario_name) return scenario except Exception as e: - logging.exception(f"Error obtaining scenario {scenario_name}: {e}") + logging.exception(f"Error obtaining scenario {payload.scenario_name}: {e}") raise HTTPException(status_code=500, detail="Internal server error") # Nodes @app.get(Routes.NODES_BY_SCENARIO_NAME) async def list_nodes_by_scenario_name_endpoint( - scenario_name: Annotated[str, Path(pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50)], + payload: ListNodesByScenarioNameRequest = Depends() ): try: - nodes = await db.list_nodes_by_scenario_name(scenario_name) + nodes = await db._list_nodes_by_scenario_name(payload.scenario_name) return nodes except Exception as e: logging.exception(f"Error obtaining nodes: {e}") @@ -199,25 +213,27 @@ async def list_nodes_by_scenario_name_endpoint( @app.post(Routes.NODES_UPDATE) -async def update_node_record(data: dict): +async def update_node_record(payload: NodesUpdateRequest): try: - # Store latitude/longitude (and any mobility details) inside extras JSONB - # Prefer explicit `extras` sent by controller; fallback to `mobility_args` for backward compatibility - extras = data.get("extras") or data.get("mobility_args", {}) - await db.update_node_record( - str(data["device_args"]["uid"]), - str(data["device_args"]["idx"]), - str(data["network_args"]["ip"]), - str(data["network_args"]["port"]), - str(data["device_args"]["role"]), - data["network_args"]["neighbors"], + # Build extras from mobility_args + extras = { + "latitude": payload.mobility_args.latitude, + "longitude": payload.mobility_args.longitude, + } + await db._update_node_record( + str(payload.device_args.uid), + str(payload.device_args.idx), + str(payload.network_args.ip), + str(payload.network_args.port), + str(payload.device_args.role), + payload.network_args.neighbors, extras, - str(data["timestamp"]), - str(data["scenario_args"]["federation"]), - str(data["federation_args"]["round"]), - str(data["scenario_args"]["name"]), - str(data["tracking_args"]["run_hash"]), - str(data["device_args"]["malicious"]), + str(payload.timestamp), + str(payload.scenario_args.federation), + str(payload.federation_args.round), + str(payload.scenario_args.name), + str(payload.tracking_args.run_hash), + bool(payload.device_args.malicious), ) return {"message": "Node updated successfully"} except Exception as e: @@ -226,10 +242,10 @@ async def update_node_record(data: dict): @app.post(Routes.NODES_REMOVE) -async def remove_nodes_by_scenario_name_endpoint(scenario_name: str = Body(..., embed=True)): +async def remove_nodes_by_scenario_name_endpoint(payload: NodesRemoveRequest): try: - await db.remove_nodes_by_scenario_name(scenario_name) - return {"message": f"Nodes for scenario {scenario_name} removed successfully"} + await db._remove_nodes_by_scenario_name(payload.scenario_name) + return {"message": f"Nodes for scenario {payload.scenario_name} removed successfully"} except Exception as e: logging.exception(f"Error removing nodes: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -238,33 +254,31 @@ async def remove_nodes_by_scenario_name_endpoint(scenario_name: str = Body(..., # Notes @app.get(Routes.NOTES_BY_SCENARIO_NAME) async def get_notes_by_scenario_name( - scenario_name: Annotated[str, Path(pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50)], + payload: GetNotesByScenarioNameRequest = Depends() ): try: - notes_record = await db.get_notes(scenario_name) - if notes_record is not None: - notes_record = dict(notes_record.items()) + notes_record = await db._get_notes(payload.scenario_name) return notes_record except Exception as e: - logging.exception(f"Error obtaining notes for scenario {scenario_name}: {e}") + logging.exception(f"Error obtaining notes for scenario {payload.scenario_name}: {e}") raise HTTPException(status_code=500, detail="Internal server error") @app.post(Routes.NOTES_UPDATE) -async def update_notes_by_scenario_name(scenario_name: str = Body(..., embed=True), notes: str = Body(..., embed=True)): +async def update_notes_by_scenario_name(payload: NotesUpdateRequest): try: - await db.save_notes(scenario_name, notes) - return {"message": f"Notes for scenario {scenario_name} updated successfully"} + await db._save_notes(payload.scenario_name, payload.notes) + return {"message": f"Notes for scenario {payload.scenario_name} updated successfully"} except Exception as e: logging.exception(f"Error updating notes: {e}") raise HTTPException(status_code=500, detail="Internal server error") @app.post(Routes.NOTES_REMOVE) -async def remove_notes_by_scenario_name_endpoint(scenario_name: str = Body(..., embed=True)): +async def remove_notes_by_scenario_name_endpoint(payload: NotesRemoveRequest): try: - await db.remove_note(scenario_name) - return {"message": f"Notes for scenario {scenario_name} removed successfully"} + await db._remove_note(payload.scenario_name) + return {"message": f"Notes for scenario {payload.scenario_name} removed successfully"} except Exception as e: logging.exception(f"Error removing notes: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -272,12 +286,9 @@ async def remove_notes_by_scenario_name_endpoint(scenario_name: str = Body(..., # Users @app.get(Routes.USER_LIST) -async def list_users_controller(all_info: bool = False): +async def list_users_controller(payload: ListUsersRequest = Depends()): try: - user_list = await db.list_users(all_info) - if all_info: - user_list = [dict(user) for user in user_list] - return {"users": user_list} + return {"users": await db._list_users(payload.all_info)} except Exception as e: logging.exception(f"Error retrieving users: {e}") raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error retrieving users: {e}") @@ -285,20 +296,20 @@ async def list_users_controller(all_info: bool = False): @app.get(Routes.USER_BY_SCENARIO_NAME) async def get_user_by_scenario_name_endpoint( - scenario_name: Annotated[str, Path(pattern="^[a-zA-Z0-9_-]+$", min_length=1, max_length=50)], + payload: GetUserByScenarioNameRequest = Depends() ): try: - user = await db.get_user_by_scenario_name(scenario_name) + user = await db._get_user_by_scenario_name(payload.scenario_name) return user except Exception as e: - logging.exception(f"Error obtaining user for scenario {scenario_name}: {e}") + logging.exception(f"Error obtaining user for scenario {payload.scenario_name}: {e}") raise HTTPException(status_code=500, detail="Internal server error") @app.post(Routes.USER_ADD) -async def add_user_controller(user: str = Body(...), password: str = Body(...), role: str = Body(...)): +async def add_user_controller(payload: UserAddRequest): try: - await db.add_user(user, password, role) + await db._add_user(payload.user, payload.password, payload.role) return {"detail": "User added successfully"} except Exception as e: logging.exception(f"Error adding user: {e}") @@ -306,9 +317,9 @@ async def add_user_controller(user: str = Body(...), password: str = Body(...), @app.post(Routes.USER_DELETE) -async def remove_user_controller(user: str = Body(..., embed=True)): +async def remove_user_controller(payload: UserDeleteRequest): try: - await db.delete_user_from_db(user) + await db._delete_user_from_db(payload.user) return {"detail": "User deleted successfully"} except Exception as e: logging.exception(f"Error deleting user: {e}") @@ -316,9 +327,9 @@ async def remove_user_controller(user: str = Body(..., embed=True)): @app.post(Routes.USER_UPDATE) -async def update_user_controller(user: str = Body(...), password: str = Body(...), role: str = Body(...)): +async def update_user_controller(payload: UserUpdateRequest): try: - await db.update_user(user, password, role) + await db._update_user(payload.user, payload.password, payload.role) return {"detail": "User updated successfully"} except Exception as e: logging.exception(f"Error updating user: {e}") @@ -326,15 +337,12 @@ async def update_user_controller(user: str = Body(...), password: str = Body(... @app.post(Routes.USER_VERIFY) -async def verify_user_controller(user: str = Body(...), password: str = Body(...)): +async def verify_user_controller(payload: UserVerifyRequest): try: - user_submitted = user.upper() - users = await db.list_users() - if users and await db.verify(user_submitted, password): - user_info = await db.get_user_info(user_submitted) - return {"user": user_submitted, "role": user_info[2]} - else: - raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED) + auth = await db._verify(payload.user, payload.password) + if auth: + return auth + raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED) except Exception as e: logging.exception(f"Error verifying user: {e}") raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error verifying user: {e}") diff --git a/nebula/database/utils_requests.py b/nebula/database/utils_requests.py index 25060322f..8d9f88da2 100644 --- a/nebula/database/utils_requests.py +++ b/nebula/database/utils_requests.py @@ -1,24 +1,17 @@ from typing import Any, Dict, List -from pydantic import BaseModel, conint, confloat +from pydantic import BaseModel, confloat, conint class Routes: # Scenarios INIT = "/" - STATUS = "/status" - RESOURCES = "/resources" - LEAST_MEMORY_GPU = "/least_memory_gpu" - AVAILABLE_GPUS = "/available_gpus/" - DISCOVER_VPN = "/discover-vpn" - - RUN = "/scenarios/run" UPDATE = "/scenarios/update" STOP = "/scenarios/stop" REMOVE = "/scenarios/remove" FINISH = "/scenarios/set_status_to_finished" RUNNING = "/scenarios/running" - CHECK_SCENARIO = "/scenarios/check/{role}/{scenario_name}" + CHECK_SCENARIO = "/scenarios/check/{user}/{role}/{scenario_name}" GET_SCENARIOS_BY_USER = "/scenarios/{user}/{role}" GET_SCENARIOS_BY_SCENARIO_NAME = "/scenarios/{scenario_name}" @@ -26,8 +19,6 @@ class Routes: NODES_BY_SCENARIO_NAME = "/nodes/{scenario_name}" NODES_UPDATE = "/nodes/update" NODES_REMOVE = "/nodes/remove" - NODES_UPDATE_BY_SCENARIO = "/nodes/{scenario_name}/update" - NODES_DONE_BY_SCENARIO = "/nodes/{scenario_name}/done" # Notes NOTES_BY_SCENARIO_NAME = "/notes/{scenario_name}" @@ -42,23 +33,6 @@ class Routes: USER_UPDATE = "/user/update" USER_VERIFY = "/user/verify" - # Physical - PHYSICAL_RUN = "/physical/run/{ip}" - PHYSICAL_STOP = "/physical/stop/{ip}" - PHYSICAL_SETUP = "/physical/setup/{ip}" - PHYSICAL_STATE = "/physical/state/{ip}" - PHYSICAL_SCENARIO_STATE = "/physical/scenario-state/{scenario_name}" - - -class RunScenarioRequest(BaseModel): - """Request model to trigger a scenario run on the controller. - - - Only requires scenario_data and user. - - Extra fields (e.g., role, federation_id) are ignored. - """ - scenario_data: Dict[str, Any] - user: str - class ScenarioUpdateRequest(BaseModel): scenario_name: str @@ -158,6 +132,40 @@ class NodesUpdateRequest(BaseModel): scenario_args: ScenarioArgs timestamp: str +class GetScenariosRequest(BaseModel): + user: str + role: str + + +class GetRunningScenarioRequest(BaseModel): + get_all: bool = False + + +class CheckScenarioRequest(BaseModel): + user: str + role: str + scenario_name: str + + +class GetScenarioByNameRequest(BaseModel): + scenario_name: str + + +class ListNodesByScenarioNameRequest(BaseModel): + scenario_name: str + + +class GetNotesByScenarioNameRequest(BaseModel): + scenario_name: str + + +class ListUsersRequest(BaseModel): + all_info: bool = False + + +class GetUserByScenarioNameRequest(BaseModel): + scenario_name: str + def factory_requests_path(resource: str, user: str = "", role: str = "", scenario_name: str = "") -> str: if resource == "init": @@ -173,7 +181,7 @@ def factory_requests_path(resource: str, user: str = "", role: str = "", scenari elif resource == "running": return Routes.RUNNING elif resource == "check_scenario": - return Routes.CHECK_SCENARIO.format(role=role, scenario_name=scenario_name) + return Routes.CHECK_SCENARIO.format(user=user, role=role, scenario_name=scenario_name) elif resource == "get_scenarios_by_user": return Routes.GET_SCENARIOS_BY_USER.format(user=user, role=role) elif resource == "get_scenarios_by_scenario_name": From cc13db465f24ee27dbdbfc5cebc6adcbededb56e Mon Sep 17 00:00:00 2001 From: FerTV Date: Fri, 19 Sep 2025 13:25:54 +0200 Subject: [PATCH 08/14] feature: added extras column for addons in the node table --- nebula/frontend/app.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/nebula/frontend/app.py b/nebula/frontend/app.py index fb065e308..59c836a1b 100755 --- a/nebula/frontend/app.py +++ b/nebula/frontend/app.py @@ -1593,6 +1593,7 @@ async def nebula_dashboard_monitor(scenario_name: str, request: Request, session # Calculate initial status based on timestamp timestamp = datetime.datetime.strptime(node["timestamp"], "%Y-%m-%d %H:%M:%S.%f") is_online = (datetime.datetime.now() - timestamp) <= datetime.timedelta(seconds=25) + mobility_args = json.loads(node["extras"]) formatted_nodes.append({ "uid": node["uid"], @@ -1601,8 +1602,8 @@ async def nebula_dashboard_monitor(scenario_name: str, request: Request, session "port": node["port"], "role": node["role"], "neighbors": " ".join(node["neighbors"]), - "latitude": node["latitude"], - "longitude": node["longitude"], + "latitude": mobility_args["latitude"], + "longitude": mobility_args["longitude"], "timestamp": node["timestamp"], "federation": node["federation"], "round": str(node["round"]), From 56c4d1a96509647678771f6f823bd516de218d24 Mon Sep 17 00:00:00 2001 From: FerTV Date: Fri, 19 Sep 2025 13:37:44 +0200 Subject: [PATCH 09/14] fix: dashboard and user issues --- nebula/database/adapters/postgress/postgress.py | 2 +- nebula/database/database_api.py | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/nebula/database/adapters/postgress/postgress.py b/nebula/database/adapters/postgress/postgress.py index 14423f94d..7d2c1a37e 100755 --- a/nebula/database/adapters/postgress/postgress.py +++ b/nebula/database/adapters/postgress/postgress.py @@ -598,7 +598,7 @@ async def _get_scenarios(self, user: str, role: str): """ Compose scenarios list and running scenario respecting role. """ - scenarios = await self._get_all_scenarios_and_check_completed(username=user, role=role) + scenarios = await self._get_all_scenarios_and_check_completed(user=user, role=role) scenario_running = await self._get_running_scenario(None if role == "admin" else user) return {"scenarios": scenarios, "scenario_running": scenario_running} diff --git a/nebula/database/database_api.py b/nebula/database/database_api.py index f716117ae..fba278a15 100644 --- a/nebula/database/database_api.py +++ b/nebula/database/database_api.py @@ -343,6 +343,9 @@ async def verify_user_controller(payload: UserVerifyRequest): if auth: return auth raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED) + except HTTPException as e: + # Propagate intended HTTP errors (e.g., 401) without wrapping + raise e except Exception as e: logging.exception(f"Error verifying user: {e}") raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error verifying user: {e}") From cd518ffe442d6fc5eb56391ba8548cbe7c36b6bb Mon Sep 17 00:00:00 2001 From: FerTV Date: Fri, 19 Sep 2025 15:47:22 +0200 Subject: [PATCH 10/14] refactor: payload for request --- nebula/database/database_api.py | 131 ++++++++++++++++---------------- 1 file changed, 65 insertions(+), 66 deletions(-) diff --git a/nebula/database/database_api.py b/nebula/database/database_api.py index fba278a15..579a11a05 100644 --- a/nebula/database/database_api.py +++ b/nebula/database/database_api.py @@ -96,54 +96,54 @@ async def read_root(): # Scenarios @app.post(Routes.UPDATE) async def update_scenario( - payload: ScenarioUpdateRequest, + request: ScenarioUpdateRequest, ): try: await db._scenario_update_record( - **payload.model_dump() + **request.model_dump() ) - return {"message": f"Scenario {payload.scenario_name} updated successfully"} + return {"message": f"Scenario {request.scenario_name} updated successfully"} except Exception as e: logging.exception( - f"Error updating scenario {payload.scenario_name}: {e}" + f"Error updating scenario {request.scenario_name}: {e}" ) raise HTTPException(status_code=500, detail="Internal server error") @app.post(Routes.STOP) async def stop_scenario( - payload: ScenarioStopRequest, + request: ScenarioStopRequest, ): try: - await db._finish_scenario(payload.scenario_name, payload.all) + await db._finish_scenario(request.scenario_name, request.all) return {"message": "Finished status set successfully"} except Exception as e: logging.exception( - f"Error stopping scenario {payload.scenario_name}: {e}" + f"Error stopping scenario {request.scenario_name}: {e}" ) raise HTTPException(status_code=500, detail="Internal server error") @app.post(Routes.REMOVE) async def remove_scenario( - payload: ScenarioRemoveRequest, + request: ScenarioRemoveRequest, ): try: - await db._remove_scenario_by_name(payload.scenario_name) - return {"message": f"Scenario {payload.scenario_name} removed successfully"} + await db._remove_scenario_by_name(request.scenario_name) + return {"message": f"Scenario {request.scenario_name} removed successfully"} except Exception as e: logging.exception( - f"Error removing scenario {payload.scenario_name}: {e}" + f"Error removing scenario {request.scenario_name}: {e}" ) raise HTTPException(status_code=500, detail="Internal server error") @app.get(Routes.GET_SCENARIOS_BY_USER) async def get_scenarios( - payload: GetScenariosRequest = Depends() + request: GetScenariosRequest = Depends() ): try: - return await db._get_scenarios(payload.user, payload.role) + return await db._get_scenarios(request.user, request.role) except Exception as e: logging.exception(f"Error obtaining scenarios: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -151,24 +151,24 @@ async def get_scenarios( @app.post(Routes.FINISH) async def set_scenario_status_to_finished( - payload: ScenarioFinishRequest, + request: ScenarioFinishRequest, ): try: await db._finish_scenario( - payload.scenario_name, payload.all + request.scenario_name, request.all ) return {"message": "Finished status set successfully"} except Exception as e: logging.exception( - f"Error setting scenario {payload.scenario_name} to finished: {e}" + f"Error setting scenario {request.scenario_name} to finished: {e}" ) raise HTTPException(status_code=500, detail="Internal server error") @app.get(Routes.RUNNING) -async def get_running_scenario_endpoint(payload: GetRunningScenarioRequest = Depends()): +async def get_running_scenario_endpoint(request: GetRunningScenarioRequest = Depends()): try: - return await db._get_running_scenario(get_all=payload.get_all) + return await db._get_running_scenario(get_all=request.get_all) except Exception as e: logging.exception(f"Error obtaining running scenario: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -176,11 +176,10 @@ async def get_running_scenario_endpoint(payload: GetRunningScenarioRequest = Dep @app.get(Routes.CHECK_SCENARIO) async def check_scenario( - payload: CheckScenarioRequest = Depends() + request: CheckScenarioRequest = Depends() ): try: - params = CheckScenarioRequest(**payload.model_dump()) - allowed = await db._check_scenario_with_role(params.role, params.scenario_name, params.user) + allowed = await db._check_scenario_with_role(**request.model_dump()) return {"allowed": allowed} except Exception as e: logging.exception(f"Error checking scenario with role: {e}") @@ -189,23 +188,23 @@ async def check_scenario( @app.get(Routes.GET_SCENARIOS_BY_SCENARIO_NAME) async def get_scenario_by_name_endpoint( - payload: GetScenarioByNameRequest = Depends(), + request: GetScenarioByNameRequest = Depends(), ): try: - scenario = await db._get_scenario_by_name(payload.scenario_name) + scenario = await db._get_scenario_by_name(request.scenario_name) return scenario except Exception as e: - logging.exception(f"Error obtaining scenario {payload.scenario_name}: {e}") + logging.exception(f"Error obtaining scenario {request.scenario_name}: {e}") raise HTTPException(status_code=500, detail="Internal server error") # Nodes @app.get(Routes.NODES_BY_SCENARIO_NAME) async def list_nodes_by_scenario_name_endpoint( - payload: ListNodesByScenarioNameRequest = Depends() + request: ListNodesByScenarioNameRequest = Depends() ): try: - nodes = await db._list_nodes_by_scenario_name(payload.scenario_name) + nodes = await db._list_nodes_by_scenario_name(request.scenario_name) return nodes except Exception as e: logging.exception(f"Error obtaining nodes: {e}") @@ -213,27 +212,27 @@ async def list_nodes_by_scenario_name_endpoint( @app.post(Routes.NODES_UPDATE) -async def update_node_record(payload: NodesUpdateRequest): +async def update_node_record(request: NodesUpdateRequest): try: # Build extras from mobility_args extras = { - "latitude": payload.mobility_args.latitude, - "longitude": payload.mobility_args.longitude, + "latitude": request.mobility_args.latitude, + "longitude": request.mobility_args.longitude, } await db._update_node_record( - str(payload.device_args.uid), - str(payload.device_args.idx), - str(payload.network_args.ip), - str(payload.network_args.port), - str(payload.device_args.role), - payload.network_args.neighbors, + str(request.device_args.uid), + str(request.device_args.idx), + str(request.network_args.ip), + str(request.network_args.port), + str(request.device_args.role), + request.network_args.neighbors, extras, - str(payload.timestamp), - str(payload.scenario_args.federation), - str(payload.federation_args.round), - str(payload.scenario_args.name), - str(payload.tracking_args.run_hash), - bool(payload.device_args.malicious), + str(request.timestamp), + str(request.scenario_args.federation), + str(request.federation_args.round), + str(request.scenario_args.name), + str(request.tracking_args.run_hash), + bool(request.device_args.malicious), ) return {"message": "Node updated successfully"} except Exception as e: @@ -242,10 +241,10 @@ async def update_node_record(payload: NodesUpdateRequest): @app.post(Routes.NODES_REMOVE) -async def remove_nodes_by_scenario_name_endpoint(payload: NodesRemoveRequest): +async def remove_nodes_by_scenario_name_endpoint(request: NodesRemoveRequest): try: - await db._remove_nodes_by_scenario_name(payload.scenario_name) - return {"message": f"Nodes for scenario {payload.scenario_name} removed successfully"} + await db._remove_nodes_by_scenario_name(request.scenario_name) + return {"message": f"Nodes for scenario {request.scenario_name} removed successfully"} except Exception as e: logging.exception(f"Error removing nodes: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -254,31 +253,31 @@ async def remove_nodes_by_scenario_name_endpoint(payload: NodesRemoveRequest): # Notes @app.get(Routes.NOTES_BY_SCENARIO_NAME) async def get_notes_by_scenario_name( - payload: GetNotesByScenarioNameRequest = Depends() + request: GetNotesByScenarioNameRequest = Depends() ): try: - notes_record = await db._get_notes(payload.scenario_name) + notes_record = await db._get_notes(request.scenario_name) return notes_record except Exception as e: - logging.exception(f"Error obtaining notes for scenario {payload.scenario_name}: {e}") + logging.exception(f"Error obtaining notes for scenario {request.scenario_name}: {e}") raise HTTPException(status_code=500, detail="Internal server error") @app.post(Routes.NOTES_UPDATE) -async def update_notes_by_scenario_name(payload: NotesUpdateRequest): +async def update_notes_by_scenario_name(request: NotesUpdateRequest): try: - await db._save_notes(payload.scenario_name, payload.notes) - return {"message": f"Notes for scenario {payload.scenario_name} updated successfully"} + await db._save_notes(**request.model_dump()) + return {"message": f"Notes for scenario {request.scenario_name} updated successfully"} except Exception as e: logging.exception(f"Error updating notes: {e}") raise HTTPException(status_code=500, detail="Internal server error") @app.post(Routes.NOTES_REMOVE) -async def remove_notes_by_scenario_name_endpoint(payload: NotesRemoveRequest): +async def remove_notes_by_scenario_name_endpoint(request: NotesRemoveRequest): try: - await db._remove_note(payload.scenario_name) - return {"message": f"Notes for scenario {payload.scenario_name} removed successfully"} + await db._remove_note(request.scenario_name) + return {"message": f"Notes for scenario {request.scenario_name} removed successfully"} except Exception as e: logging.exception(f"Error removing notes: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -286,9 +285,9 @@ async def remove_notes_by_scenario_name_endpoint(payload: NotesRemoveRequest): # Users @app.get(Routes.USER_LIST) -async def list_users_controller(payload: ListUsersRequest = Depends()): +async def list_users_controller(request: ListUsersRequest = Depends()): try: - return {"users": await db._list_users(payload.all_info)} + return {"users": await db._list_users(request.all_info)} except Exception as e: logging.exception(f"Error retrieving users: {e}") raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error retrieving users: {e}") @@ -296,20 +295,20 @@ async def list_users_controller(payload: ListUsersRequest = Depends()): @app.get(Routes.USER_BY_SCENARIO_NAME) async def get_user_by_scenario_name_endpoint( - payload: GetUserByScenarioNameRequest = Depends() + request: GetUserByScenarioNameRequest = Depends() ): try: - user = await db._get_user_by_scenario_name(payload.scenario_name) + user = await db._get_user_by_scenario_name(request.scenario_name) return user except Exception as e: - logging.exception(f"Error obtaining user for scenario {payload.scenario_name}: {e}") + logging.exception(f"Error obtaining user for scenario {request.scenario_name}: {e}") raise HTTPException(status_code=500, detail="Internal server error") @app.post(Routes.USER_ADD) -async def add_user_controller(payload: UserAddRequest): +async def add_user_controller(request: UserAddRequest): try: - await db._add_user(payload.user, payload.password, payload.role) + await db._add_user(**request.model_dump()) return {"detail": "User added successfully"} except Exception as e: logging.exception(f"Error adding user: {e}") @@ -317,9 +316,9 @@ async def add_user_controller(payload: UserAddRequest): @app.post(Routes.USER_DELETE) -async def remove_user_controller(payload: UserDeleteRequest): +async def remove_user_controller(request: UserDeleteRequest): try: - await db._delete_user_from_db(payload.user) + await db._delete_user_from_db(request.user) return {"detail": "User deleted successfully"} except Exception as e: logging.exception(f"Error deleting user: {e}") @@ -327,9 +326,9 @@ async def remove_user_controller(payload: UserDeleteRequest): @app.post(Routes.USER_UPDATE) -async def update_user_controller(payload: UserUpdateRequest): +async def update_user_controller(request: UserUpdateRequest): try: - await db._update_user(payload.user, payload.password, payload.role) + await db._update_user(**request.model_dump()) return {"detail": "User updated successfully"} except Exception as e: logging.exception(f"Error updating user: {e}") @@ -337,9 +336,9 @@ async def update_user_controller(payload: UserUpdateRequest): @app.post(Routes.USER_VERIFY) -async def verify_user_controller(payload: UserVerifyRequest): +async def verify_user_controller(request: UserVerifyRequest): try: - auth = await db._verify(payload.user, payload.password) + auth = await db._verify(**request.model_dump()) if auth: return auth raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED) From b85174031c8469cadc60bc9915013e7b7f541386 Mon Sep 17 00:00:00 2001 From: FerTV Date: Mon, 22 Sep 2025 12:03:48 +0200 Subject: [PATCH 11/14] chore: init.py added for database files --- nebula/database/__init__.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 nebula/database/__init__.py diff --git a/nebula/database/__init__.py b/nebula/database/__init__.py new file mode 100644 index 000000000..e69de29bb From 2167951adc40fb4d20224bb4907efc8a44bedd6c Mon Sep 17 00:00:00 2001 From: FerTV Date: Mon, 22 Sep 2025 12:06:21 +0200 Subject: [PATCH 12/14] chore: init file added for postgres --- nebula/database/adapters/postgress/__init__.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 nebula/database/adapters/postgress/__init__.py diff --git a/nebula/database/adapters/postgress/__init__.py b/nebula/database/adapters/postgress/__init__.py new file mode 100644 index 000000000..e69de29bb From 916c359ed0da44464fff3e02ba1108d7dab8377c Mon Sep 17 00:00:00 2001 From: FerTV Date: Mon, 22 Sep 2025 12:08:39 +0200 Subject: [PATCH 13/14] chore: init file added to adapaters --- nebula/database/adapters/__inir__.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 nebula/database/adapters/__inir__.py diff --git a/nebula/database/adapters/__inir__.py b/nebula/database/adapters/__inir__.py new file mode 100644 index 000000000..e69de29bb From c0eaaea7e4d8253e53ae10127570cb3ab44a0275 Mon Sep 17 00:00:00 2001 From: FerTV Date: Mon, 22 Sep 2025 12:09:59 +0200 Subject: [PATCH 14/14] fix: typo in init file --- nebula/database/adapters/{__inir__.py => __init__.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename nebula/database/adapters/{__inir__.py => __init__.py} (100%) diff --git a/nebula/database/adapters/__inir__.py b/nebula/database/adapters/__init__.py similarity index 100% rename from nebula/database/adapters/__inir__.py rename to nebula/database/adapters/__init__.py