#!/usr/bin/python3

import argparse
import getpass
import json
import os
import re
import socket
import subprocess
import sys
import time
import traceback

import requests

# Utility functions


def sh(cmd, debug=os.environ.get("GAN_DEBUG")):
    if debug:
        print(f"+ {cmd}")
    r = (
        subprocess.run(cmd, shell=True, capture_output=True, check=True)
        .stdout.decode()
        .strip()
    )
    if debug and r:
        print("> " + "> ".join(r.splitlines()))
    return r


def check_kavlan():
    print("Checking if the node is in a KaVLAN")
    hostname = sh("hostname")
    kavlan_number = None
    if len(hostname.split("-")) >= 4:
        try:
            kavlan_number = int(hostname.split("-")[3].split(".")[0])
        except ValueError:
            print("Failed to get the KaVLAN number from the hostname.")
    if kavlan_number is None or kavlan_number < 4:
        print(
            f"Node {hostname} is not in an isolated KaVLAN. "
            "You are probably not on an Armored Node."
        )
        sys.exit(1)


def check_home_encrypted():
    print("Checking if the user's home directory is encrypted")
    df_output = sh("df -T {}".format(f"{os.environ['HOME']}"))
    if "encrypted" not in df_output:
        print(
            "No encrypted mount point found in the user's home directory. "
            "Have you already launched g5k-armor-node.py?"
        )
        sys.exit(1)


def load_config():
    print(
        "Loading CompuVault configuration "
        "from the PROJECTNAME_cv_config.json file"
    )
    config_files = [
        file
        for file in os.listdir(f"{os.environ['HOME']}")
        if file.endswith("_cv_config.json")
    ]
    if not config_files:
        print(
            "PROJECTNAME_cv_config.json not found. "
            "Please check it in your home directory.",
            file=sys.stderr,
        )
        sys.exit(1)
    if len(config_files) > 1:
        print(
            "Multiple *_cv_config.json files found. "
            "Please ensure there is only one "
            "file ending with _cv_config.json.",
            file=sys.stderr,
        )
        sys.exit(1)
    config_file_path = os.path.join(f"{os.environ['HOME']}", config_files[0])
    if os.path.getsize(config_file_path) == 0:
        print(
            f"The configuration file {config_file_path} is empty.",
            file=sys.stderr,
        )
        sys.exit(1)
    try:
        with open(config_file_path) as f:
            config = json.load(f)
    except json.JSONDecodeError as e:
        print(
            f"Error decoding JSON from the file {config_file_path}: {e}",
            file=sys.stderr,
        )
        sys.exit(1)
    return config


def get_storage_server_ip(config):
    print("Getting the storage server IP address")
    try:
        storage_server_ip = socket.gethostbyname(config["STORAGE_SERVER"])
        return storage_server_ip
    except socket.error as e:
        print(
            f"Failed to get the IP address of {config['STORAGE_SERVER']}: {e}"
        )
        sys.exit(1)


def install_required_packages():
    print("Installing required packages")
    sh("sudo apt install -y open-iscsi cryptsetup")


def discover_iscsi_target(storage_server_ip):
    print("Discovering the iSCSI target and getting its IQN")
    sh(
        f"sudo iscsiadm --mode discoverydb --type sendtargets "
        f"--portal {storage_server_ip} --discover"
    )


def configure_iscsi_authentication(config, storage_server_ip):
    print("Configuring iSCSI authentication")
    default_path = (
        f"/etc/iscsi/nodes/{config['EXPORT_NAME']}/"
        f"{storage_server_ip},3260,1/default"
    )
    sh(
        "sudo sed -i "
        "'s/^node.session.auth.authmethod.*/"
        "node.session.auth.authmethod = CHAP/' "
        f"{default_path}"
    )
    sh(
        f"echo 'node.session.auth.username = {config['ISCSI_LOGIN']}' | "
        f"sudo tee -a {default_path}"
    )
    sh(
        f"echo 'node.session.auth.password = {config['ISCSI_PASSWORD']}' | "
        f"sudo tee -a {default_path}"
    )


def logout_iscsi_session(config, storage_server_ip):
    print("Logging out from the iSCSI session")
    try:
        sh(
            "sudo iscsiadm --mode node --targetname "
            f"\"{config['EXPORT_NAME']}\" "
            f"--portal {storage_server_ip}:3260 --logout"
        )
        log_action(site, node, user, "iscsi logout succeeded")
    except subprocess.CalledProcessError as e:
        print(f"Failed to log out from the iSCSI session: {e}")
        log_action(site, node, user, "iscsi logout failed")


def login_iscsi_session(config, storage_server_ip):
    print("Logging into the iSCSI session")
    try:
        sh(
            "sudo iscsiadm --mode node --targetname "
            f"\"{config['EXPORT_NAME']}\" "
            f"--portal {storage_server_ip}:3260 --login"
        )
        log_action(site, node, user, "iscsi login succeeded")
    except subprocess.CalledProcessError:
        print(
            "Active iSCSI sessions found. "
            "Attempting to logout and log back in."
        )
        logout_iscsi_session(config, storage_server_ip)
        try:
            sh(
                "sudo iscsiadm --mode node --targetname "
                f"\"{config['EXPORT_NAME']}\" "
                f"--portal {storage_server_ip}:3260 --login"
            )
            log_action(site, node, user, "iscsi login succeeded")
        except subprocess.CalledProcessError as e:
            print(f"Failed to log into the iSCSI session: {e}")
            log_action(site, node, user, "iscsi login failed")
            sys.exit(1)


def get_iscsi_disk():
    print("Getting the iSCSI disk")
    start_time = time.time()
    while True:
        if time.time() - start_time > 30:
            print("iSCSI disk name not found in 30 seconds")
            sys.exit(1)
        try:
            disk_output = sh("sudo iscsiadm -m session -P3")
            disk_name_match = re.search(
                r"Attached scsi disk (\w+)", disk_output
            )
            if disk_name_match:
                disk = disk_name_match.group(1)
                print("iSCSI disk name is:", disk)
                return disk
        except subprocess.CalledProcessError as e:
            print(f"Failed to get the iSCSI disk: {e}")
        time.sleep(5)


def open_encrypted_storage(disk, config, passphrase):
    print("Opening encrypted storage")
    process = subprocess.Popen(
        [
            "sudo",
            "cryptsetup",
            "open",
            f"/dev/{disk}",
            f"cv-{config['PROJECT_NAME']}",
        ],
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    stdout, stderr = process.communicate(input=f"{passphrase}".encode())
    if process.returncode != 0:
        print(f"Error opening encrypted storage: {stderr.decode()}")
        sys.exit(1)
    if not os.path.exists(f"/dev/mapper/cv-{config['PROJECT_NAME']}"):
        print("Failed to open encrypted storage")
        sys.exit(1)
    log_action(
        site,
        node,
        user,
        f"encrypted storage cv-{config['PROJECT_NAME']} opened",
    )


def api(url):
    base_url = "https://api.grid5000.fr/stable/"
    r = requests.get(base_url + url)
    r.raise_for_status()
    return r.json()


def get_node_site():
    nodesite = sh("hostname -f").split(".")[0:2]
    if len(nodesite) != 2:
        raise RuntimeError("Cannot retrieve node and site from 'hostname -f'")
    node, site = nodesite
    # Drop kavlan suffix from the node
    return "-".join(node.split("-")[0:2]), site


def get_node_user(node, site):
    try:
        n = ".".join(get_node_site()) + ".grid5000.fr"
        reservations = api(f"sites/{site}/status")["nodes"][n]["reservations"]
        return [
            item["user_uid"]
            for item in reservations
            if item["state"] == "running"
        ][0]
    except (IndexError, requests.exceptions.HTTPError):
        traceback.print_exc()
        print(
            "Cannot retrieve user for node " f"{node}.{site}.grid5000.fr",
            file=sys.stderr,
        )
        sys.exit(1)


def log_action(site, node, user, action):
    timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
    sh(
        f"logger -n syslog --tcp -P514 -p authpriv.info "
        f'"CompuVault {action} on {node}.{site} by {user} at {timestamp}"'
    )


def is_mounted(mount_point):
    print("Checking if the storage is mounted")
    try:
        result = subprocess.run(
            ["findmnt", "-rn", "-o", "TARGET", mount_point],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )
        return result.stdout.strip() == mount_point
    except Exception as e:
        print(f"Error checking if mount point is active: {e}")
        return False


# Start of the script

parser = argparse.ArgumentParser(description="Manage CompuVault operations")
parser.add_argument(
    "operation",
    choices=["init", "mount", "disconnect"],
    help="Operation to perform: 'init', 'mount' or 'disconnect'",
)
args = parser.parse_args()

check_kavlan()
check_home_encrypted()
node, site = get_node_site()
user = get_node_user(node, site)
config = load_config()
storage_server_ip = get_storage_server_ip(config)
install_required_packages()

if args.operation == "init":
    discover_iscsi_target(storage_server_ip)
    configure_iscsi_authentication(config, storage_server_ip)
    login_iscsi_session(config, storage_server_ip)
    disk = get_iscsi_disk()

    print("Checking if a LUKS signature already exists for the iSCSI disk")
    output = sh(f"sudo lsblk -o FSTYPE /dev/{disk}")
    if "LUKS" in output:
        print(
            "A LUKS signature already exists. "
            "You do not need to initialize CompuVault. "
            "Please use the 'mount' operation instead."
        )
        sys.exit(1)

    print("Encrypting storage with LUKS")
    passphrase = getpass.getpass(
        "Please enter your passphrase for LUKS encryption:"
    )
    process = subprocess.Popen(
        ["sudo", "cryptsetup", "-q", "luksFormat", f"/dev/{disk}"],
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    stdout, stderr = process.communicate(input=f"{passphrase}\n".encode())
    if process.returncode != 0:
        print(f"Error during LUKS encryption: {stderr.decode()}")
        sys.exit(1)
    log_action(
        site,
        node,
        user,
        f"encrypted iSCSI disk {disk} with LUKS successfully",
    )

    open_encrypted_storage(disk, config, passphrase)

    print("Formatting ext4")
    sh(f"sudo mkfs.ext4 /dev/mapper/cv-{config['PROJECT_NAME']}")

    print("Closing encrypted storage")
    sh(f"sudo cryptsetup close cv-{config['PROJECT_NAME']}")

    logout_iscsi_session(config, storage_server_ip)

    print("Initialization of CompuVault completed successfully!")
elif args.operation == "mount":
    discover_iscsi_target(storage_server_ip)
    configure_iscsi_authentication(config, storage_server_ip)
    login_iscsi_session(config, storage_server_ip)
    disk = get_iscsi_disk()

    if is_mounted(f"/mnt/cv-{config['PROJECT_NAME']}"):
        print(
            "The storage is already mounted. "
            "You do not need to mount it again. "
        )
        sys.exit(1)
    passphrase = getpass.getpass(
        "Please enter your passphrase for LUKS encryption:"
    )
    open_encrypted_storage(disk, config, passphrase)

    print("Mounting encrypted storage")
    sh(f"sudo mkdir -p /mnt/cv-{config['PROJECT_NAME']}")
    sh(
        f"sudo mount /dev/mapper/cv-{config['PROJECT_NAME']} "
        f"/mnt/cv-{config['PROJECT_NAME']}"
    )
    sh(f"sudo chown {os.environ['USER']}: /mnt/cv-{config['PROJECT_NAME']}")

    log_action(
        site,
        node,
        user,
        f"encrypted storage cv-{config['PROJECT_NAME']} mounted",
    )
    print("Mounting of CompuVault completed successfully!")
elif args.operation == "disconnect":
    if not is_mounted(f"/mnt/cv-{config['PROJECT_NAME']}"):
        print(
            "The storage is not mounted. "
            "You do not need to disconnect from it. "
        )
        sys.exit(1)
    print("Disconnecting the encrypted storage")
    sh(f"sudo umount /mnt/cv-{config['PROJECT_NAME']}")
    sh(f"sudo cryptsetup close cv-{config['PROJECT_NAME']}")
    logout_iscsi_session(config, storage_server_ip)
    log_action(
        site,
        node,
        user,
        f"encrypted storage cv-{config['PROJECT_NAME']} disconnected",
    )
    print("Disconnecting of CompuVault completed successfully!")