# -*- coding: utf-8 -*-
"""
Shared utilities for Control-M SaaS migration export/import scripts.

Provides constants, encryption/decryption, key handling, component registry,
subprocess helpers, archive operations, and password management.

Envelope format (binary .enc):
  - 8 bytes: magic CTM_MIG1
  - 1 byte:  flag 0 = AES key from global file only; 1 = PBKDF2(global, password)
  - 12 bytes: GCM nonce
  - remainder: AES-256-GCM ciphertext followed by 16-byte auth tag
"""

import getpass
import hashlib
import logging
import os
import platform
import shutil
import subprocess
import sys
import tarfile

if platform.system() == "Windows":
    import winreg

try:
    from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
except ImportError:
    Cipher = None
    algorithms = None
    modes = None

logger = logging.getLogger("ctmsaas")

MAGIC = b"CTM_MIG1"
FLAG_KEY_GLOBAL = 0
FLAG_KEY_PBKDF2 = 1
PBKDF2_ITERATIONS = 100000
NONCE_SIZE = 12
GLOBAL_KEY_BYTES = 32
ENCRYPT_CHUNK_SIZE = 4 * 1024 * 1024

SUCCESS = 0
FAILURE = 1

MAX_PASSWORD_ATTEMPTS = 3
ALLOWED_PW_CHARS = set(
    "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789!@#$%^&*()-_=+"
)
ALLOWED_PW_DISPLAY = "A-Z, a-z, 0-9, !@#$%^&*()-_=+"
FINAL_ARCHIVE_NAME = "CTMSExportForSaaS.tar.gz"
SUBPROCESS_TIMEOUT = 600


# ---------------------------------------------------------------------------
# Logging and paths
# ---------------------------------------------------------------------------

def setup_logging(log_filename):
    """Configure logging to a file only.

    log_filename is just the basename (e.g. 'ctmsaasexport.log').
    The file is placed in the same directory as ctmsaascommon.py.

    Console output is intentionally limited to explicit print() calls so
    that calling scripts see only clean step/result lines on stdout.
    """
    log_dir = os.path.dirname(os.path.realpath(__file__))
    log_path = os.path.join(log_dir, log_filename)
    fmt = "%(asctime)s %(name)s %(levelname)s [%(funcName)s] %(message)s"

    file_handler = logging.FileHandler(log_path, encoding="utf-8")
    file_handler.setLevel(logging.INFO)
    file_handler.setFormatter(logging.Formatter(fmt))

    root = logging.getLogger()
    root.setLevel(logging.INFO)
    root.addHandler(file_handler)

    logger.info("Log file: %s", log_path)


def _read_registry(reg_path, value_name):
    """Read a string value from HKLM; return None if the key or value is missing."""
    try:
        with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, reg_path, 0, winreg.KEY_READ) as key:
            value, _ = winreg.QueryValueEx(key, value_name)
            return value
    except WindowsError:
        return None


def get_server_base_dir():
    """Return the Control-M server root directory.

    Resolution order:
      1. CONTROLM_SERVER environment variable (explicit override)
      2. On Windows: HKLM registry value
         SOFTWARE\\BMC Software\\Control-M/Server\\SYSPRM\\CONTROLM_SERVER
      3. One level above this script's real path (standard CTM install layout)
    """
    env_val = os.getenv("CONTROLM_SERVER")
    if env_val:
        return env_val

    if platform.system() == "Windows":
        reg_path = "SOFTWARE\\BMC Software\\Control-M" + chr(0x2F) + "Server\\SYSPRM"
        server_base = _read_registry(reg_path, "CONTROLM_SERVER")
        if server_base:
            return server_base

    return os.path.join(os.path.dirname(os.path.realpath(__file__)), "..")


def make_work_dir(server_base):
    """Create and return a working directory for intermediate files.

    The directory is <server_base>/tmp/ctmsaasmigration.  It is created if it
    does not already exist.
    """
    work_dir = os.path.join(server_base, "tmp", "ctmsaasmigration")
    os.makedirs(work_dir, exist_ok=True)
    logger.info("Work directory: %s", work_dir)
    return work_dir


# ---------------------------------------------------------------------------
# Progress tracker
# ---------------------------------------------------------------------------

class StepTracker:
    def __init__(self, total):
        self._current = 0
        self._total = total

    def step(self, msg):
        self._current += 1
        text = "Step %d of %d - %s" % (self._current, self._total, msg)
        print(text)
        logger.info(text)


# ---------------------------------------------------------------------------
# Key handling
# ---------------------------------------------------------------------------

def global_key_path(server_base):
    return os.path.join(server_base, "data", "ctm_key.txt")


def load_global_key_bytes(server_base):
    path = global_key_path(server_base)
    if not os.path.isfile(path):
        raise IOError("Global key file not found: %s" % path)
    with open(path, "rb") as f:
        raw = f.read()
    if len(raw) < GLOBAL_KEY_BYTES:
        raise ValueError(
            "Global key file too short (need %d bytes): %s" % (GLOBAL_KEY_BYTES, path)
        )
    return raw[:GLOBAL_KEY_BYTES]


def derive_aes_key(global_key_bytes, use_password, password_utf8):
    """Derive the AES-256 key: raw global key bytes, or PBKDF2(password, salt=SHA256(global)[:16])."""
    if not use_password:
        return global_key_bytes
    salt = hashlib.sha256(global_key_bytes).digest()[:16]
    return hashlib.pbkdf2_hmac(
        "sha256",
        password_utf8,
        salt,
        PBKDF2_ITERATIONS,
        dklen=GLOBAL_KEY_BYTES,
    )


# ---------------------------------------------------------------------------
# Encryption / decryption
# ---------------------------------------------------------------------------

def encrypt_file_aes_gcm(plaintext_path, output_enc_path, aes_key, flag_byte):
    """Write binary .enc: MAGIC + flag + nonce + AES-GCM ciphertext + auth tag."""
    nonce = os.urandom(NONCE_SIZE)
    encryptor = Cipher(algorithms.AES(aes_key), modes.GCM(nonce)).encryptor()
    try:
        with open(plaintext_path, "rb") as src, open(output_enc_path, "wb") as out:
            out.write(MAGIC)
            out.write(bytes([flag_byte]))
            out.write(nonce)
            while True:
                chunk = src.read(ENCRYPT_CHUNK_SIZE)
                if not chunk:
                    break
                out.write(encryptor.update(chunk))
            out.write(encryptor.finalize())
            out.write(encryptor.tag)
    except Exception:
        if os.path.isfile(output_enc_path):
            try:
                os.remove(output_enc_path)
            except Exception as cleanup_err:
                logger.warning(
                    "Failed to remove partial encrypted file %s: %s",
                    output_enc_path, cleanup_err,
                )
        raise


def peek_enc_flag(enc_path):
    """Read the key-derivation flag byte from an .enc file header."""
    with open(enc_path, "rb") as f:
        magic = f.read(len(MAGIC))
        if magic != MAGIC:
            raise ValueError("Bad magic header in %s" % enc_path)
        return f.read(1)[0]


def decrypt_file_aes_gcm(enc_path, output_path, aes_key):
    """Read binary .enc envelope and decrypt to output_path."""
    with open(enc_path, "rb") as f:
        magic = f.read(len(MAGIC))
        if magic != MAGIC:
            raise ValueError("Bad magic header in %s" % enc_path)
        f.read(1)  # flag byte — key was already derived by caller
        nonce = f.read(NONCE_SIZE)

        header_size = len(MAGIC) + 1 + NONCE_SIZE
        f.seek(0, 2)
        file_size = f.tell()
        ciphertext_size = file_size - header_size - 16
        if ciphertext_size < 0:
            raise ValueError("Encrypted file too small: %s" % enc_path)

        f.seek(-16, 2)
        tag = f.read(16)
        f.seek(header_size)

        decryptor = Cipher(algorithms.AES(aes_key), modes.GCM(nonce, tag)).decryptor()
        try:
            with open(output_path, "wb") as out:
                remaining = ciphertext_size
                while remaining > 0:
                    to_read = min(ENCRYPT_CHUNK_SIZE, remaining)
                    chunk = f.read(to_read)
                    if not chunk:
                        break
                    out.write(decryptor.update(chunk))
                    remaining -= len(chunk)
                out.write(decryptor.finalize())
        except Exception:
            if os.path.isfile(output_path):
                try:
                    os.remove(output_path)
                except Exception as cleanup_err:
                    logger.warning(
                        "Failed to remove partial decrypted file %s: %s",
                        output_path, cleanup_err,
                    )
            raise


# ---------------------------------------------------------------------------
# Subprocess
# ---------------------------------------------------------------------------

def run_cmd(cmd, timeout=SUBPROCESS_TIMEOUT):
    """Run a subprocess and return (stdout, stderr) as bytes.

    Raises RuntimeError on non-zero exit or timeout.
    """
    logger.info("Running: %s", " ".join(cmd))
    try:
        result = subprocess.run(cmd, capture_output=True, timeout=timeout)
    except subprocess.TimeoutExpired:
        raise RuntimeError("%s timed out after %d seconds" % (cmd[0], timeout))
    if result.returncode != 0:
        stderr_text = result.stderr.decode("utf-8", "replace") if result.stderr else ""
        stdout_text = result.stdout.decode("utf-8", "replace") if result.stdout else ""
        raise RuntimeError(
            "%s failed (exit %s). stderr: %s stdout: %s"
            % (cmd[0], result.returncode, stderr_text, stdout_text)
        )
    return result.stdout, result.stderr


# ---------------------------------------------------------------------------
# Component registry
# ---------------------------------------------------------------------------

def find_ctmsetown(server_base):
    """Resolve ctmsetown executable: PATH first, then exe/ and bin/ under server_base."""
    names = ["ctmsetown"]
    if os.name == "nt":
        names.insert(0, "ctmsetown.exe")
    for name in names:
        found = shutil.which(name)
        if found and os.path.isfile(found):
            return found
    for subdir in ("exe", "bin"):
        for name in names:
            path = os.path.join(server_base, subdir, name)
            if os.path.isfile(path):
                return path
    return None


def export_setown(server_base, pack_path):
    """Export handler for the setown component."""
    ctmsetown = find_ctmsetown(server_base)
    if not ctmsetown:
        raise RuntimeError(
            "ctmsetown not found (PATH, %s/exe, or %s/bin)" % (server_base, server_base)
        )
    logger.info("ctmsetown found: %s", ctmsetown)
    run_cmd([ctmsetown, "-action", "export", "-filename", pack_path])
    logger.info("ctmsetown export completed")


def import_setown(server_base, pack_path):
    """Import handler for the setown component."""
    ctmsetown = find_ctmsetown(server_base)
    if not ctmsetown:
        raise RuntimeError(
            "ctmsetown not found (PATH, %s/exe, or %s/bin)" % (server_base, server_base)
        )
    run_cmd(["sed", "-i", r"s/\r//", pack_path])
    logger.info("Stripped carriage returns from: %s", pack_path)
    run_cmd([ctmsetown, "-action", "import", "-filename", pack_path, "-data", "truncate"])
    logger.info("ctmsetown import completed")


COMPONENT_REGISTRY = {
    "setown": {
        "pack_file": "migration_export_ctmsetown.txt",
        "export": export_setown,
        "import": import_setown,
    },
}
SUPPORTED_COMPONENTS = list(COMPONENT_REGISTRY)


# ---------------------------------------------------------------------------
# Component file helpers
# ---------------------------------------------------------------------------

def resolve_pack_path(base_dir, component):
    """Return <base_dir>/<pack_filename> for the given component."""
    info = COMPONENT_REGISTRY.get(component)
    filename = info["pack_file"] if info else "migration_export_%s.txt" % component
    return os.path.join(base_dir, filename)


def file_to_component(basename):
    """Map a pack filename (.enc or .txt) to (component_name, is_encrypted).

    Returns (None, None) if the filename is not recognized.
    """
    for comp, info in COMPONENT_REGISTRY.items():
        pack_name = info["pack_file"]
        if basename == pack_name + ".enc":
            return comp, True
        if basename == pack_name:
            return comp, False
    return None, None


# ---------------------------------------------------------------------------
# Password
# ---------------------------------------------------------------------------

def validate_password(pw):
    """Return None if valid, or an error message string."""
    if not pw:
        return "Password cannot be empty."
    bad_chars = set(pw) - ALLOWED_PW_CHARS
    if bad_chars:
        return "Password contains disallowed characters: %s" % " ".join(sorted(bad_chars))
    return None


def prompt_password(confirm=True):
    """Prompt for password with optional confirmation. Allows up to MAX_PASSWORD_ATTEMPTS tries."""
    for attempt in range(1, MAX_PASSWORD_ATTEMPTS + 1):
        pw1 = getpass.getpass("Password (allowed chars: %s): " % ALLOWED_PW_DISPLAY)
        error = validate_password(pw1)
        if error:
            print("ERROR: %s" % error, file=sys.stderr)
            logger.warning("Password attempt %d rejected: %s", attempt, error)
            continue
        if confirm:
            pw2 = getpass.getpass("Confirm password: ")
            if pw1 != pw2:
                print("ERROR: Passwords do not match. Please try again.", file=sys.stderr)
                logger.warning("Password attempt %d: mismatch", attempt)
                continue
        return pw1.encode("utf-8")
    return None


# ---------------------------------------------------------------------------
# Archive
# ---------------------------------------------------------------------------

def compress_to_targz(source_paths, archive_path):
    """Compress one or more files into a single .tar.gz archive."""
    logger.info("Compressing %d file(s) -> %s", len(source_paths), archive_path)
    with tarfile.open(archive_path, "w:gz") as tar:
        for src in source_paths:
            arcname = os.path.basename(src)
            tar.add(src, arcname=arcname)
            logger.info("  Added: %s as %s", src, arcname)
    logger.info(
        "Compression complete: %s (%d bytes)",
        archive_path, os.path.getsize(archive_path),
    )


def extract_targz(archive_path, dest_dir):
    """Extract all regular, flat-named files from a .tar.gz archive to dest_dir.

    Members with path components are skipped to prevent path-traversal attacks.
    Returns a list of extracted file paths.
    """
    logger.info("Extracting %s -> %s", archive_path, dest_dir)
    with tarfile.open(archive_path, "r:gz") as tar:
        all_members = tar.getmembers()
        safe_members = [
            m for m in all_members
            if m.isfile() and os.path.basename(m.name) == m.name
        ]
        skipped = len(all_members) - len(safe_members)
        if skipped:
            logger.warning(
                "Skipped %d archive member(s) with path components or non-regular type",
                skipped,
            )
        tar.extractall(dest_dir, members=safe_members)
    extracted = [os.path.join(dest_dir, m.name) for m in safe_members]
    logger.info("Extracted %d file(s)", len(extracted))
    return extracted


# ---------------------------------------------------------------------------
# Cleanup
# ---------------------------------------------------------------------------

def cleanup_work_dir(work_dir):
    """Remove the temporary working directory and all its contents."""
    if work_dir and os.path.isdir(work_dir):
        try:
            shutil.rmtree(work_dir)
            logger.info("Removed work directory: %s", work_dir)
        except Exception as e:
            logger.warning("Failed to remove work directory %s: %s", work_dir, e)


def cleanup_files(paths):
    """Best-effort removal of a list of file paths."""
    for p in paths:
        try:
            if os.path.isfile(p):
                os.remove(p)
                logger.info("Removed: %s", p)
        except Exception as e:
            logger.warning("Failed to remove %s: %s", p, e)
