94 lines
3.1 KiB
Python
Executable File
94 lines
3.1 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Rotate the initialized PostgreSQL bootstrap role and verify before returning success."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import psycopg2
|
|
from psycopg2 import sql
|
|
|
|
|
|
def read_secret(path: str) -> str:
|
|
value = Path(path).read_text()
|
|
if not value or "\x00" in value or any(character.isspace() for character in value):
|
|
raise ValueError("secret must be non-empty and contain no whitespace or NUL bytes")
|
|
return value
|
|
|
|
|
|
def connect(password: str):
|
|
return psycopg2.connect(
|
|
host=os.environ.get("THT_VECTOR_HOST", "vector-db"),
|
|
port=int(os.environ.get("THT_VECTOR_PORT", "5432")),
|
|
dbname=os.environ.get("THT_VECTOR_DATABASE", "thoth"),
|
|
user=os.environ.get("THT_VECTOR_BOOTSTRAP_USER", "postgres"),
|
|
password=password,
|
|
connect_timeout=5,
|
|
)
|
|
|
|
|
|
def alter_current_role(connection, password: str) -> None:
|
|
with connection.cursor() as cursor:
|
|
cursor.execute("SELECT current_user")
|
|
current_user = cursor.fetchone()[0]
|
|
expected = os.environ.get("THT_VECTOR_BOOTSTRAP_USER", "postgres")
|
|
if current_user != expected:
|
|
raise RuntimeError("authenticated role does not match THT_VECTOR_BOOTSTRAP_USER")
|
|
cursor.execute(
|
|
sql.SQL("ALTER ROLE {} PASSWORD {}").format(
|
|
sql.Identifier(current_user), sql.Literal(password)
|
|
)
|
|
)
|
|
connection.commit()
|
|
|
|
|
|
def main() -> int:
|
|
if len(sys.argv) != 3:
|
|
print("usage: rotate-bootstrap-password.py OLD_SECRET NEW_SECRET", file=sys.stderr)
|
|
return 2
|
|
try:
|
|
old_password = read_secret(sys.argv[1])
|
|
new_password = read_secret(sys.argv[2])
|
|
if old_password == new_password:
|
|
raise ValueError("old and new bootstrap passwords must differ")
|
|
old_connection = connect(old_password)
|
|
except Exception as exc:
|
|
print(f"bootstrap rotation refused before change: {type(exc).__name__}", file=sys.stderr)
|
|
return 1
|
|
|
|
try:
|
|
alter_current_role(old_connection, new_password)
|
|
try:
|
|
verification = connect(new_password)
|
|
verification.close()
|
|
except Exception as verify_exc:
|
|
try:
|
|
alter_current_role(old_connection, old_password)
|
|
except Exception as restore_exc:
|
|
print(
|
|
"bootstrap rotation verification failed and password restore failed: "
|
|
f"{type(verify_exc).__name__}/{type(restore_exc).__name__}",
|
|
file=sys.stderr,
|
|
)
|
|
return 3
|
|
print(
|
|
f"bootstrap rotation verification failed; old password restored: "
|
|
f"{type(verify_exc).__name__}",
|
|
file=sys.stderr,
|
|
)
|
|
return 1
|
|
except Exception as exc:
|
|
print(f"bootstrap rotation failed: {type(exc).__name__}", file=sys.stderr)
|
|
return 1
|
|
finally:
|
|
old_connection.close()
|
|
|
|
print("bootstrap database password rotated and new login verified")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|