#!/usr/bin/env python3
"""
LUKS Container Manager for Hermes Secret Vault

Creates, opens, closes, and manages a LUKS2-encrypted container file
that houses the Vaultwarden data directory.

Usage:
    from luks_manager import LuksContainer

    container = LuksContainer("/opt/rook-vault.img", "/mnt/rook-vault")
    container.create(passphrase, size_mb=500)
    container.open(passphrase)
    # ... Vaultwarden runs inside /mnt/rook-vault ...
    container.close()

Requires: cryptsetup, sudo access
"""

import logging
import os
import subprocess
import sys
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)

DEFAULT_CONTAINER_PATH = "/opt/rook-vault.img"
DEFAULT_MOUNT_POINT = "/mnt/rook-vault"
DEFAULT_MAPPER_NAME = "rook-vault"
DEFAULT_SIZE_MB = 500


class LuksError(Exception):
    """LUKS operation failed."""

    pass


def _run(cmd: list[str], input_data: Optional[str] = None, check: bool = True) -> subprocess.CompletedProcess:
    """Run a command, optionally with stdin input."""
    logger.debug("Running: %s", " ".join(cmd))
    result = subprocess.run(
        cmd,
        input=input_data.encode() if input_data else None,
        capture_output=True,
        text=True,
    )
    if check and result.returncode != 0:
        logger.error("Command failed: %s\nstderr: %s", " ".join(cmd), result.stderr)
        raise LuksError(f"Command failed: {' '.join(cmd)}: {result.stderr.strip()}")
    return result


class LuksContainer:
    """Manages a LUKS2-encrypted file container."""

    def __init__(
        self,
        container_path: str = DEFAULT_CONTAINER_PATH,
        mount_point: str = DEFAULT_MOUNT_POINT,
        mapper_name: str = DEFAULT_MAPPER_NAME,
    ):
        self.container_path = Path(container_path)
        self.mount_point = Path(mount_point)
        self.mapper_name = mapper_name
        self.mapper_path = Path(f"/dev/mapper/{mapper_name}")

    @property
    def exists(self) -> bool:
        """Check if the container file exists."""
        return self.container_path.exists()

    @property
    def is_open(self) -> bool:
        """Check if the LUKS container is currently open (mapped)."""
        return self.mapper_path.exists()

    @property
    def is_mounted(self) -> bool:
        """Check if the container filesystem is mounted."""
        if not self.mount_point.exists():
            return False
        result = subprocess.run(
            ["mountpoint", "-q", str(self.mount_point)],
            capture_output=True,
        )
        return result.returncode == 0

    def create(self, passphrase: str, size_mb: int = DEFAULT_SIZE_MB) -> None:
        """
        Create a new LUKS2-encrypted container.

        1. Allocate a file of the specified size
        2. Format as LUKS2 with the passphrase
        3. Open the container
        4. Create ext4 filesystem inside
        5. Set up mount point with correct permissions
        6. Close the container (caller re-opens when needed)
        """
        if self.exists:
            raise LuksError(f"Container already exists at {self.container_path}")

        logger.info("Creating %dMB LUKS2 container at %s", size_mb, self.container_path)

        # Step 1: Allocate file
        _run(["sudo", "fallocate", "-l", f"{size_mb}M", str(self.container_path)])
        _run(["sudo", "chmod", "600", str(self.container_path)])

        # Step 2: Format as LUKS2
        # --batch-mode skips confirmation, --key-file=- reads passphrase from stdin
        _run(
            [
                "sudo",
                "cryptsetup",
                "luksFormat",
                "--type",
                "luks2",
                "--cipher",
                "aes-xts-plain64",
                "--key-size",
                "512",
                "--hash",
                "sha256",
                "--pbkdf",
                "argon2id",
                "--pbkdf-memory",
                "262144",  # 256MB
                "--pbkdf-time",
                "4000",  # 4 seconds
                "--batch-mode",
                "--key-file=-",
                str(self.container_path),
            ],
            input_data=passphrase,
        )
        logger.info("LUKS2 format complete")

        # Step 3: Open
        _run(
            [
                "sudo",
                "cryptsetup",
                "open",
                "--type",
                "luks2",
                "--key-file=-",
                str(self.container_path),
                self.mapper_name,
            ],
            input_data=passphrase,
        )

        # Step 4: Create ext4 filesystem
        _run(["sudo", "mkfs.ext4", "-q", "-L", "rook-vault", str(self.mapper_path)])

        # Step 5: Mount and set permissions
        self.mount_point.parent.mkdir(parents=True, exist_ok=True)
        _run(["sudo", "mkdir", "-p", str(self.mount_point)])
        _run(["sudo", "mount", str(self.mapper_path), str(self.mount_point)])

        # Set ownership to current user
        user = os.environ.get("USER", "howsa")
        _run(["sudo", "chown", "-R", f"{user}:{user}", str(self.mount_point)])

        # Create Vaultwarden data directory structure
        vw_data = self.mount_point / "vaultwarden"
        vw_data.mkdir(exist_ok=True)
        (vw_data / "data").mkdir(exist_ok=True)

        # Step 6: Close (caller re-opens when ready)
        self.close()

        logger.info("LUKS container created successfully at %s", self.container_path)

    def open(self, passphrase: str) -> None:
        """Open the LUKS container and mount it."""
        if not self.exists:
            raise LuksError(f"Container not found at {self.container_path}")
        if self.is_open:
            logger.info("Container already open")
            if not self.is_mounted:
                _run(["sudo", "mount", str(self.mapper_path), str(self.mount_point)])
            return

        _run(
            [
                "sudo",
                "cryptsetup",
                "open",
                "--type",
                "luks2",
                "--key-file=-",
                str(self.container_path),
                self.mapper_name,
            ],
            input_data=passphrase,
        )

        _run(["sudo", "mkdir", "-p", str(self.mount_point)])
        _run(["sudo", "mount", str(self.mapper_path), str(self.mount_point)])

        user = os.environ.get("USER", "howsa")
        _run(["sudo", "chown", "-R", f"{user}:{user}", str(self.mount_point)])

        logger.info("Container opened and mounted at %s", self.mount_point)

    def close(self) -> None:
        """Unmount and close the LUKS container."""
        if self.is_mounted:
            _run(["sudo", "umount", str(self.mount_point)], check=False)

        if self.is_open:
            _run(["sudo", "cryptsetup", "close", self.mapper_name], check=False)

        logger.info("Container closed")

    def status(self) -> dict:
        """Return container status information."""
        info = {
            "container_path": str(self.container_path),
            "mount_point": str(self.mount_point),
            "exists": self.exists,
            "is_open": self.is_open,
            "is_mounted": self.is_mounted,
        }

        if self.exists:
            info["size_mb"] = self.container_path.stat().st_size // (1024 * 1024)

        if self.is_open:
            result = _run(
                ["sudo", "cryptsetup", "status", self.mapper_name],
                check=False,
            )
            info["luks_status"] = result.stdout.strip()

        if self.is_mounted:
            result = _run(["df", "-h", str(self.mount_point)], check=False)
            lines = result.stdout.strip().split("\n")
            if len(lines) >= 2:
                info["disk_usage"] = lines[1]

        return info

    def change_passphrase(self, old_passphrase: str, new_passphrase: str) -> None:
        """Change the LUKS passphrase."""
        if not self.exists:
            raise LuksError("Container does not exist")

        # Add new key
        _run(
            [
                "sudo",
                "cryptsetup",
                "luksAddKey",
                "--key-file=-",
                str(self.container_path),
            ],
            input_data=f"{old_passphrase}\n{new_passphrase}",
        )

        # Remove old key
        _run(
            [
                "sudo",
                "cryptsetup",
                "luksRemoveKey",
                "--key-file=-",
                str(self.container_path),
            ],
            input_data=old_passphrase,
        )

        logger.info("Passphrase changed successfully")

    def destroy(self) -> None:
        """Permanently destroy the container. IRREVERSIBLE."""
        self.close()
        if self.exists:
            # Overwrite with random data before deletion
            size = self.container_path.stat().st_size
            _run(["sudo", "shred", "-n", "3", "-z", str(self.container_path)])
            _run(["sudo", "rm", "-f", str(self.container_path)])
            logger.warning(
                "Container destroyed: %s (%d MB)",
                self.container_path,
                size // (1024 * 1024),
            )


# =========================================================================
# Systemd integration
# =========================================================================

SYSTEMD_UNIT = """[Unit]
Description=Rook Vault (LUKS container)
Before=rook-agent.service rook-nerve.service
After=local-fs.target

[Service]
Type=oneshot
RemainAfterExit=yes
ExecStart=/usr/bin/python3 {script} open
ExecStop=/usr/bin/python3 {script} close
TimeoutStartSec=30

[Install]
WantedBy=multi-user.target
"""


def install_systemd_unit(script_path: str) -> str:
    """Generate a systemd unit for auto-mounting the vault at boot."""
    return SYSTEMD_UNIT.format(script=script_path)


# =========================================================================
# CLI
# =========================================================================


def main():
    import argparse

    logging.basicConfig(level=logging.INFO, format="%(message)s")

    parser = argparse.ArgumentParser(description="Rook LUKS Vault Manager")
    parser.add_argument("--container", default=DEFAULT_CONTAINER_PATH, help="Container file path")
    parser.add_argument("--mount", default=DEFAULT_MOUNT_POINT, help="Mount point")

    sub = parser.add_subparsers(dest="command")
    sub.add_parser("status", help="Show container status")

    create_cmd = sub.add_parser("create", help="Create new container")
    create_cmd.add_argument("--size", type=int, default=DEFAULT_SIZE_MB, help="Size in MB")

    sub.add_parser("open", help="Open and mount container")
    sub.add_parser("close", help="Unmount and close container")
    sub.add_parser("destroy", help="Permanently destroy container")

    args = parser.parse_args()

    container = LuksContainer(args.container, args.mount)

    if args.command == "status":
        import json

        print(json.dumps(container.status(), indent=2))

    elif args.command == "create":
        import getpass

        pw = getpass.getpass("Enter passphrase for new vault: ")
        pw2 = getpass.getpass("Confirm passphrase: ")
        if pw != pw2:
            print("Passphrases don't match")
            sys.exit(1)
        container.create(pw, args.size)
        print(f"Vault created at {args.container} ({args.size}MB)")

    elif args.command == "open":
        import getpass

        pw = getpass.getpass("Vault passphrase: ")
        container.open(pw)
        print(f"Vault mounted at {args.mount}")

    elif args.command == "close":
        container.close()
        print("Vault closed")

    elif args.command == "destroy":
        confirm = input("This will PERMANENTLY DESTROY the vault. Type 'yes' to confirm: ")
        if confirm == "yes":
            container.destroy()
            print("Vault destroyed")
        else:
            print("Aborted")

    else:
        parser.print_help()


if __name__ == "__main__":
    main()