#!/usr/bin/env python3
""" 
Module for handling LUKS encryption and decryption of partitions.

Author: Arjen Balfoort
Date: 2026-02-22

Dependencies:
cryptsetup
dmsetup
e2fsprogs (mkfs.ext4)
xfsprogs (mkfs.xfs)
"""

import sys
import os
import subprocess
import time
from typing import Optional
import psutil


def is_encrypted(target: str) -> bool:
    """
    Check if a partition or mount point is LUKS-encrypted using cryptsetup.

    Args:
        target: Path to the partition (e.g., "/dev/sda1") or mount point (e.g., "/boot").

    Returns:
        bool: True if the partition is LUKS-encrypted, False otherwise.
    """
    # If target is a mount point, resolve to block device
    if os.path.isdir(target):
        try:
            result = subprocess.run(
                ['findmnt', '-n', '-o', 'SOURCE', '--target', target],
                check=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True
            )
            old_target = target
            target = result.stdout.strip()
            print(f"Resolved mount point {old_target} to block device {target}")

            # Check if the device is a mapped device
            if target.startswith("/dev/mapper/"):
                return True

        except subprocess.CalledProcessError:
            return False

    # Check if the block device is LUKS-encrypted
    try:
        subprocess.run(
            ["cryptsetup", "isLuks", target],
            capture_output=True,
            check=True
        )
        return True
    except (subprocess.CalledProcessError, FileNotFoundError):
        return False

def is_mapped(partition: str) -> bool:
    """
    Check if a partition is mapped using cryptsetup.

    Args:
        partition: Path to the partition (e.g., "/dev/sda1").

    Returns:
        bool: True if the partition is mapped, False otherwise.
    """
    try:
        command = ["cryptsetup", "status", partition]
        print(f"Running: {' '.join(command)}")
        subprocess.run(
            command,
            capture_output=True,
            text=True,
            check=True
        )
        # If the command succeeds, the partition is mapped.
        return True
    except subprocess.CalledProcessError:
        # If the command fails, the partition is not mapped.
        return False

def mount_partition(device_path: str, mount_point: str) -> bool:
    """
    Mount a partition.

    Args:
        device_path: Path to the device (e.g., "/dev/sda1").
        mount_point: Path to the mountpoint (e.g., "/mnt/data").

    Returns:
        bool: True if mounted successfully, False otherwise.
    """
    try:
        subprocess.run(["mount", device_path, mount_point], check=True)
        print(f"Mounted {device_path} to {mount_point}")
        return True
    except subprocess.CalledProcessError as e:
        print(f"Failed to mount {device_path}: {e}")
        return False

def unmount_partition(partition: str) -> Optional[str]:
    """
    Unmounts a partition if it is mounted.
    Returns:
        str: The partition path if unmounted successfully, None otherwise.
             In case of a mapped device, returns the underlying device path.
    """
    # Check if the partition is mounted
    result = subprocess.run(["findmnt", "--noheadings", "--output=TARGET", "--source", partition],
                            capture_output=True,
                            text=True,
                            check=False)

    if result.stdout.strip() != "":
        # Try to unmount
        try:
            subprocess.run(["umount", partition], check=True)
            print(f"Unmounted {partition}")

            # Check if the partition is in use by any process
            seconds = 0
            for seconds in range(5):
                if is_partition_in_use(partition):
                    time.sleep(1)
                else:
                    break

            if seconds < 5:
                return partition
            return None
        except subprocess.CalledProcessError:
            print(f"Failed to unmount {partition}")
            return None

    # If the partition is a mapped device: close the mapped device
    if "/dev/mapper" in partition:
        try:
            # Get the underlying device
            command = ["dmsetup", "table", "--target", "crypt", partition]
            print(f"Running: {' '.join(command)}")
            underlying_device = subprocess.run(
                command,
                capture_output=True,
                text=True,
                check=True
            )
            underlying_device = underlying_device.stdout.split()[-1].strip()

            # Close the mapped device
            command = ["cryptsetup", "close", partition]
            print(f"Running: {' '.join(command)}")
            subprocess.run(command, check=True)
            print(f"Closed mapped device {partition}")

            return underlying_device
        except subprocess.CalledProcessError:
            print(f"Failed to close mapped device {partition}")
            return None

    return partition

def get_uuid_from_device(device_path: str) -> Optional[str]:
    """
    Get the UUID from a device path.
    If the device is a mapped/encrypted device (e.g., /dev/mapper/luks-UUID),
    the UUID of the underlying device (e.g., /dev/sda1) is returned.

    Args:
        device_path: Path to the device (e.g., "/dev/sda1" or "/dev/mapper/luks-UUID").

    Returns:
        Optional[str]: UUID of the device or underlying device, or None if not found.
    """
    is_root = os.geteuid() == 0

    # If the device is a mapped/encrypted device, resolve the underlying device
    if "/dev/mapper/" in device_path:
        if is_root:
            try:
                command = ["dmsetup", "table", "--target", "crypt", device_path]
                print(f"Running: {' '.join(command)}")
                result = subprocess.run(
                    command,
                    capture_output=True,
                    text=True,
                    check=True
                )
                underlying_device = result.stdout.split()[-1].strip()
                print(f"Underlying device of {device_path}: {underlying_device}")
                device_path = underlying_device
            except subprocess.CalledProcessError:
                return None
        else:
            return None  # Cannot resolve mapped devices without root

    # Try to get UUID with blkid (requires root for most devices)
    if is_root:
        try:
            command = ["blkid", "-s", "UUID", "-o", "value", device_path]
            print(f"Running: {' '.join(command)}")
            result = subprocess.run(
                command,
                capture_output=True,
                text=True,
                check=True
            )
            uuid = result.stdout.strip()
            print(f"UUID of {device_path}: {uuid}")
            return uuid
        except subprocess.CalledProcessError:
            pass  # Fall through to non-root method

    # Fallback: Check /dev/disk/by-uuid/ (no root needed)
    try:
        for uuid in os.listdir("/dev/disk/by-uuid"):
            uuid_path = os.path.join("/dev/disk/by-uuid", uuid)
            if os.path.realpath(uuid_path) == device_path:
                print(f"UUID of {device_path} by /dev/disk/by-uuid: {uuid}")
                return uuid
    except FileNotFoundError:
        pass

    return None

def get_partition_info(device_path: str) -> tuple[str, str]:
    """
    Get the filesystem type and label of a device or mapped device.

    Args:
        device_path: Path to the device (e.g., "/dev/sda1" or "/dev/mapper/luks-UUID").

    Returns:
        Tuple[str, str]: Tuple containing the filesystem type and label, or ('', '') if not found.
    """
    fs_type = ''
    label = ''
    try:
        command = ["blkid", "-s", "TYPE", "-s", "LABEL", "-o", "value", device_path]
        print(f"Running: {' '.join(command)}")
        result = subprocess.run(
            command,
            capture_output=True,
            text=True,
            check=True
        )
        output_list = result.stdout.split()
        # Extract the filesystem type and label
        if len(output_list) > 0:
            fs_type = output_list[0]
        if len(output_list) > 1:
            label = output_list[1]
    except subprocess.CalledProcessError:
        pass

    print(f"Filesystem type of {device_path}: {fs_type}, Label: {label}")
    return (fs_type, label)

def is_partition_in_use(partition):
    """ Check if a partition is in use by any process. """
    for part in psutil.disk_partitions():
        if part.device == partition:
            return True
    return False

def encrypt_partition(
        partition: str,
        passphrase: str
    ) -> Optional[str]:
    """
    Encrypts the given partition with LUKS and returns the mapped partition path.

    Args:
        partition: Path to the partition (e.g., "/dev/sda1").
        passphrase: Passphrase for encryption.

    Returns:
        str: Mapped partition path if successful, else None.
    """
    if not os.path.exists(partition):
        print(f"Partition {partition} does not exist.")
        return None

    # Unmount the partition if mounted
    partition = unmount_partition(partition)
    if partition is None:
        print(f"Failed to unmount {partition}.")
        return None

    # Save current file system type and label
    fs_type, label = get_partition_info(partition)
    # Not all file systems are supported by LUKS
    supported_filesystems = ["ext2", "ext3", "ext4", "xfs", "btrfs"]
    if fs_type not in supported_filesystems:
        print(f"Filesystem {fs_type} is not supported by LUKS, default to ext4.")
        fs_type = "ext4"

    # Wipe the partition
    try:
        print(f"Wiping {partition}...")
        subprocess.run(
            ["dd", "if=/dev/zero", f"of={partition}", "bs=1M", "count=1"],
            check=True,
            stderr=subprocess.DEVNULL,  # Mute stderr
        )
    except subprocess.CalledProcessError as e:
        print(f"Failed to wipe {partition}: {e}")
        return None

    # Check if the partition is already LUKS-encrypted
    is_luks = subprocess.run(
        ["cryptsetup", "isLuks", partition],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        check=False
    ).returncode == 0

    if is_luks:
        command = ["cryptsetup", "reencrypt", "--encrypt", "--reduce-device-size", "16M", partition]
    else:
        command = ["cryptsetup", "luksFormat", partition]

    print(f"Running command: {' '.join(command)}")

    try:
        # If you want to use input= in subprocess.run() and also use stdout=subprocess.PIPE
        # or stderr=subprocess.PIPE, you are doing text-mode I/O with pipes.
        # This can cause encoding/decoding issues if you mix bytes and strings.
        # It's generally better to use universal_newlines=True when using pipes.
        subprocess.run(
            command,
            input=f"{passphrase}\n{passphrase}\n",
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            universal_newlines=True,
            check=True,
        )
    except subprocess.CalledProcessError as e:
        print(f"Failed to encrypt {partition}: {e}")
        return None

    # Connect the encrypted partition
    print("Connecting encrypted partition...")
    try:
        mapped_name = f"luks-{get_uuid_from_device(partition)}"
        command = ["cryptsetup", "open", partition, mapped_name]
        print(f"Running: {' '.join(command)}")
        subprocess.run(
            command,
                input=f"{passphrase}\n",
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                universal_newlines=True,
                check=True
        )
        mapped_partition = f"/dev/mapper/{mapped_name}"
    except subprocess.CalledProcessError as e:
        print(f"Failed to open {partition}: {e}")
        return None

    # Format the mapped filesystem
    try:
        command = ["mkfs", "-t", fs_type, "-L", label, mapped_partition]
        print(f"Running: {' '.join(command)}")
        subprocess.run(
            command,
            check=True,
            stderr=subprocess.DEVNULL,  # Mute stderr
        )
        print(f"Formatted {mapped_partition}: {fs_type}")
    except subprocess.CalledProcessError as e:
        print(f"Failed to format {mapped_partition}: {e}")

    return mapped_partition

def get_partition_from_mapped_path(mapped_path: str) -> str:
    """
    Get the underlying partition path from a mapped device path.

    Args:
        mapped_path: Path to the mapped device (e.g., "/dev/mapper/luks-123").

    Returns:
        str: Underlying partition path (e.g., "/dev/sda1"), or None if not found.
    """
    try:
        result = subprocess.run(
            ["dmsetup", "table", "--target", "crypt", mapped_path],
            capture_output=True,
            text=True,
            check=True
        )
        # Extract the last field, which is the underlying device
        return result.stdout.split()[-1].strip()
    except subprocess.CalledProcessError:
        return None

def read_passphrase(fifo_path: str) -> str:
    """ Reads the passphrase from the specified FIFO file. """
    if os.path.exists(fifo_path):
        with open(file=fifo_path, mode="r", encoding="utf-8") as fifo:
            return fifo.read().strip()
    print(f"Error: FIFO {fifo_path} does not exist.")
    return ""

def main():
    if len(sys.argv) != 3:
        print("Usage: sudo python3 encryption_sudo.py <fifo_path> <device_path> <device_path> ...")
        sys.exit(1)

    fifo_path = sys.argv[1]
    device_paths = sys.argv[2:]
    passphrase = read_passphrase(fifo_path)

    # Clean up FIFO
    try:
        os.unlink(fifo_path)
    except OSError:
        pass

    if not passphrase:
        print("Error: Failed to read passphrase from FIFO.")
        sys.exit(1)

    for device_path in device_paths:
        if not os.path.exists(device_path):
            print(f"Error: Device {device_path} does not exist - skipping.")
            continue

        print(f"Encrypting {device_path} with passphrase (hidden)")
        mapped_device_path = encrypt_partition(device_path, passphrase)
        if mapped_device_path:
            print(f"Encrypted device path: {mapped_device_path}")
        else:
            print(f"Error: Failed to encrypt {device_path}")

if __name__ == "__main__":
    # Start logging
    print("encryption_sudo.py started.")
    main()
