#!/usr/bin/python3
"""Trigger a DB credential rotation via the csl-setup IPC socket and wait.

Used by csl-credential-rotation.service (timer-driven). Sends
``rotateDatabaseCredentials`` then polls ``getRotationStatus`` until the
rotation finishes, exiting non-zero if it failed so systemd records the failure.
"""
import argparse
import json
import os
import socket
import sys
import time

SOCKET_PATH = os.environ.get("CSL_SOCKET", "/run/csl-setup/csl.sock")
POLL_INTERVAL = 5
MAX_WAIT = int(os.environ.get("CSL_ROTATION_CLIENT_TIMEOUT", "1800"))  # 30 min


def _send(request, timeout=60):
    client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    client.settimeout(timeout)
    try:
        client.connect(SOCKET_PATH)
        client.sendall(json.dumps(request).encode())
        chunks = []
        while True:
            chunk = client.recv(65536)
            if not chunk:
                break
            chunks.append(chunk)
    finally:
        client.close()
    raw = b"".join(chunks)
    return json.loads(raw.decode()) if raw else {}


def main():
    parser = argparse.ArgumentParser(description="Rotate CSL DB credentials via the host agent.")
    parser.add_argument("--targets", default="mysql_user,mongo,mongo_scan,mongo_autocrypt",
                        help="Comma-separated rotation targets "
                             "(mysql_user,mongo,mongo_scan,mongo_autocrypt,mysql_root).")
    args = parser.parse_args()
    targets = [t.strip() for t in args.targets.split(",") if t.strip()]

    started = _send({"action": "rotateDatabaseCredentials", "params": {"targets": targets}})
    if not started.get("success"):
        print(f"Failed to start rotation: {started.get('error', started)}", file=sys.stderr)
        return 1
    print(f"Rotation started (targets={targets}).")

    deadline = time.time() + MAX_WAIT
    while time.time() < deadline:
        time.sleep(POLL_INTERVAL)
        try:
            status = _send({"action": "getRotationStatus"})
        except Exception as exc:  # transient socket issue; keep polling
            print(f"status poll error (retrying): {exc}", file=sys.stderr)
            continue
        if status.get("running"):
            continue
        result = status.get("result") or {}
        if result.get("success"):
            print(f"Rotation completed: rotation_id={result.get('rotation_id')}")
            return 0
        print(f"Rotation failed: {result.get('error', 'unknown error')}", file=sys.stderr)
        return 1

    print("Timed out waiting for rotation to complete.", file=sys.stderr)
    return 1


if __name__ == "__main__":
    sys.exit(main())
