#!/bin/sh
#
# genluks - Lay out an encrypted install in one command
# rawnix.org - https://rawnix.org
#
# Usage: genluks <mountpoint> <dev>:<target>[:<opts>] ...
#
#   genluks /mnt /dev/sdX2:/
#   genluks /mnt /dev/sdX2:/ /dev/sdX3:/var /dev/sdX6:/home /dev/sdX8:swap
#   genluks /mnt /dev/sdX2:/ /dev/sdX9:/data2:plain,xfs
#
#   <target>   a path, or the bare word "swap" (not a mountpoint, so no
#              leading slash, and nothing is mounted for it)
#   <opts>     comma-separated, any order:
#                plain            no encryption at all
#                ext4|xfs|btrfs   filesystem, default ext4
#
# Encrypted is the default; "plain" is the opt-out, for a partition
# another OS has to be able to read.
#
# For each volume: luksFormat, open, mkfs, mount -- in the order the
# mountpoints nest, not the order given. Then every non-root LUKS
# volume gets a keyfile on the encrypted root, and /etc/crypttab is
# written. The result is one passphrase at boot: the initramfs unlocks
# root, /etc/runit/1 opens the rest from the keyfiles.
#
#     nvme0n1p1  /boot/efi   not handled here -- the ESP is never
#                            encrypted, and is a plain mkfs.fat + mount
#     nvme0n1p2  /           passphrase, via the initramfs
#     nvme0n1p3  /var    ]
#     nvme0n1p4  swap    ]-- keyfile in /etc/cryptkeys
#     nvme0n1p6  /home   ]
#     nvme0n1p9  /data2      plain, shared with another OS
#
# THIS DESTROYS DATA, the way mkfs does. It prints what it is about to
# erase and waits for a typed "yes".
#
# On the keyfiles: they live on the encrypted root at mode 400, and
# that is their only protection. It is also enough -- with the machine
# off they are inside the LUKS container, and once it is running and
# unlocked anyone with root can read the volume keys from the kernel
# keyring anyway. The consequence worth understanding is that the root
# passphrase now guards every volume, so it is the one that has to be
# good: separate partitions are a quota and fsck boundary here, not a
# security one.
#
# The passphrase is also kept as a slot on every volume, so there is a
# way in when root itself cannot be mounted.

set -e

die()  { echo "ERROR: $*" >&2; exit 1; }
say()  { printf '>>> %s\n' "$*"; }
info() { printf '    %s\n' "$*"; }

[ "$(id -u)" = 0 ] || die "Must run as root"
command -v cryptsetup >/dev/null 2>&1 || die "cryptsetup not found"
[ $# -ge 2 ] || die "Usage: genluks <mountpoint> <dev>:<target>[:<opts>] ..."

MNT="${1%/}"; shift
[ -n "$MNT" ] || MNT="/"
[ -d "$MNT" ] || die "Directory not found: $MNT"
! mountpoint -q "$MNT" || die "$MNT is already a mountpoint — unmount it first"

# ── The disk the running system booted from ───────────────────
# Formatting it takes the installer out from under itself, and in a VM
# or off a USB stick the target is one letter away.
RUNDISK=""
_rootsrc=$(findmnt -n -o SOURCE / 2>/dev/null | head -1)
case "$_rootsrc" in
    /dev/*) RUNDISK=$(lsblk -n -o PKNAME "$_rootsrc" 2>/dev/null | head -1) ;;
esac

# ── Parse and validate everything before touching anything ────
# Finding the typo in argument five after one to four have been erased
# is not a useful place to find it.
SPECS=""
HAVE_ROOT=0

for arg in "$@"; do
    dev=${arg%%:*}
    rest=${arg#*:}
    [ "$dev" != "$arg" ] || die "$arg: expected <dev>:<target>[:<opts>]"

    case "$rest" in
        *:*) target=${rest%%:*}; opts=${rest#*:} ;;
        *)   target=$rest;       opts="" ;;
    esac
    [ -n "$target" ] || die "$arg: empty target"

    fstype=""
    plain=0
    _o=$opts
    while [ -n "$_o" ]; do
        case "$_o" in *,*) one=${_o%%,*}; _o=${_o#*,} ;; *) one=$_o; _o="" ;; esac
        case "$one" in
            "")    ;;
            plain) plain=1 ;;
            ext4|xfs|btrfs)
                   [ -z "$fstype" ] || die "$arg: two filesystems given"
                   fstype=$one ;;
            *)     die "$arg: unknown option '$one' (plain, ext4, xfs, btrfs)" ;;
        esac
    done
    [ -n "$fstype" ] || fstype=ext4

    # Checked here, with everything else, rather than at mkfs time: a
    # missing mkfs discovered on the fourth volume leaves the first
    # three formatted and the run half done.
    if [ "$target" != swap ]; then
        command -v "mkfs.$fstype" >/dev/null 2>&1 || \
            die "$arg: mkfs.$fstype is not installed"
    fi

    case "$target" in
        swap)      [ "$opts" = "" ] || [ "$plain" -eq 1 ] || die "$arg: swap takes no filesystem" ;;
        /boot/efi) die "$arg: the ESP is not handled here — mkfs.fat -F32 and mount it yourself" ;;
        /*)        ;;
        *)         die "$arg: target must be an absolute path, or the word 'swap'" ;;
    esac

    if [ "$target" = "/" ]; then
        [ "$plain" -eq 0 ] || \
            die "$arg: root cannot be plain — every other volume's keyfile lives on it"
        HAVE_ROOT=1
    fi

    [ -b "$dev" ] || die "$dev is not a block device"
    ! findmnt -n -S "$dev" >/dev/null 2>&1 || die "$dev is mounted — unmount it first"
    ! lsblk -n -o TYPE "$dev" 2>/dev/null | grep -q '^crypt$' || \
        die "$dev already has an open mapping — close it first"
    if [ -n "$RUNDISK" ]; then
        _pk=$(lsblk -n -o PKNAME "$dev" 2>/dev/null | head -1)
        [ "$_pk" != "$RUNDISK" ] || \
            die "$dev is on $RUNDISK, the disk the running system booted from"
    fi

    for s in $SPECS; do
        [ "${s%%|*}" != "$dev" ] || die "$dev given twice"
        _t=${s#*|}; _t=${_t%%|*}
        [ "$_t" != "$target" ] || die "$target given twice"
    done

    SPECS="$SPECS $dev|$target|$fstype|$plain"
done

[ "$HAVE_ROOT" -eq 1 ] || die "no volume for / — the install has nowhere to go"

# ── Names ─────────────────────────────────────────────────────
# The root name ends up in the kernel cmdline and all of them end up in
# crypttab, so they are derived, never invented.
mapname() {
    case "$1" in
        /)    echo croot ;;
        swap) echo cswap ;;
        *)    echo "c$(echo "$1" | tr -dc 'a-zA-Z0-9')" ;;
    esac
}

spec_dev()    { echo "${1%%|*}"; }
spec_target() { _r=${1#*|}; echo "${_r%%|*}"; }
spec_fs()     { _r=${1#*|}; _r=${_r#*|}; echo "${_r%%|*}"; }
spec_plain()  { _r=${1#*|}; _r=${_r#*|}; echo "${_r#*|}"; }

# ── Show and confirm ──────────────────────────────────────────
echo ""
say "This will ERASE the following, losing everything on them:"
echo ""
printf '    %-16s %-12s %-8s %s\n' DEVICE TARGET FS MODE
for s in $SPECS; do
    target=$(spec_target "$s"); fstype=$(spec_fs "$s")
    [ "$target" != swap ] || fstype=swap
    if   [ "$(spec_plain "$s")" -eq 1 ]; then mode="plain"
    elif [ "$target" = "/" ];            then mode="LUKS (passphrase)"
    else                                      mode="LUKS (keyfile)"
    fi
    printf '    %-16s %-12s %-8s %s\n' "$(spec_dev "$s")" "$target" "$fstype" "$mode"
done
echo ""
printf 'Type yes to continue: '
IFS= read -r _yes || true
[ "$_yes" = "yes" ] || die "aborted"
echo ""

# ── Passphrase, once ──────────────────────────────────────────
# Written to a file in /run rather than piped: `--key-file -` has its
# own newline handling, and a passphrase that silently gains or loses a
# trailing byte fails in a way that looks exactly like a typo. /run is
# tmpfs, so this never reaches a disk, and the trap is armed before the
# file exists so no exit path can leave one behind.
PASSFILE=""
cleanup() {
    if [ -n "$PASSFILE" ]; then
        rm -f "$PASSFILE"
        PASSFILE=""
    fi
}
trap cleanup EXIT
trap 'cleanup; exit 130' INT
trap 'cleanup; exit 143' TERM

NEED_PASS=0
for s in $SPECS; do
    [ "$(spec_plain "$s")" -eq 1 ] || NEED_PASS=1
done

if [ "$NEED_PASS" -eq 1 ]; then
    PASSFILE=$(mktemp /run/genluks.XXXXXX) || die "cannot create a temp file in /run"
    chmod 600 "$PASSFILE"

    # Typed at boot before any keymap is loaded, so it is entered on a
    # US layout whatever the installed system ends up using. Better
    # said here than discovered at the first reboot.
    echo "The root passphrase is typed at boot before any keymap loads:"
    echo "use ASCII, and remember it is entered as US."
    echo ""
    stty -echo 2>/dev/null || true
    printf 'Passphrase: '         >&2; IFS= read -r _p1 || true; printf '\n' >&2
    printf 'Passphrase (again): ' >&2; IFS= read -r _p2 || true; printf '\n' >&2
    stty echo 2>/dev/null || true

    [ -n "$_p1" ]       || die "empty passphrase"
    [ "$_p1" = "$_p2" ] || die "passphrases do not match"
    printf '%s' "$_p1" > "$PASSFILE"
    _p1=""; _p2=""
    echo ""
fi

# ── Format ────────────────────────────────────────────────────
mkfs_on() {
    case "$2" in
        ext4)  mkfs.ext4  -q -F "$1" ;;
        xfs)   mkfs.xfs   -q -f "$1" ;;
        btrfs) mkfs.btrfs -q -f "$1" ;;
        *)     die "unknown filesystem $2" ;;
    esac
}

for s in $SPECS; do
    dev=$(spec_dev "$s"); target=$(spec_target "$s")
    fstype=$(spec_fs "$s"); plain=$(spec_plain "$s")
    name=$(mapname "$target")

    if [ "$plain" -eq 1 ]; then
        say "$dev -> $target (plain, $fstype)"
        mkfs_on "$dev" "$fstype"
        continue
    fi

    say "$dev -> $target (LUKS, $name)"
    # --batch-mode: the "YES" prompt is the confirmation already given
    cryptsetup luksFormat --type luks2 --batch-mode --key-file "$PASSFILE" "$dev"
    cryptsetup open --key-file "$PASSFILE" "$dev" "$name"

    if [ "$target" = swap ]; then
        mkswap "/dev/mapper/$name" >/dev/null
    else
        mkfs_on "/dev/mapper/$name" "$fstype"
    fi
done

# ── Mount, shallowest first ───────────────────────────────────
# /mnt must exist before /mnt/var can be mounted into it, and the
# argument order is the user's, not necessarily the nesting order.
echo ""
say "Mounting"
_order=$(for s in $SPECS; do
             t=$(spec_target "$s")
             [ "$t" != swap ] || continue
             printf '%s\t%s\n' "$(printf '%s' "$t" | tr -cd / | wc -c)" "$s"
         done | sort -n -k1,1 | cut -f2)

for s in $_order; do
    target=$(spec_target "$s")
    name=$(mapname "$target")
    if [ "$(spec_plain "$s")" -eq 1 ]; then
        src=$(spec_dev "$s")
    else
        src="/dev/mapper/$name"
    fi
    case "$target" in
        /) dest="$MNT" ;;
        *) dest="$MNT$target" ;;
    esac
    mkdir -p "$dest"
    mount "$src" "$dest"
    info "$src -> $dest"
done

# ── Keyfiles and crypttab ─────────────────────────────────────
KDIR="$MNT/etc/cryptkeys"
CRYPTTAB="$MNT/etc/crypttab"
mkdir -p "$MNT/etc"

tmp="$CRYPTTAB.tmp.$$"
cat > "$tmp" << 'EOF'
# /etc/crypttab: non-root encrypted volumes
# rawnix.org - https://rawnix.org
#
# <name>  <device>  <keyfile|none>  <options>
#
# Opened by /etc/runit/1 before fstab is mounted, closed by /etc/runit/3.
# The root volume is NOT listed here — the initramfs opens it.
#
# Generated by genluks. A keyfile named here must also be a real
# keyslot on that device; every volume additionally keeps the
# passphrase, which is the way in when root cannot be mounted.
#
# "none" prompts on the console at boot.  Add ",noauto" to skip a volume.
EOF

count=0
for s in $SPECS; do
    dev=$(spec_dev "$s"); target=$(spec_target "$s")
    [ "$(spec_plain "$s")" -eq 0 ] || continue
    [ "$target" != "/" ] || continue

    name=$(mapname "$target")
    uuid=$(cryptsetup luksUUID "$dev")

    # mkdir + chmod, not `install -d`: this script runs from the
    # installer image, where /bin/install is the netinstaller and PATH
    # is /sbin:/bin:/usr/bin.  A bare `install` re-executes that, which
    # reads -d as a group name and exits 1 -- and with `set -e` genluks
    # would stop here, after the volumes are formatted and mounted but
    # before a single keyfile or crypttab line exists.
    mkdir -p "$KDIR"
    chmod 700 "$KDIR"
    kf="$KDIR/$name.key"
    head -c 64 /dev/urandom > "$kf"
    chmod 400 "$kf"
    cryptsetup luksAddKey --key-file "$PASSFILE" "$dev" "$kf"

    printf '%-12s %-45s %-26s %s\n' \
        "$name" "UUID=$uuid" "/etc/cryptkeys/$name.key" "luks" >> "$tmp"
    count=$((count + 1))
done

if [ "$count" -gt 0 ]; then
    mv "$tmp" "$CRYPTTAB"
    chmod 600 "$CRYPTTAB"
else
    rm -f "$tmp"
fi

# ── Summary ───────────────────────────────────────────────────
echo ""
say "Done."
findmnt -l -R -n -o TARGET,SOURCE,FSTYPE "$MNT" | sed 's/^/    /'
if [ "$count" -gt 0 ]; then
    echo ""
    info "$count keyfile(s) in /etc/cryptkeys, listed in /etc/crypttab"
fi
echo ""
info "Next: mount the ESP, then setup-filesystem $MNT"
echo ""