DataOps-platform/app/commands/bootstrap_admin.py

from __future__ import annotations

import argparse
import getpass
import logging
import os
import re

from sqlalchemy import text

from app import create_app, db
from app.core.common.identifiers import new_governance_uid
from app.core.system.auth import hash_password


logger = logging.getLogger(__name__)


def validate_password(password: str) -> None:
    if len(password) < 10:
        raise ValueError("password must contain at least 10 characters")
    if not re.search(r"[a-z]", password):
        raise ValueError("password must contain a lowercase letter")
    if not re.search(r"[A-Z]", password):
        raise ValueError("password must contain an uppercase letter")
    if not re.search(r"[0-9]", password):
        raise ValueError("password must contain a digit")


def log_bootstrap_result(username: str, *, created: bool) -> None:
    if created:
        logger.info("Created first platform administrator: %s", username)
    else:
        logger.info("Platform administrator already exists; bootstrap skipped")


def bootstrap_first_admin(session, *, username: str, password: str) -> str:
    validate_password(password)
    existing = session.execute(
        text(
            "SELECT COUNT(*) FROM public.users u "
            "JOIN public.user_roles ur ON ur.user_id = u.id "
            "JOIN public.roles r ON r.id = ur.role_id "
            "WHERE r.name = 'admin' AND u.status = 'active'"
        )
    ).scalar_one()
    if existing:
        raise RuntimeError("an active administrator already exists")

    user_id = new_governance_uid()
    session.execute(
        text(
            "INSERT INTO public.users "
            "(id, username, display_name, password_hash) "
            "VALUES (CAST(:id AS uuid), :username, :username, :password_hash)"
        ),
        {"id": user_id, "username": username.strip(), "password_hash": hash_password(password)},
    )
    session.execute(
        text(
            "INSERT INTO public.user_roles (user_id, role_id, assigned_by) "
            "SELECT CAST(:id AS uuid), id, CAST(:id AS uuid) "
            "FROM public.roles WHERE name = 'admin'"
        ),
        {"id": user_id},
    )
    return user_id


def main() -> None:
    parser = argparse.ArgumentParser(description="Create the first DataOps administrator")
    parser.add_argument("--username", default=os.environ.get("BOOTSTRAP_ADMIN_USERNAME"))
    parser.add_argument("--password", default=os.environ.get("BOOTSTRAP_ADMIN_PASSWORD"))
    parser.add_argument("--if-missing", action="store_true")
    args = parser.parse_args()
    username = args.username or input("Administrator username: ").strip()
    password = args.password or getpass.getpass("Administrator password: ")

    app = create_app()
    with app.app_context():
        try:
            bootstrap_first_admin(db.session, username=username, password=password)
            db.session.commit()
            log_bootstrap_result(username, created=True)
        except RuntimeError:
            db.session.rollback()
            if not args.if_missing:
                raise
            log_bootstrap_result(username, created=False)


if __name__ == "__main__":
    main()