#!/bin/bash

### Unlocking and mounting an acquired LUKS encrypted volume bound to Tang
###
### Based on previous works by other authors, tuned for this purpose by
### Brian McDermott (bmcdermott at census-labs.com)

# Unlocks and mounts a LUKS encrypted volume using a raw disk image
# and the private “derive” key from a Tang server to which the volume was bound

######################################################
# Positional args
# 
# 1 - Path to RAW disk image
# 2 - path to Tang server private "derive" key - e.g /home/user/key.jwk
#
# USAGE: ./unlock_mount_volume.sh /path/to/disk.img /path/to/key.jwk

######################################################
# Pre-requisites
# - Disk file must be in RAW format
# - sudo apt install kpartx lvm2 jose
# - This has been tested on Ubuntu 22.04

######################################################
# Caveats:
# - If there already exists a Volume Group with the same name
#   as that on the encrypted volume, there may be a clash
#   In such case it is possible to change the name of a volume group.
# - Currently only tested on luks2 volumes
# - Does not currently unmount / do cleanup 
#     + Below is a template for a quick umount/cleanup script
#
# #!/bin/bash
#
# $DEV="${1}"
# $VG="${2}"
# $MNT="${3}"
# NAME="${NAME:-luks-"$(cryptsetup luksUUID $DEV)"}"
#
# # Unmount and detach drives 
#
# umount $MNT && \
# vgchange -an $VG && \
# cryptsetup luksClose $NAME && \
# kpartx -d $DEV && \
# losetup -d $DEV && \
# losetup
#
######################################################
#
# Example output on successful execution:
#
# Retrieving password for LUKS device: /dev/mapper/loop9p4
#
# Password retrieved: aBOs0HAL^ULJYz&olAbYbUcUpahZuL;4Nol90pDIRYs3m6IbVYq3Nv
#
# Device luks-a75ab556-a45e-40d2-a836-fd6780db7902 unlocked successfully
#
# 2 logical volume(s) in volume group "vgubuntu" now active
# ACTIVE            '/dev/vgubuntu/root' [20.31 GiB] inherit
# ACTIVE            '/dev/vgubuntu/swap_1' [<2.50 GiB] inherit
#
#
# Device mounted at:    /mnt/vgubuntu-root
#
######################################################

IMG="${1}"
KEY="${2}"

get_passphrase() {

    hdr="${1}"

    # B64 decoded 'header'
    jhd="$(jose b64 dec -i- <<< "$hdr")"; 

    # Client public key
    clt="$(jose fmt -j- -Og epk -Oo- <<< "$jhd")"; 

    # Private "derive" key leaked from Tang server                                    
    jwk="$(cat $KEY)"
        
    # ECDH -- c * S = K (see https://github.com/latchset/tang)
    jwk="$(jose jwk exc -l- -r- <<< "$jwk$clt")"      

    # DECRYPTION STEP
    password="$(jose jwe dec -k- -i- <<< "${jwk}${hdr}")"

    echo "$password"

}

mount_volume() {

    local DEV="${1}"
    PWD="${2}"

    # Get luksUUID of LUKS volume
    NAME="${NAME:-luks-"$(cryptsetup luksUUID $DEV)"}"

    # Unlock LUKS volume
    echo "$PWD" | cryptsetup luksOpen "$DEV" "$NAME" && echo "Device $NAME unlocked succesfully"
    echo

    # Mount device
    vgchange -ay

    # Output useful info
    lvscan
    echo

    # Get volume group name
    out=$(sudo vgdisplay --short)
    VG=$(echo $out | cut -d' ' -f1| cut -d'"' -f 2)

    # Create directory and mount volume
    [ ! -e /mnt/"$VG"-root ] && mkdir /mnt/"$VG"-root

    mount /dev/"$VG"/root /mnt/"$VG"-root

    echo
    echo "Device mounted at:    /mnt/"$VG"-root"
    echo

    return 0

}

# Taken from 
# https://github.com/latchset/clevis/blob/master/src/luks/clevis-luks-common-functions.in
clevis_luks_used_slots() {
    local DEV="${1:-}"
    [ -z "${DEV}" ] && return 1

    local used_slots
    if cryptsetup isLuks --type luks1 "${DEV}"; then
        if ! used_slots=$(cryptsetup luksDump "${DEV}" 2>/dev/null \
                          | sed -rn 's|^Key Slot ([0-7]): ENABLED$|\1|p'); then
            return 1
        fi
    elif cryptsetup isLuks --type luks2 "${DEV}"; then
        if ! used_slots=$(cryptsetup luksDump "${DEV}" 2>/dev/null \
                          | sed -rn 's|^\s+([0-9]+): luks2$|\1|p'); then
            return 1
        fi
    else
        echo "${DEV} is not a supported LUKS device!" >&2
        return 1
    fi
    echo "${used_slots}"
}

# Taken from 
# https://github.com/latchset/clevis/blob/master/src/luks/clevis-luks-common-functions.in
valid_slot() {
    local SLT="${1}"
    local MAX_SLOTS="${2}"
    case "${SLT}" in
        ''|*[!0-9]*)
            return 1
            ;;
        *)
            # We got an integer, now let's make sure it is within the
            # supported range.
            if [ "${SLT}" -ge "${MAX_SLOTS}" ]; then
                return 1
            fi
            ;;
    esac
}

# MODIFIED from 
# https://github.com/latchset/clevis/blob/master/src/luks/clevis-luks-common-functions.in
clevis_luks_read_slot() {
    local DEV="${1}"
    local SLT="${2}"

    if [ -z "${DEV}" ] || [ -z "${SLT}" ]; then
        echo "something missing"
        echo "Need both a device and a slot as arguments." >&2
        return 1
    fi

    # echo "reading slot $SLT"
    local DATA_CODED=''
    local MAX_LUKS2_SLOTS=32
    if cryptsetup isLuks --type luks2 "${DEV}"; then
        if ! valid_slot "${SLT}" "${MAX_LUKS2_SLOTS}"; then
            echo "Please, provide a valid key slot number; 0-31 for LUKS2" >&2
            return 1
        fi

        local token_id
        token_id=$(cryptsetup luksDump "${DEV}" | grep -E -B1 "^\s+Keyslot:\s+${SLT}$" | sed -n 1p | sed -rn 's|^\s+([0-9]+): clevis|\1|p')
        if [ -z "${token_id}" ]; then
            echo "Cannot load data from ${DEV} slot:${SLT}. No token found!" >&2
            return 1
        fi

        local token
        token=$(cryptsetup token export --token-id "${token_id}" "${DEV}")              
        DATA_CODED=$(jose fmt -j- -Og jwe -o- <<< "${token}" | jose jwe fmt -i- -c)

        if [ -z "${DATA_CODED}" ]; then
            echo "Cannot load data from ${DEV} slot:${SLT}!" >&2
            return 1
        fi
    else
        echo "${DEV} is not a supported LUKS device!" >&2
        return 1
    fi
    echo "${DATA_CODED}"
}

# MODIFIED from
# https://github.com/latchset/clevis/blob/master/src/luks/clevis-luks-common-functions.in
clevis_luks_unlock_device_by_slot() {
    local DEV="${1}"
    local SLT="${2}"

    [ -z "${DEV}" ] && return 1
    [ -z "${SLT}" ] && return 1

    local jwe passphrase
    if ! jwe="$(clevis_luks_read_slot "${DEV}" "${SLT}" 2>/dev/null)" || [ -z "${jwe}" ]; then
        return 1
    fi

    if ! passphrase="$(get_passphrase "${jwe}")" || [ -z "${passphrase}" ]; then
        return 1
    fi

    #clevis_luks_check_valid_key_or_keyfile "${DEV}" "${passphrase}" || return 1
    echo "${passphrase}"
}

# MODIFIED from
# https://github.com/latchset/clevis/blob/master/src/luks/clevis-luks-common-functions.in
clevis_luks_unlock_device() {
    local DEV="${1}"
    [ -z "${DEV}" ] && return 1

    local used_slots
    if ! used_slots=$(clevis_luks_used_slots "${DEV}") \
                      || [ -z "${used_slots}" ]; then
        return 1
    fi

    local slt pt
    for slt in ${used_slots}; do
        if ! pt=$(clevis_luks_unlock_device_by_slot "${DEV}" "${slt}") \
                  || [ -z "${pt}" ]; then
             continue
        fi
        break
    done

    echo "$pt"
}

# mount device on available /dev/loopX
losetup -f $IMG

# Find out which loop device was used
DEV=$(losetup -l | grep $IMG | cut -d' ' -f1)

# Add partition mappings - see /dev/mapper 
kpartx -a $DEV

SUBSTR=$(echo $DEV | cut -d'/' -f3) # e.g. 'loop9'

for file in /dev/mapper/*; do       # e.g. 'loop9p1, loop9p2, ... etc.'

    if [[ "$file" == *"$SUBSTR"* ]]; then
        if cryptsetup isLuks --type luks2 $file; then
            echo "Retrieving password for LUKS device: $file"
            echo

            if ! password=$(clevis_luks_unlock_device "${file}") \
                        || [ -z "${password}" ]; then
                echo "[!] Could not retrieve password"
                exit 1
            fi 


            echo "Password retrieved: $password"
            echo

            mount_volume $file $password

            exit 0
            
        fi
    fi

done

exit 1