bootstrap_admin.py 3.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293
  1. from __future__ import annotations
  2. import argparse
  3. import getpass
  4. import logging
  5. import os
  6. import re
  7. from sqlalchemy import text
  8. from app import create_app, db
  9. from app.core.common.identifiers import new_governance_uid
  10. from app.core.system.auth import hash_password
  11. logger = logging.getLogger(__name__)
  12. def validate_password(password: str) -> None:
  13. if len(password) < 10:
  14. raise ValueError("password must contain at least 10 characters")
  15. if not re.search(r"[a-z]", password):
  16. raise ValueError("password must contain a lowercase letter")
  17. if not re.search(r"[A-Z]", password):
  18. raise ValueError("password must contain an uppercase letter")
  19. if not re.search(r"[0-9]", password):
  20. raise ValueError("password must contain a digit")
  21. def log_bootstrap_result(username: str, *, created: bool) -> None:
  22. if created:
  23. logger.info("Created first platform administrator: %s", username)
  24. else:
  25. logger.info("Platform administrator already exists; bootstrap skipped")
  26. def bootstrap_first_admin(session, *, username: str, password: str) -> str:
  27. validate_password(password)
  28. existing = session.execute(
  29. text(
  30. "SELECT COUNT(*) FROM public.users u "
  31. "JOIN public.user_roles ur ON ur.user_id = u.id "
  32. "JOIN public.roles r ON r.id = ur.role_id "
  33. "WHERE r.name = 'admin' AND u.status = 'active'"
  34. )
  35. ).scalar_one()
  36. if existing:
  37. raise RuntimeError("an active administrator already exists")
  38. user_id = new_governance_uid()
  39. session.execute(
  40. text(
  41. "INSERT INTO public.users "
  42. "(id, username, display_name, password_hash) "
  43. "VALUES (CAST(:id AS uuid), :username, :username, :password_hash)"
  44. ),
  45. {"id": user_id, "username": username.strip(), "password_hash": hash_password(password)},
  46. )
  47. session.execute(
  48. text(
  49. "INSERT INTO public.user_roles (user_id, role_id, assigned_by) "
  50. "SELECT CAST(:id AS uuid), id, CAST(:id AS uuid) "
  51. "FROM public.roles WHERE name = 'admin'"
  52. ),
  53. {"id": user_id},
  54. )
  55. return user_id
  56. def main() -> None:
  57. parser = argparse.ArgumentParser(description="Create the first DataOps administrator")
  58. parser.add_argument("--username", default=os.environ.get("BOOTSTRAP_ADMIN_USERNAME"))
  59. parser.add_argument("--password", default=os.environ.get("BOOTSTRAP_ADMIN_PASSWORD"))
  60. parser.add_argument("--if-missing", action="store_true")
  61. args = parser.parse_args()
  62. username = args.username or input("Administrator username: ").strip()
  63. password = args.password or getpass.getpass("Administrator password: ")
  64. app = create_app()
  65. with app.app_context():
  66. try:
  67. bootstrap_first_admin(db.session, username=username, password=password)
  68. db.session.commit()
  69. log_bootstrap_result(username, created=True)
  70. except RuntimeError:
  71. db.session.rollback()
  72. if not args.if_missing:
  73. raise
  74. log_bootstrap_result(username, created=False)
  75. if __name__ == "__main__":
  76. main()