#!/usr/bin/env python3
import argparse
import configparser
import os
import subprocess
import sys
from datetime import datetime

import redfish
import urllib3

try:
    from cryptography.fernet import Fernet, InvalidToken
except ImportError:
    sys.exit("Error: 'cryptography' package is required.  Run: pip install cryptography")

DEFAULT_CREDENTIALS_FILE = "/root/cron_serverlists/credentials.ini"
DEFAULT_KEYFILE_NAME = "/root/cron_serverlists/secret.key"
DEFAULT_SERVER_LIST_FILE = "/root/cron_serverlists/serverlist.ini"

# Examples:
#   ./multi_powercycle.py on
#   ./multi_powercycle.py on --keyfile /root/cron_serverlists/secret.key
#   ./multi_powercycle.py off -f /root/cron_serverlists/serverlist.ini --keyfile /root/cron_serverlists/secret.key

#   ./multi_powercycle.py on
#   ./multi_powercycle.py off
#   ./multi_powercycle.py on -f server_groups.txt
#
# File format:
#   [ipmi]
#   192.168.129.240
#   192.168.129.241
#
#   [ilo]
#   192.168.129.242

DEFAULT_EK = "0000000000000000000000000000000000000000"


def load_fernet(keyfile: str):
    """Load Fernet instance from the key file created by set_password.py."""
    if not os.path.exists(keyfile):
        log(f"Error: key file not found: {keyfile}", is_error=True)
        sys.exit(1)

    with open(keyfile, "rb") as fh:
        key = fh.read()

    try:
        return Fernet(key)
    except ValueError:
        log(f"Error: invalid key format in key file: {keyfile}", is_error=True)
        sys.exit(1)


def resolve_keyfile(cli_keyfile, ipmi_credentials_path, ilo_credentials_path):
    """Resolve keyfile path. Prefer explicit --keyfile, else search next to credentials files."""
    if cli_keyfile:
        return cli_keyfile if os.path.exists(cli_keyfile) else None

    candidates = []
    for cred_path in (ipmi_credentials_path, ilo_credentials_path):
        key_candidate = os.path.join(os.path.dirname(cred_path), DEFAULT_KEYFILE_NAME)
        if key_candidate not in candidates:
            candidates.append(key_candidate)

    # Backward-compatible fallback
    legacy_default = os.path.join(os.path.dirname(DEFAULT_CREDENTIALS_FILE), DEFAULT_KEYFILE_NAME)
    if legacy_default not in candidates:
        candidates.append(legacy_default)

    for candidate in candidates:
        if os.path.exists(candidate):
            return candidate

    return None


def resolve_credentials_path(cli_path, server_file_path):
    """Resolve credentials path. Prefer explicit CLI path, then search common locations."""
    if cli_path and os.path.exists(cli_path):
        return cli_path

    candidates = []

    if cli_path:
        candidates.append(cli_path)

    # Credentials next to server list file is a common deployment pattern
    candidates.append(os.path.join(os.path.dirname(server_file_path), "credentials.ini"))

    # Current working directory fallback
    candidates.append(os.path.join(os.getcwd(), "credentials.ini"))

    # Legacy default fallback
    if DEFAULT_CREDENTIALS_FILE not in candidates:
        candidates.append(DEFAULT_CREDENTIALS_FILE)

    for candidate in candidates:
        if os.path.exists(candidate):
            return candidate

    log(
        "Error: credentials file not found. Looked in: " + ", ".join(candidates),
        is_error=True,
    )
    sys.exit(1)


def is_fernet_token(value: str) -> bool:
    """Best-effort check for Fernet token format from set_password.py."""
    return isinstance(value, str) and value.startswith("gAAAA")


def decrypt_password(encrypted_password: str, fernet):
    """Decrypt password encrypted by set_password.py, or pass through plaintext."""
    if not is_fernet_token(encrypted_password):
        return encrypted_password

    if fernet is None:
        log(
            "Error: encrypted password detected but key file was not found. "
            "Use --keyfile /path/to/secret.key.",
            is_error=True,
        )
        sys.exit(1)

    try:
        return fernet.decrypt(encrypted_password.encode()).decode()
    except InvalidToken:
        log("Error: failed to decrypt password. Check that --keyfile matches the credentials file.", is_error=True)
        sys.exit(1)


def log(message, is_error=False):
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    stream = sys.stderr if is_error else sys.stdout
    print(f"[{timestamp}] {message}", file=stream)


def parse_server_groups(path):
    if not os.path.exists(path):
        log(f"Error: server list file not found: {path}", is_error=True)
        sys.exit(1)

    groups = {"ipmi": [], "ilo": []}
    current_group = None
    aliases = {
        "ipmi": "ipmi",
        "dell": "ipmi",
        "ilo": "ilo",
    }

    try:
        with open(path, "r", encoding="utf-8") as handle:
            for raw_line in handle:
                line = raw_line.strip()
                if not line or line.startswith("#"):
                    continue

                if line.startswith("[") and line.endswith("]"):
                    section_name = line[1:-1].strip().lower()
                    current_group = aliases.get(section_name)
                    if current_group is None:
                        log(f"Error: unsupported server group [{section_name}] in {path}", is_error=True)
                        sys.exit(1)
                    continue

                if current_group is None:
                    log(f"Error: host entry found before any section header in {path}: {line}", is_error=True)
                    sys.exit(1)

                groups[current_group].append(line)
    except OSError as exc:
        log(f"Error reading server list file: {exc}", is_error=True)
        sys.exit(1)

    return groups


def load_ipmi_credentials(path, fernet):
    if not os.path.exists(path):
        log(f"Error: IPMI credentials file not found: {path}", is_error=True)
        sys.exit(1)

    config = configparser.ConfigParser()
    config.read(path)

    section = "ipmi_credentials"
    if not config.has_section(section):
        log(f"Error reading IPMI credentials file: missing [{section}] section in {path}", is_error=True)
        sys.exit(1)

    username = config.get(section, "username", fallback="root")
    encrypted_password = config.get(section, "password", fallback=None)
    encryption_key = config.get(
        section,
        "ipmiek",
        fallback=config.get(section, "encryption_key", fallback=DEFAULT_EK),
    )

    if not encrypted_password:
        log(f"Error reading IPMI credentials file: missing password in [{section}] section of {path}", is_error=True)
        sys.exit(1)

    password = decrypt_password(encrypted_password, fernet)
    return username, password, encryption_key


def load_ilo_credentials(path, fernet):
    if not os.path.exists(path):
        log(f"Error: iLO credentials file not found: {path}", is_error=True)
        sys.exit(1)

    config = configparser.ConfigParser()
    config.read(path)

    section = "ilo_credentials"
    if not config.has_section(section):
        log(f"Error reading iLO credentials file: missing [{section}] section in {path}", is_error=True)
        sys.exit(1)

    username = config.get(section, "username", fallback="root")
    encrypted_password = config.get(section, "password", fallback=None)
    if not encrypted_password:
        log(f"Error reading iLO credentials file: missing password in [{section}] section of {path}", is_error=True)
        sys.exit(1)

    password = decrypt_password(encrypted_password, fernet)
    return username, password


def ipmi_power(host, action, username, password, encryption_key):
    cmd = [
        "ipmitool",
        "-I",
        "lanplus",
        "-H",
        host,
        "-U",
        username,
        "-P",
        password,
        "-y",
        encryption_key,
        "power",
        action,
    ]

    log(f"Sending IPMI power {action} to {host}")
    proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    if proc.returncode != 0:
        detail = proc.stderr.strip() or proc.stdout.strip() or "unknown error"
        log(f"IPMI {action} failed for {host}: {detail}", is_error=True)
        return False

    detail = proc.stdout.strip() or f"power {action} request accepted"
    log(f"IPMI {action} succeeded for {host}: {detail}")
    return True


def ilo_powercycle(host, username, password):
    base_url = f"https://{host}"
    client = redfish.RedfishClient(base_url=base_url, username=username, password=password)

    try:
        client.login(auth="session")
        sys_response = client.get("/redfish/v1/Systems/1/")
        reset_uri = sys_response.dict["Actions"]["#ComputerSystem.Reset"]["target"]
        response = client.post(reset_uri, body={"ResetType": "PushPowerButton"})

        log(f"iLO powercycle status for {host}: {response.status}")
        if response.status >= 400:
            log(f"iLO powercycle failed for {host}: {response.read}", is_error=True)
            return False
        return True
    except Exception as exc:
        log(f"iLO powercycle failed for {host}: {exc}", is_error=True)
        return False
    finally:
        try:
            client.logout()
        except Exception:
            pass


def main():
    urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

    parser = argparse.ArgumentParser(
        description="Dispatch power operations to IPMI hosts and iLO hosts from a grouped server list"
    )
    parser.add_argument("action", choices=("on", "off"), help="Power action for IPMI hosts")
    parser.add_argument(
        "-f",
        "--file",
        dest="server_file",
        default=DEFAULT_SERVER_LIST_FILE,
        help=f"Path to server list file (default: {DEFAULT_SERVER_LIST_FILE})",
    )
    parser.add_argument(
        "--ipmi-credentials",
        default=DEFAULT_CREDENTIALS_FILE,
        help=f"Path to IPMI credentials file (default: {DEFAULT_CREDENTIALS_FILE})",
    )
    parser.add_argument(
        "--ilo-credentials",
        default=DEFAULT_CREDENTIALS_FILE,
        help=f"Path to iLO credentials file (default: {DEFAULT_CREDENTIALS_FILE})",
    )
    parser.add_argument(
        "--keyfile",
        default=None,
        help="Path to Fernet key file created by set_password.py (default: auto-detect next to credentials file)",
    )
    args = parser.parse_args()

    args.ipmi_credentials = resolve_credentials_path(args.ipmi_credentials, args.server_file)
    args.ilo_credentials = resolve_credentials_path(args.ilo_credentials, args.server_file)

    keyfile = resolve_keyfile(args.keyfile, args.ipmi_credentials, args.ilo_credentials)
    fernet = load_fernet(keyfile) if keyfile else None

    groups = parse_server_groups(args.server_file)
    total_hosts = len(groups["ipmi"]) + len(groups["ilo"])
    if total_hosts == 0:
        log(f"Error: no servers found in {args.server_file}", is_error=True)
        sys.exit(1)

    log(
        f"Starting run for {total_hosts} host(s): {len(groups['ipmi'])} IPMI host(s), {len(groups['ilo'])} iLO host(s)"
    )

    ipmi_creds = None
    ilo_creds = None
    failures = 0

    if groups["ipmi"]:
        ipmi_creds = load_ipmi_credentials(args.ipmi_credentials, fernet)
        for host in groups["ipmi"]:
            if not ipmi_power(host, args.action, *ipmi_creds):
                failures += 1

    if groups["ilo"]:
        ilo_creds = load_ilo_credentials(args.ilo_credentials, fernet)
        for host in groups["ilo"]:
            if not ilo_powercycle(host, *ilo_creds):
                failures += 1

    successes = total_hosts - failures
    log(f"Completed run: {successes} successful, {failures} failed")
    if failures:
        sys.exit(1)


if __name__ == "__main__":
    main()