diff --git a/Makefile b/Makefile index c1e0cd768..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/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 7b624358e..889421330 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.web_app_controller import TermEscapeCodeFormatter +from nebula.controller.hub import TermEscapeCodeFormatter from nebula.controller.scenarios import ScenarioManagement from nebula.utils import DockerUtils, FileUtils, SocketUtils @@ -638,7 +638,7 @@ def __init__(self, args): sys.exit(1) self.controller_port = int(args.controllerport) if hasattr(args, "controllerport") else 5050 - self.federation_controller_port = int(args.federationcontrollerport) if hasattr(args, "federationcontrollerport") else 5051 + self.federation_controller_port = int(args.federationcontrollerport) if hasattr(args, "federationcontrollerport") else 5052 self.waf_port = int(args.wafport) if hasattr(args, "wafport") else 6000 self.frontend_port = int(args.webport) if hasattr(args, "webport") else 6060 self.grafana_port = int(args.grafanaport) if hasattr(args, "grafanaport") else 6040 @@ -853,7 +853,7 @@ def start(self): # Check ports available if not SocketUtils.is_port_open(self.controller_port): self.controller_port = SocketUtils.find_free_port(start_port=self.controller_port) - + if not SocketUtils.is_port_open(self.federation_controller_port): self.federation_controller_port = SocketUtils.find_free_port(start_port=self.federation_controller_port) @@ -1045,17 +1045,24 @@ 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/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) 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")} @@ -1068,6 +1075,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) @@ -1132,11 +1140,7 @@ def run_controller(self): "NEBULA_FEDERATION_CONTROLLER_PORT" : self.federation_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/app/main.py b/app/main.py index e822bfb4d..5d14308af 100755 --- a/app/main.py +++ b/app/main.py @@ -21,8 +21,8 @@ "-fcp", "--federationcontrollerport", dest="federationcontrollerport", - default=5051, - help="federation controller port port (default: 5051)", + default=5052, + help="federation controller port port (default: 5052)", ) argparser.add_argument( diff --git a/nebula/controller/database.py b/nebula/controller/database.py deleted file mode 100755 index ed4ead816..000000000 --- a/nebula/controller/database.py +++ /dev/null @@ -1,638 +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 - finally: - if conn: - await conn.close() - - -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 (same as get_all_scenarios) - 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/web_app_controller.py b/nebula/controller/hub.py similarity index 76% rename from nebula/controller/web_app_controller.py rename to nebula/controller/hub.py index 7835f2a1a..495feb529 100755 --- a/nebula/controller/web_app_controller.py +++ b/nebula/controller/hub.py @@ -15,17 +15,13 @@ import uvicorn 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.controller.http_helpers import remote_get, remote_post_form -from nebula.utils import DockerUtils, APIUtils -from nebula.controller.federation.utils_requests import RunScenarioRequest, StopScenarioRequest, factory_requests_path +from nebula.utils import APIUtils, DockerUtils +import nebula.controller.federation.utils_requests as federation_requests +import nebula.controller.utils_requests as controller_requests + +# URL for the database API +DATABASE_API_URL = os.environ.get("NEBULA_DATABASE_API_URL", "http://nebula-database:5051") # Setup controller logger @@ -107,28 +103,20 @@ def configure_logger(controller_log): handler.setFormatter(logging.Formatter("[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s")) logger.addHandler(handler) -id_counter = 1 - @asynccontextmanager 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 init_db_pool() - await insert_default_admin() - yield # Code to run on shutdown - await close_db_pool() + pass # Initialize FastAPI app outside the Controller class @@ -136,7 +124,7 @@ async def lifespan(app: FastAPI): # Define endpoints outside the Controller class -@app.get("/") +@app.get(controller_requests.Routes.INIT) async def read_root(): """ Root endpoint of the NEBULA Controller API. @@ -147,7 +135,7 @@ async def read_root(): return {"message": "Welcome to the NEBULA Controller API"} -@app.get("/status") +@app.get(controller_requests.Routes.STATUS) async def get_status(): """ Check the status of the NEBULA Controller API. @@ -158,7 +146,7 @@ async def get_status(): return {"status": "NEBULA Controller API is running"} -@app.get("/resources") +@app.get(controller_requests.Routes.RESOURCES) async def get_resources(): """ Get system resource usage including RAM and GPU memory usage. @@ -200,7 +188,7 @@ async def get_resources(): } -@app.get("/least_memory_gpu") +@app.get(controller_requests.Routes.LEAST_MEMORY_GPU) async def get_least_memory_gpu(): """ Identify the GPU with the highest memory usage above a threshold (50%). @@ -242,7 +230,7 @@ async def get_least_memory_gpu(): } -@app.get("/available_gpus/") +@app.get(controller_requests.Routes.AVAILABLE_GPUS) async def get_available_gpu(): """ Get the list of GPUs with memory usage below 5%. @@ -301,10 +289,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(controller_requests.Routes.RUN) +async def run_scenario(run_scenario_request: controller_requests.RunScenarioRequest): """ Launches a new scenario based on the provided configuration. @@ -316,58 +302,40 @@ async def run_scenario( Returns: str: The name of the scenario that was started. """ - - import subprocess - global id_counter - from nebula.controller.scenarios import ScenarioManagement try: fed_controller_port = os.environ.get("NEBULA_FEDERATION_CONTROLLER_PORT") fed_controller_host = os.environ.get("NEBULA_CONTROLLER_HOST") - url_init_fed_controller = f"http://{fed_controller_host}:{fed_controller_port}" + factory_requests_path("init") - url_run_scenario = f"http://{fed_controller_host}:{fed_controller_port}" + factory_requests_path("run") + url_init_fed_controller = f"http://{fed_controller_host}:{fed_controller_port}" + federation_requests.factory_requests_path("init") + url_run_scenario = f"http://{fed_controller_host}:{fed_controller_port}" + federation_requests.factory_requests_path("run") #init_fed_req = InitFederationRequest(experiment_type="docker") - run_scenario_req = RunScenarioRequest(scenario_data=scenario_data, federation_id=f"id_nebula_{id_counter}", user=user) #TODO ID per experiment - id_counter += 1 + run_scenario_req = federation_requests.RunScenarioRequest(scenario_data=run_scenario_request.scenario_data, federation_id=f"id_nebula_{1}", user=run_scenario_request.user) #TODO ID per experiment #await APIUtils.post(url_init_fed_controller, init_fed_req.model_dump()) await APIUtils.post(url_run_scenario, run_scenario_req.model_dump()) except Exception as e: logging.info(e) - validate_physical_fields(scenario_data) + # Unpack request data (role is intentionally ignored for now) + scenario_data = run_scenario_request.scenario_data + user = run_scenario_request.user - db_scenario = copy.deepcopy(scenario_data) + # validate_physical_fields(scenario_data) # Manager for the actual scenario #scenarioManagement = ScenarioManagement(scenario_data, user) # await update_scenario( - # scenario_name=scenarioManagement.scenario_name, - # start_time=scenarioManagement.start_date_scenario, + # scenario_name="", #TODO scenario_name + # start_time="", #TODO start_time # end_time="", # scenario=scenario_data, # status="running", - # role=role, # username=user, # ) - # Run the actual scenario - # try: - # if scenarioManagement.scenario.mobility: - # additional_participants = scenario_data["additional_participants"] - # schema_additional_participants = scenario_data["schema_additional_participants"] - # await scenarioManagement.load_configurations_and_start_nodes( - # additional_participants, schema_additional_participants - # ) - # else: - # await scenarioManagement.load_configurations_and_start_nodes() - # except subprocess.CalledProcessError as e: - # logging.exception(f"Error docker-compose up: {e}") - # return - - return ""#scenarioManagement.scenario_name + return ""#scenarioManagement.scenario_name #TODO return -@app.post("/scenarios/stop") +@app.post(controller_requests.Routes.STOP) #TODO redo method async def stop_scenario( scenario_name: str = Body(..., embed=True), all: bool = Body(False, embed=True), @@ -394,13 +362,16 @@ async def stop_scenario( """ fed_controller_port = os.environ.get("NEBULA_FEDERATION_CONTROLLER_PORT") fed_controller_host = os.environ.get("NEBULA_CONTROLLER_HOST") - url_stop_scenario = f"http://{fed_controller_host}:{fed_controller_port}" + factory_requests_path("stop") - stop_scenario_req = StopScenarioRequest(federation_id="id_nebula") + url_stop_scenario = f"http://{fed_controller_host}:{fed_controller_port}" + federation_requests.factory_requests_path("stop") + stop_scenario_req = federation_requests.StopScenarioRequest(federation_id="id_nebula") try: + path = federation_requests.factory_requests_path("stop") + payload = federation_requests.StopScenarioRequest(scenario_name=scenario_name, all=all).model_dump() + await APIUtils.post(f"{DATABASE_API_URL}{path}", data=payload) await APIUtils.post(url_stop_scenario, stop_scenario_req.model_dump()) except Exception as e: logging.info(f"ERROR: sending stop scenario to federation Controller: {e}") - + # from nebula.controller.scenarios import ScenarioManagement # ScenarioManagement.cleanup_scenario_containers() @@ -414,7 +385,7 @@ async def stop_scenario( # raise HTTPException(status_code=500, detail="Internal server error") -@app.post("/scenarios/remove") +@app.post(controller_requests.Routes.REMOVE) async def remove_scenario( scenario_name: str = Body(..., embed=True), ): @@ -427,11 +398,12 @@ 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) + path = controller_requests.factory_requests_path("remove") + payload = controller_requests.ScenarioRemoveRequest(scenario_name=scenario_name).model_dump() + await APIUtils.post(f"{DATABASE_API_URL}{path}", data=payload) ScenarioManagement.remove_files_by_scenario(scenario_name) except Exception as e: @@ -441,7 +413,7 @@ async def remove_scenario( return {"message": f"Scenario {scenario_name} removed successfully"} -@app.get("/scenarios/{user}/{role}") +@app.get(controller_requests.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")], @@ -456,30 +428,21 @@ 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) - - if role == "admin": - scenario_running = await get_running_scenario() - else: - scenario_running = await get_running_scenario(username=user) + path = controller_requests.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") - return {"scenarios": scenarios, "scenario_running": scenario_running} - -@app.post("/scenarios/update") +@app.post(controller_requests.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), - role: str = Body(..., embed=True), username: str = Body(..., embed=True), ): """ @@ -491,24 +454,28 @@ 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. """ - from nebula.controller.database import scenario_update_record - try: - await scenario_update_record(scenario_name, start_time, end_time, scenario, status, username) + payload = controller_requests.ScenarioUpdateRequest( + scenario_name=scenario_name, + start_time=start_time, + end_time=end_time, + scenario=scenario, + status=status, + username=username, + ).model_dump() + path = controller_requests.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") - return {"message": f"Scenario {scenario_name} updated successfully"} - -@app.post("/scenarios/set_status_to_finished") +@app.post(controller_requests.Routes.FINISH) async def set_scenario_status_to_finished( scenario_name: str = Body(..., embed=True), all: bool = Body(False, embed=True) ): @@ -522,22 +489,17 @@ 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() - else: - await scenario_set_status_to_finished(scenario_name) + payload = controller_requests.ScenarioFinishRequest(scenario_name=scenario_name, all=all).model_dump() + path = controller_requests.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") - return {"message": f"Scenario {scenario_name} status set to finished successfully"} - -@app.get("/scenarios/running") -async def get_running_scenario(get_all: bool = False): +@app.get(controller_requests.Routes.RUNNING) +async def get_running_scenario_endpoint(get_all: bool = False): """ Retrieves the currently running scenario(s). @@ -547,16 +509,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) + path = controller_requests.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/{user}/{role}/{scenario_name}") +@app.get(controller_requests.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")], @@ -574,18 +535,16 @@ 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) - return {"allowed": allowed} + path = controller_requests.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}") raise HTTPException(status_code=500, detail="Internal server error") -@app.get("/scenarios/{scenario_name}") -async def get_scenario_by_name( +@app.get(controller_requests.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") ], @@ -599,19 +558,16 @@ 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) + path = controller_requests.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") - return scenario - -@app.get("/nodes/{scenario_name}") -async def list_nodes_by_scenario_name( +@app.get(controller_requests.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") ], @@ -625,18 +581,15 @@ 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) + path = controller_requests.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") - return nodes - -@app.post("/nodes/{scenario_name}/update") +@app.post(controller_requests.Routes.NODES_UPDATE_BY_SCENARIO) async def update_nodes( scenario_name: Annotated[ str, @@ -654,28 +607,24 @@ 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( - str(config["data"]["device_args"]["uid"]), - str(config["data"]["device_args"]["idx"]), - str(config["data"]["network_args"]["ip"]), - str(config["data"]["network_args"]["port"]), - str(config["data"]["device_args"]["role"]), - config["data"]["network_args"]["neighbors"], - str(config["data"]["addons"]["mobility"]["latitude"]), - str(config["data"]["addons"]["mobility"]["longitude"]), - str(timestamp), - str(config["data"]["data"]["scenario_args"]["federation"]), - str(config["data"]["federation_args"]["round"]), - str(config["data"]["scenario_args"]["name"]), - str(config["data"]["tracking_args"]["run_hash"]), - str(config["data"]["device_args"]["malicious"]), - ) + config:dict = await request.json() + config["timestamp"] = str(datetime.datetime.now()) + + mobility_args = config.get("mobility_args", None) + if not mobility_args: + # default Murcia coordinates if none provided + config["mobility_args"] = {"latitude": "38.0235", "longitude": "-1.1744"} + # Validate and normalize payload + validated = controller_requests.NodesUpdateRequest(**config) + + # 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 = controller_requests.factory_requests_path("update_nodes") + await APIUtils.post(f"{DATABASE_API_URL}{path}", data=payload) except Exception as e: logging.exception(f"Error updating nodes: {e}") raise HTTPException(status_code=500, detail="Internal server error") @@ -684,19 +633,10 @@ 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) + return await APIUtils.post(url, data=config) - 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"} - - -@app.post("/nodes/{scenario_name}/done") +@app.post(controller_requests.Routes.NODES_DONE_BY_SCENARIO) async def node_done( scenario_name: Annotated[ str, @@ -720,18 +660,11 @@ 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") -async def remove_nodes_by_scenario_name(scenario_name: str = Body(..., embed=True)): +@app.post(controller_requests.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. @@ -740,10 +673,10 @@ 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) + path = controller_requests.factory_requests_path("remove_nodes") + payload = controller_requests.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}") raise HTTPException(status_code=500, detail="Internal server error") @@ -751,7 +684,7 @@ async def remove_nodes_by_scenario_name(scenario_name: str = Body(..., embed=Tru return {"message": f"Nodes for scenario {scenario_name} removed successfully"} -@app.get("/notes/{scenario_name}") +@app.get(controller_requests.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") @@ -760,22 +693,15 @@ 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) - - if notes_record is not None: - notes_record = dict(notes_record.items()) - - return notes_record - + path = controller_requests.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(controller_requests.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. @@ -786,19 +712,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. """ - from nebula.controller.database import save_notes - try: - await save_notes(scenario_name, notes) + payload = controller_requests.NotesUpdateRequest(scenario_name=scenario_name, notes=notes).model_dump() + path = controller_requests.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") - return {"message": f"Notes for scenario {scenario_name} updated successfully"} - -@app.post("/notes/remove") -async def remove_notes_by_scenario_name(scenario_name: str = Body(..., embed=True)): +@app.post(controller_requests.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. @@ -807,10 +731,10 @@ 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) + path = controller_requests.factory_requests_path("remove_notes") + payload = controller_requests.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}") raise HTTPException(status_code=500, detail="Internal server error") @@ -818,7 +742,7 @@ async def remove_notes_by_scenario_name(scenario_name: str = Body(..., embed=Tru return {"message": f"Notes for scenario {scenario_name} removed successfully"} -@app.get("/user/list") +@app.get(controller_requests.Routes.USER_LIST) async def list_users_controller(all_info: bool = False): """ Endpoint to list all users in the database. @@ -828,20 +752,16 @@ 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) - 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} + path = controller_requests.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}") -async def get_user_by_scenario_name( +@app.get(controller_requests.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") ], @@ -854,18 +774,15 @@ 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) + path = controller_requests.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 {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") +@app.get(controller_requests.Routes.DISCOVER_VPN) async def discover_vpn(): """ Calls the Tailscale CLI to fetch the current status in JSON format, @@ -906,7 +823,7 @@ async def discover_vpn(): raise HTTPException(status_code=500, detail="No devices discovered") -@app.get("/physical/run/{ip}", tags=["physical"]) +@app.get(controller_requests.Routes.PHYSICAL_RUN, tags=["physical"]) async def physical_run(ip: str): status, data = await remote_get(ip, "/run/") @@ -917,7 +834,7 @@ async def physical_run(ip: str): raise HTTPException(status_code=status, detail=data) -@app.get("/physical/stop/{ip}", tags=["physical"]) +@app.get(controller_requests.Routes.PHYSICAL_STOP, tags=["physical"]) async def physical_stop(ip: str): status, data = await remote_get(ip, "/stop/") if status == 200: @@ -927,7 +844,7 @@ async def physical_stop(ip: str): raise HTTPException(status_code=status, detail=data) -@app.put("/physical/setup/{ip}", tags=["physical"], +@app.put(controller_requests.Routes.PHYSICAL_SETUP, tags=["physical"], status_code=status.HTTP_201_CREATED) async def physical_setup( ip: str, @@ -960,7 +877,7 @@ async def physical_setup( # โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ # Physical ยท single-node state # โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ -@app.get("/physical/state/{ip}", tags=["physical"]) +@app.get(controller_requests.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. @@ -997,7 +914,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(controller_requests.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. @@ -1018,11 +935,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 get_scenario_by_name_endpoint(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 list_nodes_by_scenario_name_endpoint(scenario_name) if not nodes: raise HTTPException(status_code=404, detail="No nodes found for scenario") @@ -1047,7 +964,7 @@ async def get_physical_scenario_state(scenario_name: str): } -@app.post("/user/add") +@app.post(controller_requests.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. @@ -1059,17 +976,16 @@ 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) - return {"detail": "User added successfully"} + payload = controller_requests.UserAddRequest(user=user, password=password, role=role).model_dump() + path = controller_requests.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(controller_requests.Routes.USER_DELETE) async def remove_user_controller(user: str = Body(..., embed=True)): """ Controller endpoint that inserts a new user into the database. @@ -1079,17 +995,16 @@ 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) - return {"detail": "User deleted successfully"} + path = controller_requests.factory_requests_path("delete_user") + payload = controller_requests.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}") raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error deleting user: {e}") -@app.post("/user/update") +@app.post(controller_requests.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. @@ -1101,17 +1016,16 @@ 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) - return {"detail": "User updated successfully"} + payload = controller_requests.UserUpdateRequest(user=user, password=password, role=role).model_dump() + path = controller_requests.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(controller_requests.Routes.USER_VERIFY) async def verify_user_controller(user: str = Body(...), password: str = Body(...)): """ Endpoint to verify user credentials. @@ -1122,16 +1036,13 @@ 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) - return {"user": user_submitted, "role": user_info[2]} - else: - raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED) - except Exception as e: + payload = controller_requests.UserVerifyRequest(user=user, password=password).model_dump() + path = controller_requests.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 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/controller/start_services.sh b/nebula/controller/start_services.sh index c0ecb234b..d506eab54 100644 --- a/nebula/controller/start_services.sh +++ b/nebula/controller/start_services.sh @@ -14,11 +14,11 @@ NEBULA_SOCK=nebula.sock echo "NEBULA_PRODUCTION: $NEBULA_PRODUCTION" if [ "$NEBULA_PRODUCTION" = "False" ]; then echo "Starting Gunicorn in dev mode..." - uvicorn nebula.controller.web_app_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 "*" & uvicorn nebula.controller.federation.federation_api:app --host 0.0.0.0 --port $NEBULA_FEDERATION_CONTROLLER_PORT --log-level debug --proxy-headers --forwarded-allow-ips "*" & else echo "Starting Gunicorn in production mode..." - uvicorn nebula.controller.web_app_controller:app --host 0.0.0.0 --port $NEBULA_CONTROLLER_PORT --log-level info --proxy-headers --forwarded-allow-ips "*" & + uvicorn nebula.controller.hub:app --host 0.0.0.0 --port $NEBULA_CONTROLLER_PORT --log-level info --proxy-headers --forwarded-allow-ips "*" & uvicorn nebula.controller.federation.federation_api:app --host 0.0.0.0 --port $NEBULA_FEDERATION_CONTROLLER_PORT --log-level debug --proxy-headers --forwarded-allow-ips "*" & fi 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/__init__.py b/nebula/database/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/nebula/database/adapters/__init__.py b/nebula/database/adapters/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/nebula/database/adapters/postgress/__init__.py b/nebula/database/adapters/postgress/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/nebula/database/Dockerfile b/nebula/database/adapters/postgress/docker/Dockerfile similarity index 84% rename from nebula/database/Dockerfile rename to nebula/database/adapters/postgress/docker/Dockerfile index 04ed953e9..2859de8da 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/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 new file mode 100644 index 000000000..911961294 --- /dev/null +++ b/nebula/database/adapters/postgress/docker/docker-entrypoint.sh @@ -0,0 +1,21 @@ +#!/bin/sh +set -x + +# 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." + + cd nebula + NEBULA_SOCK=nebula.sock + + 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/init-configs.sql b/nebula/database/adapters/postgress/docker/init-configs.sql similarity index 78% rename from nebula/database/init-configs.sql rename to nebula/database/adapters/postgress/docker/init-configs.sql index a34b17841..9370a31d7 100644 --- a/nebula/database/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, @@ -21,17 +21,26 @@ 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 ); --- 3) Configs como JSONB +-- 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; CREATE TABLE configs ( diff --git a/nebula/database/adapters/postgress/postgress.py b/nebula/database/adapters/postgress/postgress.py new file mode 100755 index 000000000..7d2c1a37e --- /dev/null +++ b/nebula/database/adapters/postgress/postgress.py @@ -0,0 +1,746 @@ +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: 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 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: str): + """ + 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: str, password: str): + """ + 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: + 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: str): + """ + 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: str): + """ + 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:str, password:str, role:str): + """ + 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:str, password:str, role:str): + """ + 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:str=None, sort_by:str="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) + + # 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 + + + 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. + """ + 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) + 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, + 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. + """ + 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({}) + + # 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;", + node_uid, scenario + ) + + if result is None: + # Insert new node + await conn.execute( + """ + 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, + timestamp, federation, federation_round, scenario, run_hash, extras_payload, malicious_payload, + ) + else: + # Update existing node + await conn.execute( + """ + UPDATE nodes SET idx = $1, ip = $2, port = $3, role = $4, neighbors = $5, + timestamp = $6, federation = $7, round = $8, + hash = $9, extras = $10::jsonb, malicious = $11 + WHERE uid = $12 AND scenario = $13; + """, + idx, ip, port, role, neighbors, + timestamp, federation, federation_round, + run_hash, extras_payload, malicious_payload, + 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:str): + """ + 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: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. + """ + 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, 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. + """ + # 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(user) + + 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(user, role, sort_by) + + return scenarios_to_return + + + 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 is a dictionary before dumping to JSON + if not isinstance(scenario, dict): + try: + scenario = json.loads(scenario) + except (json.JSONDecodeError, TypeError): + logging.error("scenario 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, scenario_name, start_time, end_time, username, status, json.dumps(scenario)) + + + 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:str): + """ + 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:str): + """ + 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 _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). + """ + 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_scenarios(self, user: str, role: str): + """ + Compose scenarios list and running scenario respecting 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} + + + async def _get_scenario_by_name(self, scenario_name:str): + """ + 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:str): + """ + 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:str): + """ + 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:str): + """ + 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: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) + + if not scenario_info: + return False # Scenario does not exist + + if role == "admin": + return True # Admins can access any existing scenario + + if user is None: + logging.warning( + "check_scenario_with_role called for non-admin role without user." + ) + return False + + return scenario_info.get("username") == user + + # --- Notes Management Functions --- + + async def _save_notes(self, scenario: str, notes: str): + """ + 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: str): + """ + Retrieve notes associated with a specific scenario. + """ + async with self.pool.acquire() as conn: + 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: str): + """ + 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) 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..250464fb1 --- /dev/null +++ b/nebula/database/database_adapter_interface.py @@ -0,0 +1,198 @@ +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, + extras, + timestamp, + federation, + federation_round, + scenario, + run_hash, + malicious, + ): + """Inserts or updates a node record. Latitude/longitude must be included in `extras` (JSON).""" + 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 + + # --- 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 new file mode 100644 index 000000000..579a11a05 --- /dev/null +++ b/nebula/database/database_api.py @@ -0,0 +1,350 @@ + +import logging +import os +import sys + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../'))) + +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, + 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") + + +# 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(Routes.INIT) +async def read_root(): + return {"message": "Welcome to the NEBULA Database API"} + + +# Scenarios +@app.post(Routes.UPDATE) +async def update_scenario( + request: ScenarioUpdateRequest, +): + try: + await db._scenario_update_record( + **request.model_dump() + ) + return {"message": f"Scenario {request.scenario_name} updated successfully"} + except Exception as e: + logging.exception( + 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( + request: ScenarioStopRequest, +): + try: + 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 {request.scenario_name}: {e}" + ) + raise HTTPException(status_code=500, detail="Internal server error") + + +@app.post(Routes.REMOVE) +async def remove_scenario( + request: ScenarioRemoveRequest, +): + try: + 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 {request.scenario_name}: {e}" + ) + raise HTTPException(status_code=500, detail="Internal server error") + + +@app.get(Routes.GET_SCENARIOS_BY_USER) +async def get_scenarios( + request: GetScenariosRequest = Depends() +): + try: + 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") + + +@app.post(Routes.FINISH) +async def set_scenario_status_to_finished( + request: ScenarioFinishRequest, +): + try: + await db._finish_scenario( + request.scenario_name, request.all + ) + return {"message": "Finished status set successfully"} + except Exception as e: + logging.exception( + 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(request: GetRunningScenarioRequest = Depends()): + try: + 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") + + +@app.get(Routes.CHECK_SCENARIO) +async def check_scenario( + request: CheckScenarioRequest = Depends() +): + try: + 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}") + raise HTTPException(status_code=500, detail="Internal server error") + + +@app.get(Routes.GET_SCENARIOS_BY_SCENARIO_NAME) +async def get_scenario_by_name_endpoint( + request: GetScenarioByNameRequest = Depends(), +): + try: + scenario = await db._get_scenario_by_name(request.scenario_name) + return scenario + except Exception as 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( + request: ListNodesByScenarioNameRequest = Depends() +): + try: + nodes = await db._list_nodes_by_scenario_name(request.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(Routes.NODES_UPDATE) +async def update_node_record(request: NodesUpdateRequest): + try: + # Build extras from mobility_args + extras = { + "latitude": request.mobility_args.latitude, + "longitude": request.mobility_args.longitude, + } + await db._update_node_record( + 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(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: + logging.exception(f"Error updating node: {e}") + raise HTTPException(status_code=500, detail="Internal server error") + + +@app.post(Routes.NODES_REMOVE) +async def remove_nodes_by_scenario_name_endpoint(request: NodesRemoveRequest): + try: + 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") + + +# Notes +@app.get(Routes.NOTES_BY_SCENARIO_NAME) +async def get_notes_by_scenario_name( + request: GetNotesByScenarioNameRequest = Depends() +): + try: + notes_record = await db._get_notes(request.scenario_name) + return notes_record + except Exception as 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(request: NotesUpdateRequest): + try: + 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(request: NotesRemoveRequest): + try: + 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") + + +# Users +@app.get(Routes.USER_LIST) +async def list_users_controller(request: ListUsersRequest = Depends()): + try: + 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}") + + +@app.get(Routes.USER_BY_SCENARIO_NAME) +async def get_user_by_scenario_name_endpoint( + request: GetUserByScenarioNameRequest = Depends() +): + try: + 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 {request.scenario_name}: {e}") + raise HTTPException(status_code=500, detail="Internal server error") + + +@app.post(Routes.USER_ADD) +async def add_user_controller(request: UserAddRequest): + try: + await db._add_user(**request.model_dump()) + 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(Routes.USER_DELETE) +async def remove_user_controller(request: UserDeleteRequest): + try: + await db._delete_user_from_db(request.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(Routes.USER_UPDATE) +async def update_user_controller(request: UserUpdateRequest): + try: + await db._update_user(**request.model_dump()) + 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(Routes.USER_VERIFY) +async def verify_user_controller(request: UserVerifyRequest): + try: + auth = await db._verify(**request.model_dump()) + 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}") diff --git a/nebula/database/docker-entrypoint.sh b/nebula/database/docker-entrypoint.sh deleted file mode 100644 index a298ca23b..000000000 --- a/nebula/database/docker-entrypoint.sh +++ /dev/null @@ -1,18 +0,0 @@ -#!/bin/sh -set -e - -# 1) Run the original entrypoint and wait for it to finish initialization -/usr/local/bin/docker-entrypoint.sh.orig "$@" - -# 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 - -# 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 diff --git a/nebula/database/utils_requests.py b/nebula/database/utils_requests.py new file mode 100644 index 000000000..8d9f88da2 --- /dev/null +++ b/nebula/database/utils_requests.py @@ -0,0 +1,217 @@ +from typing import Any, Dict, List + +from pydantic import BaseModel, confloat, conint + + +class Routes: + # Scenarios + INIT = "/" + 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_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" + + +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 + +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": + 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/frontend/app.py b/nebula/frontend/app.py index 42f5b3a68..59c836a1b 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) @@ -1595,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"], @@ -1603,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"]), diff --git a/nebula/utils.py b/nebula/utils.py index 34b0c9621..11f127c8b 100644 --- a/nebula/utils.py +++ b/nebula/utils.py @@ -2,12 +2,20 @@ import os import socket +import aiohttp 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 +import re +from typing import Optional + from fastapi import HTTPException from aiohttp import ClientConnectorError from aiohttp.client_exceptions import ClientError @@ -209,10 +217,10 @@ 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, @@ -267,9 +275,9 @@ def configure_logger( logger.propagate = False return logger - + class APIUtils(): - + @staticmethod async def retry_with_backoff(func, *args, max_retries=5, initial_delay=1): """ @@ -304,12 +312,13 @@ async def retry_with_backoff(func, *args, max_retries=5, initial_delay=1): raise last_exception @staticmethod - async def get(url): + 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. @@ -320,11 +329,12 @@ async def get(url): async def _get(): async with aiohttp.ClientSession() as session: - async with session.get(url) as response: + async with session.get(url, params=params) as response: if response.status == 200: return await response.json() else: - raise HTTPException(status_code=response.status, detail="Error fetching data") + detail = await response.text() + raise HTTPException(status_code=response.status, detail=detail) return await APIUtils.retry_with_backoff(_get) @@ -354,5 +364,3 @@ async def _post(): raise HTTPException(status_code=response.status, detail=detail) return await APIUtils.retry_with_backoff(_post) - - \ No newline at end of file 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",