#!/usr/bin/bash

set -e

usage() {
    echo "Usage: $0 --size <> [--location <>] [--password <>]"
    echo ""
    echo "Generate a luks container to store sensitive data on Cosmian VM"
    echo "Mount this container into $CONTAINER_MOUNT_PATH"
    echo "Configure auto-mounting when rebooting by decrypting using the TPM"
    echo ""
    echo "Example: $0 --size 500MB"
    echo "Example: echo mypassword | $0 --size 500MB --password -"
    echo "Example: $0 --size 500MB --password mypassword"
    echo ""
    echo "Required arguments:"
    echo -e "\t--size         The size of the container (example: 500MB or 10GB)"
    echo ""
    echo "Optional arguments:"
    echo -e "\t--location     The file path to store the container (default: $CONTAINER_PATH)"
    echo -e "\t--password     The password to encrypt/decrypt the luks"
    echo -e "\t               If the argument is not passed, a prompt will ask for it"
    echo -e "\t               If the value is '-', the stdin is read."
    exit 1
}

set_default_variables() {
    # Mandatory args (initialized empty)
    CONTAINER_SIZE=""
    PASSWORD=""

    # Optional args
    DEFAULT_ROOT="/var/lib/cosmian_vm"
    CONTAINER_PATH="$DEFAULT_ROOT/container"
    HEADER_PATH="$DEFAULT_ROOT/header"
    CONTAINER_MAPPING_NAME="cosmian_vm_container"
    CONTAINER_MAPPING_PATH="/dev/mapper/$CONTAINER_MAPPING_NAME"
    CONTAINER_MOUNT_PATH="$DEFAULT_ROOT/data"
}

parse_args() {
    # Parse args
    while [[ $# -gt 0 ]]; do
        case $1 in
        --size)
            CONTAINER_SIZE="$2"
            shift # past argument
            shift # past value
            ;;

        --location)
            CONTAINER_PATH="$2"
            shift # past argument
            shift # past value
            ;;

        --password)
            PASSWORD="$2"
            shift # past argument
            shift # past value
            ;;

        -*)
            usage
            ;;
        esac
    done

    if [ -z "$CONTAINER_SIZE" ] || [ -z "$CONTAINER_PATH" ]; then
        usage
    fi

    if [ "$PASSWORD" = "-" ]; then
        PASSWORD=$(cat /dev/stdin)
    fi

    while [ -z "$PASSWORD" ]; do
        read -s -r -p "Enter a luks password: " PASSWORD
        echo ""
    done
}

set_default_variables
parse_args "$@"

if [ "$EUID" -ne 0 ]; then
    echo "Please run as root"
    exit
fi

mkdir -p "$(dirname "$CONTAINER_PATH")"
mkdir -p "$(dirname "$CONTAINER_MOUNT_PATH")"

if [ -e "$CONTAINER_PATH" ]; then
    echo "A container already exists in $CONTAINER_PATH (remove it before going any further)"
    exit 1
fi

if [ -e "$HEADER_PATH" ]; then
    echo "A LUKS header already exists in $HEADER_PATH (remove it before going any further)"
    exit 1
fi

# Make sure to close/umount existing container
if [ -e "$CONTAINER_MAPPING_PATH" ]; then
    echo "Closing previous mounted container..."
    sync || true
    umount "$CONTAINER_MAPPING_PATH" || true
    cryptsetup close "$CONTAINER_MAPPING_NAME" || true
fi

# Allocate the container
echo "Creating a $CONTAINER_SIZE container..."
fallocate -l "$CONTAINER_SIZE" "$CONTAINER_PATH"

# Encrypt the container (a password is required to run this command)
echo "Encrypting the container (with password=${PASSWORD})..."
echo -n "$PASSWORD" | cryptsetup luksFormat "$CONTAINER_PATH" --type luks2 --integrity hmac-sha256 --header "$HEADER_PATH" --key-file -

# Open the container and map it (a password is required to run this command)
echo "Opening the container at $CONTAINER_MAPPING_PATH..."
echo -n "$PASSWORD" | cryptsetup luksOpen --header "$HEADER_PATH" --key-file - "$CONTAINER_PATH" "$CONTAINER_MAPPING_NAME"

# Format it
echo "Formatting the container in Ext4..."
mkfs -t ext4 "$CONTAINER_MAPPING_PATH"

# Mount it
echo "Mounting the container at $CONTAINER_MOUNT_PATH..."
mkdir -p "$CONTAINER_MOUNT_PATH"
mount "$CONTAINER_MAPPING_PATH" "$CONTAINER_MOUNT_PATH"

# Determine the block device
BLOCK_DEVICE=$(losetup -ln --raw -O NAME,BACK-FILE | grep "$CONTAINER_PATH" | cut -f1 -d" ")

if [ -z "$BLOCK_DEVICE" ]; then
    echo "Can't find the block device attached to $CONTAINER_PATH"
    exit 1
fi

# Enroll the TPM to decrypt the luks without password (a password is required to run this command)
# Details on PCRs: https://uapi-group.org/specifications/specs/linux_tpm_pcr_registry/
echo "Enrolling the TPM for this container on block device $BLOCK_DEVICE..."

set +e
PASSWORD=$PASSWORD systemd-cryptenroll --tpm2-device=auto --wipe-slot=tpm2 "$HEADER_PATH"
if [ $? -ne 0 ]; then
    # Need to clean container after failure
    rm -f "$CONTAINER_PATH"
    rm -f "$HEADER_PATH"
    exit 1
fi

# Display debug information
set -x
cryptsetup luksDump "$HEADER_PATH"
set +x

echo "Process completed with success!"
