| 1 | import os |
| 2 | |
| 3 | from loguru import logger |
| 4 | from sqlalchemy import create_engine |
| 5 | from sqlalchemy import text |
| 6 | from sqlalchemy.exc import OperationalError |
| 7 | from sqlalchemy.exc import SQLAlchemyError |
| 8 | from sqlalchemy.ext.asyncio import AsyncSession |
| 9 | |
| 10 | # ! New with Async |
| 11 | from sqlmodel import SQLModel |
| 12 | |
| 13 | from alembic import command |
| 14 | from alembic.config import Config |
| 15 | from app.auth.services.universal import create_admin_user |
| 16 | from app.auth.services.universal import create_scheduler_user |
| 17 | from app.auth.services.universal import remove_scheduler_user |
| 18 | from app.db.db_populate import add_available_integrations_auth_keys_if_not_exist |
| 19 | from app.db.db_populate import add_available_integrations_if_not_exist |
| 20 | from app.db.db_populate import add_available_network_connectors_auth_keys_if_not_exist |
| 21 | from app.db.db_populate import add_available_network_connectors_if_not_exist |
| 22 | from app.db.db_populate import add_connectors_if_not_exist |
| 23 | from app.db.db_populate import add_roles_if_not_exist |
| 24 | from app.db.db_populate import delete_connectors_if_exist |
| 25 | from app.db.db_session import SQLALCHEMY_DATABASE_URI |
| 26 | from app.db.db_session import db_password |
| 27 | from app.schedulers.routes.scheduler import delete_job |
| 28 | |
| 29 | |
| 30 | async def create_database_if_not_exists(db_url: str, db_name: str): |
| 31 | """ |
| 32 | Create a database if it does not already exist. |
| 33 | |
| 34 | Args: |
| 35 | db_url (str): Database URL to connect to MySQL server (without database part). |
| 36 | db_name (str): The name of the database to create. |
| 37 | """ |
| 38 | engine = create_engine(db_url) |
| 39 | conn = engine.connect() |
| 40 | try: |
| 41 | # Check if database exists |
| 42 | conn.execute(text("commit")) |
| 43 | exists = conn.execute(text(f"SHOW DATABASES LIKE '{db_name}';")).fetchone() |
| 44 | if not exists: |
| 45 | # Create database if it does not exist |
| 46 | conn.execute(text("commit")) |
| 47 | conn.execute(text(f"CREATE DATABASE {db_name};")) |
| 48 | logger.info(f"Database '{db_name}' created successfully.") |
| 49 | else: |
| 50 | logger.info(f"Database '{db_name}' already exists.") |
| 51 | except SQLAlchemyError as e: |
| 52 | print(f"An error occurred: {e}") |
| 53 | finally: |
| 54 | conn.close() |
| 55 | engine.dispose() |
| 56 | |
| 57 | |
| 58 | async def create_copilot_user_if_not_exists(db_url: str, db_user_name: str): |
| 59 | """ |
| 60 | Create a user if it does not already exist. |
| 61 | |
| 62 | Args: |
| 63 | db_url (str): Database URL to connect to MySQL server (without database part). |
| 64 | db_user_name (str): The name of the user to create. |
| 65 | """ |
| 66 | db_name = "copilot" |
| 67 | engine = create_engine(db_url) |
| 68 | conn = engine.connect() |
| 69 | try: |
| 70 | # Check if user exists |
| 71 | conn.execute(text("commit")) |
| 72 | exists = conn.execute(text(f"SELECT * FROM mysql.user WHERE user = '{db_user_name}';")).fetchone() |
| 73 | if not exists: |
| 74 | # Create user if it does not exist |
| 75 | conn.execute(text("commit")) |
| 76 | conn.execute(text(f"CREATE USER '{db_user_name}'@'%' IDENTIFIED BY '{db_password}';")) |
| 77 | logger.info(f"User '{db_user_name}' created successfully with password '{db_password}'.") |
| 78 | conn.execute(text(f"GRANT ALL PRIVILEGES ON {db_name}.* TO '{db_user_name}'@'%';")) |
| 79 | logger.info(f"User '{db_user_name}' created successfully and granted all privileges to the '{db_name}' database.") |
| 80 | else: |
| 81 | logger.info(f"User '{db_user_name}' already exists.") |
| 82 | except SQLAlchemyError as e: |
| 83 | logger.info(f"An error occurred: {e}") |
| 84 | |
| 85 | |
| 86 | # def apply_migrations(): |
| 87 | # """ |
| 88 | # Applies Alembic migrations to ensure the database schema is up to date. |
| 89 | # """ |
| 90 | # logger.info("Applying migrations") |
| 91 | |
| 92 | # # Navigate up three levels from db_setup.py to the backend directory, then to the alembic directory |
| 93 | # base_dir = os.path.dirname(os.path.dirname(os.path.dirname(__file__))) |
| 94 | # alembic_directory = os.path.join(base_dir, "alembic") |
| 95 | |
| 96 | # logger.info(f"base_dir: {base_dir}") |
| 97 | # logger.info(f"Alembic directory: {alembic_directory}") |
| 98 | |
| 99 | # alembic_cfg = Config(os.path.join(alembic_directory, "alembic.ini")) |
| 100 | # alembic_cfg.set_main_option("sqlalchemy.url", SQLALCHEMY_DATABASE_URI.replace("+aiomysql", "+pymysql")) |
| 101 | # alembic_cfg.set_main_option("script_location", alembic_directory) |
| 102 | |
| 103 | # # Apply migrations to the latest revision |
| 104 | # try: |
| 105 | # command.upgrade(alembic_cfg, "head") |
| 106 | # except Exception as e: # Catch any exception |
| 107 | # logger.error(f"Error applying migrations: {e}") |
| 108 | # raise e |
| 109 | |
| 110 | |
| 111 | def apply_migrations(): |
| 112 | """ |
| 113 | Applies Alembic migrations to ensure the database schema is up to date. |
| 114 | """ |
| 115 | logger.info("Applying migrations") |
| 116 | |
| 117 | # Navigate up three levels from db_setup.py to the backend directory, then to the alembic directory |
| 118 | base_dir = os.path.dirname(os.path.dirname(os.path.dirname(__file__))) |
| 119 | alembic_directory = os.path.join(base_dir, "alembic") |
| 120 | |
| 121 | logger.info(f"base_dir: {base_dir}") |
| 122 | logger.info(f"Alembic directory: {alembic_directory}") |
| 123 | |
| 124 | alembic_cfg = Config(os.path.join(alembic_directory, "alembic.ini")) |
| 125 | alembic_cfg.set_main_option("sqlalchemy.url", SQLALCHEMY_DATABASE_URI.replace("+aiomysql", "+pymysql")) |
| 126 | alembic_cfg.set_main_option("script_location", alembic_directory) |
| 127 | |
| 128 | # Check current revision first |
| 129 | logger.info("Checking current database revision...") |
| 130 | try: |
| 131 | from sqlalchemy import create_engine |
| 132 | |
| 133 | from alembic.script import ScriptDirectory |
| 134 | |
| 135 | # Get current revision |
| 136 | engine = create_engine(SQLALCHEMY_DATABASE_URI.replace("+aiomysql", "+pymysql")) |
| 137 | with engine.connect() as connection: |
| 138 | from alembic.runtime.migration import MigrationContext |
| 139 | |
| 140 | context = MigrationContext.configure(connection) |
| 141 | current_rev = context.get_current_revision() |
| 142 | logger.info(f"Current database revision: {current_rev}") |
| 143 | |
| 144 | # Get head revision |
| 145 | script = ScriptDirectory.from_config(alembic_cfg) |
| 146 | head_rev = script.get_current_head() |
| 147 | logger.info(f"Target head revision: {head_rev}") |
| 148 | |
| 149 | if current_rev == head_rev: |
| 150 | logger.info("Database is already up to date!") |
| 151 | return |
| 152 | |
| 153 | except Exception as e: |
| 154 | logger.warning(f"Could not check current revision: {e}") |
| 155 | |
| 156 | # Apply migrations to the latest revision |
| 157 | logger.info("Starting migration upgrade...") |
| 158 | try: |
| 159 | command.upgrade(alembic_cfg, "head") |
| 160 | logger.info("Migrations completed successfully!") |
| 161 | except Exception as e: # Catch any exception |
| 162 | logger.error(f"Error applying migrations: {e}") |
| 163 | raise e |
| 164 | |
| 165 | |
| 166 | async def add_connectors(async_engine): |
| 167 | """ |
| 168 | Adds connectors to the database. |
| 169 | |
| 170 | Args: |
| 171 | async_engine (AsyncEngine): The async engine used to connect to the database. |
| 172 | |
| 173 | Returns: |
| 174 | None |
| 175 | """ |
| 176 | logger.info("Adding connectors") |
| 177 | async with AsyncSession( |
| 178 | async_engine, |
| 179 | ) as session: # Create an AsyncSession, not just a connection |
| 180 | async with session.begin(): # Start a transaction |
| 181 | await add_connectors_if_not_exist(session) |
| 182 | logger.info("Connectors added successfully") |
| 183 | |
| 184 | |
| 185 | async def delete_connectors(async_engine): |
| 186 | """ |
| 187 | Deletes connectors from the database. |
| 188 | |
| 189 | Args: |
| 190 | async_engine (AsyncEngine): The async engine used to connect to the database. |
| 191 | |
| 192 | Returns: |
| 193 | None |
| 194 | """ |
| 195 | logger.info("Deleting connectors") |
| 196 | async with AsyncSession( |
| 197 | async_engine, |
| 198 | ) as session: # Create an AsyncSession, not just a connection |
| 199 | async with session.begin(): # Start a transaction |
| 200 | await delete_connectors_if_exist(session) |
| 201 | logger.info("Connectors deleted successfully") |
| 202 | |
| 203 | |
| 204 | async def create_tables(async_engine): |
| 205 | """ |
| 206 | Creates tables in the database. |
| 207 | |
| 208 | Args: |
| 209 | async_engine (AsyncEngine): The async engine to connect to the database. |
| 210 | |
| 211 | Returns: |
| 212 | None |
| 213 | """ |
| 214 | logger.info("Creating tables") |
| 215 | async with async_engine.begin() as conn: |
| 216 | # This will create all tables |
| 217 | await conn.run_sync(SQLModel.metadata.create_all) |
| 218 | # Use AsyncSession for adding connectors |
| 219 | async with AsyncSession(async_engine) as session: |
| 220 | async with session.begin(): |
| 221 | await add_connectors_if_not_exist(session) |
| 222 | |
| 223 | |
| 224 | async def update_tables(async_engine): |
| 225 | """ |
| 226 | Updates tables in the database. Needed for adding new columns to existing tables. |
| 227 | |
| 228 | Args: |
| 229 | async_engine (AsyncEngine): The async engine to connect to the database. |
| 230 | |
| 231 | Returns: |
| 232 | None |
| 233 | """ |
| 234 | logger.info("Updating tables") |
| 235 | |
| 236 | # Define the new columns to be added |
| 237 | new_columns = { |
| 238 | "scheduled_job_metadata": ["extra_data TEXT"], |
| 239 | "customer_provisioning_default_settings": ["wazuh_worker_hostname TEXT"], |
| 240 | "agents": ["wazuh_agent_status TEXT"], |
| 241 | } |
| 242 | |
| 243 | async with async_engine.begin() as conn: |
| 244 | for table_name, columns in new_columns.items(): |
| 245 | for column in columns: |
| 246 | alter_table_query = text(f"ALTER TABLE {table_name} ADD COLUMN {column}") |
| 247 | try: |
| 248 | await conn.execute(alter_table_query) |
| 249 | except OperationalError as e: |
| 250 | if "duplicate column name" in str(e): |
| 251 | logger.info(f"Column {column} already exists in {table_name}") |
| 252 | else: |
| 253 | raise |
| 254 | |
| 255 | |
| 256 | async def create_roles(async_engine): |
| 257 | """ |
| 258 | Creates roles in the database. |
| 259 | |
| 260 | Args: |
| 261 | async_engine (AsyncEngine): The async engine used to connect to the database. |
| 262 | |
| 263 | Returns: |
| 264 | None |
| 265 | """ |
| 266 | logger.info("Creating roles") |
| 267 | async with AsyncSession( |
| 268 | async_engine, |
| 269 | ) as session: # Create an AsyncSession, not just a connection |
| 270 | async with session.begin(): # Start a transaction |
| 271 | await add_roles_if_not_exist(session) |
| 272 | |
| 273 | |
| 274 | async def create_available_integrations(async_engine): |
| 275 | """ |
| 276 | Creates available integrations in the database. |
| 277 | |
| 278 | Args: |
| 279 | async_engine (AsyncEngine): The async engine used to connect to the database. |
| 280 | |
| 281 | Returns: |
| 282 | None |
| 283 | """ |
| 284 | logger.info("Creating available integrations") |
| 285 | async with AsyncSession( |
| 286 | async_engine, |
| 287 | ) as session: # Create an AsyncSession, not just a connection |
| 288 | try: |
| 289 | await add_available_integrations_if_not_exist(session) |
| 290 | await add_available_integrations_auth_keys_if_not_exist(session) |
| 291 | except Exception as e: |
| 292 | logger.error(f"Error creating available integrations: {e}") |
| 293 | await session.rollback() # Explicit rollback on error |
| 294 | raise # Re-raise the exception to handle it further up the call stack |
| 295 | else: |
| 296 | await session.commit() # Explicit commit if all operations are successful |
| 297 | |
| 298 | |
| 299 | async def create_available_network_connectors(async_engine): |
| 300 | """ |
| 301 | Creates available network connectors in the database. |
| 302 | |
| 303 | Args: |
| 304 | async_engine (AsyncEngine): The async engine used to connect to the database. |
| 305 | |
| 306 | Returns: |
| 307 | None |
| 308 | """ |
| 309 | logger.info("Creating available network connectors") |
| 310 | async with AsyncSession( |
| 311 | async_engine, |
| 312 | ) as session: # Create an AsyncSession, not just a connection |
| 313 | try: |
| 314 | await add_available_network_connectors_if_not_exist(session) |
| 315 | await add_available_network_connectors_auth_keys_if_not_exist(session) |
| 316 | except Exception as e: |
| 317 | logger.error(f"Error creating available integrations: {e}") |
| 318 | await session.rollback() # Explicit rollback on error |
| 319 | raise # Re-raise the exception to handle it further up the call stack |
| 320 | else: |
| 321 | await session.commit() # Explicit commit if all operations are successful |
| 322 | |
| 323 | |
| 324 | async def ensure_admin_user(async_engine): |
| 325 | """ |
| 326 | Ensures that an admin user exists in the database. |
| 327 | |
| 328 | Args: |
| 329 | async_engine (AsyncEngine): The async engine used to connect to the database. |
| 330 | |
| 331 | Returns: |
| 332 | None |
| 333 | """ |
| 334 | logger.info("Ensuring admin user exists") |
| 335 | async with AsyncSession(async_engine) as session: |
| 336 | async with session.begin(): |
| 337 | # Pass the session to the inner function |
| 338 | await create_admin_user(session) |
| 339 | |
| 340 | |
| 341 | async def ensure_scheduler_user(async_engine): |
| 342 | """ |
| 343 | Ensures that the scheduler user exists in the database. |
| 344 | |
| 345 | Args: |
| 346 | async_engine (AsyncEngine): The async engine used to connect to the database. |
| 347 | |
| 348 | Returns: |
| 349 | None |
| 350 | """ |
| 351 | logger.info("Ensuring scheduler user exists") |
| 352 | async with AsyncSession(async_engine) as session: |
| 353 | async with session.begin(): |
| 354 | # Pass the session to the inner function |
| 355 | await create_scheduler_user(session) |
| 356 | |
| 357 | |
| 358 | async def delete_job_if_exists(async_engine): |
| 359 | """ |
| 360 | Deletes a job from the database if it exists. |
| 361 | |
| 362 | Args: |
| 363 | job_id (str): The ID of the job to delete. |
| 364 | |
| 365 | Returns: |
| 366 | None |
| 367 | """ |
| 368 | job_id = "wazuh_index_fields_resize" |
| 369 | logger.info(f"Deleting job with ID {job_id}") |
| 370 | async with AsyncSession(async_engine) as session: |
| 371 | async with session.begin(): |
| 372 | # Pass the session to the inner function |
| 373 | await delete_job(session, job_id) |
| 374 | |
| 375 | |
| 376 | async def ensure_scheduler_user_removed(async_engine): |
| 377 | """ |
| 378 | Ensures that the scheduler user is removed from the database. |
| 379 | |
| 380 | Args: |
| 381 | async_engine (AsyncEngine): The async engine used to connect to the database. |
| 382 | |
| 383 | Returns: |
| 384 | None |
| 385 | """ |
| 386 | logger.info("Ensuring scheduler user exists") |
| 387 | async with AsyncSession(async_engine) as session: |
| 388 | async with session.begin(): |
| 389 | # Pass the session to the inner function |
| 390 | await remove_scheduler_user(session) |