import os
import subprocess
from crypt_tools import decrypt_passphrase, create_luks_mount_config
import argparse


# Configuration file containing encrypted LUKS passphrases
CONFIG_FILE = "luks_mount.conf"




def get_device_uuid(device):
    """
    Get the UUID of a device using blkid.
    """
    try:
        output = subprocess.check_output(["blkid", "-o", "value", "-s", "UUID", device])
        return output.decode().strip()
    except subprocess.CalledProcessError as e:
        print(f"Failed to get UUID for {device}: {e}")
        return None

def get_mount_point(device_id):
    """
    Get the mount point for a given device ID from the configuration file.
    """
    if not os.path.exists(CONFIG_FILE):
        print(f"Configuration file '{CONFIG_FILE}' not found!")
        return None

    with open(CONFIG_FILE, "r") as f:
        for line in f:
            if line.startswith("#") or len(line) == 0:
                continue            
            parts = line.strip().split()
            if len(parts) != 3:
                continue

            uuid, mount_point, _ = parts
            if uuid == device_id:
                return mount_point

    return None


def mount_luks_filesystems():
    """
    Read the configuration file, decrypt the LUKS passphrases, and mount the filesystems.
    """
    if not os.path.exists(CONFIG_FILE):
        print(f"Configuration file '{CONFIG_FILE}' not found!")
        return

    with open(CONFIG_FILE, "r") as f:
        for line in f:
            if line.startswith("#") or len(line) == 0:
                continue            
            # Parse the configuration line: uuid mount_point encrypted_passphrase
            parts = line.strip().split()
            if len(parts) != 3:
                print(f"Skipping invalid line: {line}")
                continue

            uuid, mount_point, encrypted_passphrase = parts

            # Get the device path from the UUID
            device = f"/dev/disk/by-uuid/{uuid}"
            if not os.path.exists(device):
                print(f"Device '{device}' not found!")
                continue

            # Decrypt the LUKS passphrase using the master passkey
            try:
                passphrase = decrypt_passphrase(encrypted_passphrase, private_key_file="private_key.pem")
#                print ("---> %s" % passphrase)
            except Exception as e:
                print(f"Failed to decrypt passphrase for {device}: {e}")
                continue

            # Unlock the LUKS device
            try:
                subprocess.run(
                    ["cryptsetup", "luksOpen", device, f"{os.path.basename(device)}_crypt"],
                    input=passphrase.encode(),
                    check=True
                )
                print(f"Unlocked LUKS device: {device}")
            except subprocess.CalledProcessError as e:
                print(f"Failed to unlock {device}: {e}")
                continue

            # Mount the filesystem
            try:
                os.makedirs(mount_point, exist_ok=True)
                subprocess.run(
                    ["mount", f"/dev/mapper/{os.path.basename(device)}_crypt", mount_point],
                    check=True
                )
                print(f"Mounted {device} at {mount_point}")
            except subprocess.CalledProcessError as e:
                print(f"Failed to mount {device} at {mount_point}: {e}")
                continue


def unmount_and_lock_filesystem(device_id):
    """
    Unmount and lock a filesystem given the device ID.
    """
    device = f"/dev/disk/by-uuid/{device_id}"
    crypt_device = f"/dev/mapper/{os.path.basename(device)}_crypt"

    # Get the mount point from the configuration file
    mount_point = get_mount_point(device_id)
    if not mount_point:
        print(f"Mount point for device '{device_id}' not found in configuration file.")
        return

    # Unmount the filesystem
    try:
        subprocess.run(["sudo", "umount", crypt_device], check=True)
        print(f"Unmounted filesystem: {crypt_device}")
    except subprocess.CalledProcessError as e:
        print(f"Failed to unmount filesystem {crypt_device}: {e}")
        return
    

    # Remove the mount point directory
    try:
        os.rmdir(mount_point)
        print(f"Removed mount point directory: {mount_point}")
    except OSError as e:
        print(f"Failed to remove mount point directory {mount_point}: {e}")

    # Lock the LUKS device
    try:
        subprocess.run(["sudo", "cryptsetup", "luksClose", f"{os.path.basename(device)}_crypt"], check=True)
        print(f"Locked LUKS device: {device}")
    except subprocess.CalledProcessError as e:
        print(f"Failed to lock LUKS device {device}: {e}")

def close_all_filesystems():
    """
    Close all filesystems listed in the configuration file.
    """
    if not os.path.exists(CONFIG_FILE):
        print(f"Configuration file '{CONFIG_FILE}' not found!")
        return

    with open(CONFIG_FILE, "r") as f:
        for line in f:
            parts = line.strip().split()
            if len(parts) != 3:
                print(f"Skipping invalid line: {line}")
                continue

            uuid, _, _ = parts
            unmount_and_lock_filesystem(uuid)


def main():
    """
    Main function to read the master passkey and mount LUKS filesystems.
    """
    parser = argparse.ArgumentParser(description="Automount LUKS filesystems.")
    parser.add_argument("-m", "--mount", action="store_true", help="Mount all LUKS devices")
    parser.add_argument("-c", "--close", help="Close a specific LUKS device", metavar="DEVICE_ID")
    parser.add_argument("-f", "--config", help="Specify a configuration file", metavar="CONFIG_FILE")
    parser.add_argument("-k", "--close-all", action="store_true", help="Close all LUKS devices")
    parser.add_argument("-n", "--new-config", action="store_true", help="Create a new luks_mount.conf file")


    args = parser.parse_args()
    global CONFIG_FILE
    if args.config:

        CONFIG_FILE = args.config
    if args.mount:        
        mount_luks_filesystems()
    elif args.close:
        unmount_and_lock_filesystem(args.close)
    elif args.close_all:
        close_all_filesystems()
    elif args.new_config:
        create_luks_mount_config(config_file=CONFIG_FILE)
    else:
        parser.print_help()

if __name__ == "__main__":
    main()