# Functions commonly used by fripost-partman shell scripts
#
# Copyright 2013 Guilhem Moulin <guilhem@fripost.org>
#
# Licensed under the GNU GPL version 3 or higher.

true ${DEBIAN_HAS_FRONTEND:=} ${DEBCONF_REDIR:=}
. /usr/share/debconf/confmodule

# TODO:
# * support sofware RAID (mdadm)
# * encryption without SSH daemon (eg when /home is on a separate device)
# * encryption of multiple disks with the same key (eg / on a software raid)


# Log a message.
log() {
    logger -t fripost -p user.info -- "$@"
}

# Log a fatal error and quit.
fatal() {
    logger -t fripost -p user.err -- "$@"
    exit 1
}

# Ensure stdout is opened with line buffering. If some day stdbuf(1) is
# available in busybox, we should replace the LD_PRELOAD by 'stdbuf -oL -eL'.
stdbuf() {
    LD_PRELOAD=/lib/fripost-partman/stdbuf.so "$@"
}

# Wait for a device to be created
wait_for_device() {
    local device="$1" timeout=40

    until [ -b "$device" -o $timeout -le 0 ]; do
        log "Waiting for ${device}..."
        sleep 0.25
        timeout=$(( $timeout - 1 ))
    done
    [ -b "$device" ] || fatal "Error: $device not found!"
}


##############################################################################
# Wipe the disk (unless d-i's 'fripost/wipe-device' is 'none')
#
# Usage: fripost_wipe <device>

fripost_wipe() {
    local device="$1" source bs=8192 blockdir size wipe fifo pid
    blockdir="/sys/block/${device#/dev/}"

    db_get fripost/wipe-device
    if [ "$RET" = none ]; then
        log "Not wiping $device"
        return 0
    fi

    source="/dev/$RET"
    log "Want to wipe $device using source $source"
    [ -b "$device" -a -d "$blockdir" -a -c "$source" ] || \
        fatal "Invalid device $device or source $source"

    size=$(( `cat $blockdir/size` * `cat $blockdir/queue/physical_block_size` ))

    wipe="dd if=$source of=$device bs=$bs"
    db_subst fripost/wipe-device_progress_title DISK "$device"
    db_subst fripost/wipe-device_progress_title SIZE "$size"
    if [ "$source" = /dev/zero ]; then
        db_subst fripost/wipe-device_progress_title WHAT "zeroes"
    elif [ "$source" = /dev/random -o "$source" = /dev/urandom ]; then
        db_subst fripost/wipe-device_progress_title WHAT "bytes of random data"
    else
        db_subst fripost/wipe-device_progress_title WHAT "bytes"
    fi
    db_progress START 0 $(( $size / $bs )) fripost/wipe-device_progress_title
    db_subst fripost/wipe-device_progress_info COMMAND "$wipe"
    db_progress INFO fripost/wipe-device_progress_info

    fifo=$(mktemp -u)
    /bin/mknod "$fifo" p || exit 1
    trap 'kill $pid' EXIT
    $wipe 2> "$fifo" & pid=$!
    heartbeat $pid USR1 &

    local n records dir
    while read -u 7 n records dir; do
        [ "$records" = records -a "$dir" = out ] && db_progress SET ${n%+*}
    done 7< "$fifo"
    db_progress SET $(( $size / $bs )); sleep 0.25

    rm -f "$fifo"
    trap '' EXIT

    db_progress STOP
    db_unregister fripost/wipe-device_progress_title
    db_unregister fripost/wipe-device_progress_info
}

heartbeat () {
    local pid=$1 sig=${2:-SIGHUP} sleep=${3:-1}
    until false; do
        sleep $sleep
        /bin/kill -$sig $pid 2>/dev/null || break
    done
}



##############################################################################
# Partitioning
#
# Usage: fripost_mkpart <device> <name> <size in bytes or %> [flag | +flag | -flag | ...]
#
# Prints the bloc device for the new partition. <size> can be given
# using SI prefixes T,G,M, or k. If <size> is a percentage, it is
# relative to the free capacity.

fripost_mkpart() {
    local device="$1" name="$2" size="$3" grain=$(( 256*32 )) offset bs
    local start stop part flag free
    # TODO: names are supported on GPT only. What is the second argument
    # of mkpart on msdos?

    bs=$(cat /sys/block/${device#/dev/}/queue/physical_block_size)

    if /sbin/parted -sm "$device" p | grep -q '^[1-9][0-9]*:'; then
        # The last offset, in sector; it is *always* supposed to be aligned.
        start=$( /sbin/parted -sm "$device" u s p \
               | grep '^[1-9][0-9]*:' \
               | sort -t: -nk2,2 \
               | tail -1 \
               | sed -r 's/^[1-9][0-9]*:[^:]*:([1-9][0-9]*)s:.*/\1/' )
        start=$(( 1 + $start ))
    else
        # There are no partitions yet
        offset=$(cat /sys/block/${device#/dev/}/alignment_offset)
        if [ $offset -eq 0 ]; then
            # 64 is arbitrary here; $grain may be better since aligned,
            # but the 1st partition is typically /boot or something
            # without heavy IO, so choosing a smaller value means we get
            # more space without significant performance penalty
            start=64
        else
            start=$(( $offset / $bs ))
        fi
    fi

    # Remove the SI prefix, or compute the size relative to a percentage
    case "$size" in
        *T) size=$(( ${size%T} * 1024**4 ));;
        *G) size=$(( ${size%G} * 1024**3 ));;
        *M) size=$(( ${size%M} * 1024**2 ));;
        *K) size=$(( ${size%K} * 1024    ));;
        *%) free=$(( `cat /sys/block/${device#/dev/}/size` - 1 - $start ))
            size=$(( $bs * $free * ${size%\%} / 100 ))
        ;;
        *) fatal "Unknown size $size" ;;
    esac

    # Round toward 0 to preserve the alignement
    stop=$(( $start + $size / $bs ))
    stop=$(( $stop - $stop % $grain - 1 ))
    log "Creating partition '$name' on $device, between sectors $start and $stop"
    [ $start -ge $stop ] && fatal "Cannot create partition: $start >= $stop"
    /sbin/parted -a minimal -s $device mkpart "$name" ${start}s ${stop}s

    # The partition index that is to be created
    part=$( /sbin/parted -sm $device u s p \
          | grep -m1 "^[0-9]:${start}s:" \
          | sed 's/:.*//' )

    while [ $# -gt 3 ]; do
        flag="$4"
        log "Partition flag for ${device}${part} ($name): $flag"
        case "$flag" in
            +*) /sbin/parted -s $device set $part "${flag#+}" on  ;;
            -*) /sbin/parted -s $device set $part "${flag#-}" off ;;
            *)  /sbin/parted -s $device toggle $part "$flag" ;;
        esac
        shift
    done

    wait_for_device "${device}${part}"
    echo ${device}${part}
}



##############################################################################
# Encryption
#
# Usage: fripost_encrypt <device> <name> [<cryptsetup options>]

fripost_encrypt() {
    local device="$1" name="$2" keyfile sshHostKey import=/cdrom/include
    local fifo pid
    shift; shift
    keyfile=$(mktemp) || exit 1

    log "Encryting device $device and sets up a mapping $name"

    db_input high fripost/encryption-password || true
    db_go
    db_get fripost/encryption-password

    if [ -n "$RET" ]; then
        log "The password can be found from debconf (hey, that's insecure!)"
        # Since we're using the builtin, the password shouldn't be
        # visible in ps(1).
        echo -n "$RET" > "$keyfile"
        db_set fripost/encryption-password ''

    else
        log "Starting a SSH daemon to slurp the volume key..."
        local type=rsa
        db_subst fripost/ssh-keypair-generation_progress_title TYPE $type
        db_progress START 0 1 fripost/ssh-keypair-generation_progress_title

        mkdir -pm0755 /etc/ssh/
        sshHostKey=/etc/ssh/ssh_host_rsa_key
        /usr/bin/ssh-keygen -b 4096 -t $type -N '' -C $sshHostKey -f $sshHostKey
        db_progress SET 1; sleep 0.25
        db_progress STOP
        db_unregister fripost/ssh-keypair-generation_progress_title

        cat > /etc/ssh/sshd_config <<- EOF
			Port 22
			Protocol 2
			HostKey $sshHostKey
			UsePrivilegeSeparation no

			PasswordAuthentication no
			ChallengeResponseAuthentication no
			HostbasedAuthentication no
			PubkeyAuthentication yes

			PermitRootLogin yes
			AllowUsers root
			StrictModes yes

			AllowAgentForwarding no
			AllowTcpForwarding no
			ForceCommand /bin/cat >$keyfile
		EOF

        # Populate the authorized keys.
        [ -d ~root/.ssh ] || mkdir -m0700 ~root/.ssh
        copy_authorized_keys $import/authorized_keys ~root/.ssh/authorized_keys 'no-pty'

        # Start the SSH daemon
        touch /var/log/lastlog
        /usr/sbin/sshd

        # Tell the user we're ready
        db_subst fripost/encryption-slurpkey_text IPv4 "$(getIPv4)"
        db_subst fripost/encryption-slurpkey_text SSHFPR_SERVER \
            "$(/usr/bin/ssh-keygen -lf $sshHostKey)"
        db_subst fripost/encryption-slurpkey_text SSHFPR_AUTHORIZED \
            "$(sshfprs ~root/.ssh/authorized_keys '    - ')"

        # Anything sent to the SSH is stored into $keyfile, which is our LUKS key.
        until test -s "$keyfile"; do
            db_settitle fripost/encryption-slurpkey_title
            db_input critical fripost/encryption-slurpkey_text
            db_go
        done

        # Kill the daemon, we're done
        log "Killing the SSH daemon..."
        kill `cat /var/run/sshd.pid` || true
    fi

    # Load dm-crypt
    /sbin/dmsetup targets | grep -q '^crypt\s' || /sbin/modprobe -v dm_crypt

    log "cryptsetup options: $@"
    db_progress START 0 100 fripost/cryptsetup-genkey_progress_title
    db_progress INFO fripost/cryptsetup-genkey_progress_info

    fifo=$(mktemp -u)
    /bin/mknod "$fifo" p || exit 1
    trap 'kill $pid' EXIT
    stdbuf /sbin/cryptsetup -q "$@" --key-file="$keyfile" luksFormat "$device" \
    | stdbuf sed -nr 's/^Generating key \(([0-9]+)% done\)\.$/\1/p' > "$fifo" & pid=$!

    local n
    while read -u 7 n; do
        db_progress SET $n
    done 7< "$fifo"
    db_progress SET 100; sleep 0.25

    rm -f "$fifo"
    trap '' EXIT

    db_progress STOP
    db_unregister fripost/cryptsetup-genkey_progress_title
    db_unregister fripost/cryptsetup-genkey_progress_info

    # Open the LUKS device and set up the mapping
    /sbin/cryptsetup --key-file="$keyfile" luksOpen "$device" "$name"

    log "Removing the key file"
    rm -f "$keyfile" # We are on a ramdisk, so it's good enough to unlink(2)

    # Add an entry to the crypttab
    printf "%-15s %-41s %-15s %s\n" \
        "$name" UUID=$(/bin/block-attr --uuid "$device") none luks \
        >> /tmp/crypttab

    local m _
    [ -d /var/lib/fripost ] || mkdir /var/lib/fripost
    # The modules required to fire up dropbear and start cryptsetup in the ramdisk.
    echo dm_crypt > /tmp/initramfs-modules
    while read m _; do /sbin/modinfo -F filename "$m"; done < /proc/modules \
        | sed -nr "s@^/lib/modules/`uname -r`/kernel/((arch/[^/]+/)?crypto|drivers/(ata|scsi))(/.*)?/([^/]+)\.ko\$@\5@p" \
        >> /var/lib/fripost/initrd-modules
    /bin/apt-install busybox cryptsetup || true

    wait_for_device "/dev/mapper/$name"
}


# Like ssh-keygen -lf, but for a file such as authorized_keys, which
# may contain multiple keys.
#
# Usage: sshfprs.sh file [prefix]

sshfprs() {
    local file="$1" prefix="${2:-}" pk

    while read pk; do
        # /usr/bin/ssh-keygen can't read from STDIN, and the '<<<' is
        # not POSIX, so we save each pubkey in a temporary file
        pkf=$(mktemp)
        echo "$pk" > "$pkf"
        echo "${prefix}$(/usr/bin/ssh-keygen -lf $pkf)"
        rm -f "$pkf"
    done < "$file"
}

# Copy an authorized_keys file, possibly adding some options. The input
# format is sshd(8)'s; see the section called "AUTHORIZED_KEYS FILE
# FORMAT" in the manpage. The character '#' cannot occur in the options.

copy_authorized_keys() {
    local from="$1" to="$2"
    if [ $# -gt 2 ]; then
        sed -r "s#^([^#]+\s)?(ssh-(dss|rsa)|ecdsa-sha2-nistp(256|384|521))\s#$3 \2 #" \
            "$from" > "$to"
    else
        cp "$from" "$to"
    fi
    chmod 0600 "$to"
}



##############################################################################
# Add a progress bar to mkfs(8).
#
# Usage: fripost_mkfs <type> <device> [<options>]

fripost_mkfs() {
    local type="$1" device="$2" fifo pid line stage step total=
    shift; shift

    log "Formatting $device as $type; options $@"
    db_subst fripost/mkfs_progress_title DEVICE "$device"
    db_subst fripost/mkfs_progress_title TYPE   "$type"
    wait_for_device "$device"

    fifo=$(mktemp -u)
    /bin/mknod "$fifo" p || exit 1
    trap 'kill $pid' EXIT
    /sbin/mkfs."$type" "$@" "$device" | tr '\b' '\n' > "$fifo" & pid=$!

    while read -u 7 line; do
        # XXX: The parsing has been tested with mkfs.ext{2,3,4} only.
        case "$line" in
            *:*/*)
                [ -z "$total" ] || db_progress STOP

                stage="${line%: *}"
                step="${line##*: }"
                log "mkfs(8) stage '$stage', $step steps"

                total="${step#*/}"
                db_progress START 0 $total fripost/mkfs_progress_title
                db_progress SET ${step%/*}

                db_subst fripost/mkfs_progress_info STAGE "$stage"
                db_progress INFO fripost/mkfs_progress_info
            ;;
            */$total)
                [ -z "$total" ] || db_progress SET ${line%/*}
            ;;
            done)
                [ -z "$total" ] || db_progress SET $total
            ;;
        esac
    done 7< "$fifo"

    rm -f "$fifo"
    trap '' EXIT

    [ -z "$total" ] || db_progress STOP
    db_unregister fripost/mkfs_progress_title
    db_unregister fripost/mkfs_progress_info
}


##############################################################################
# Set up the fstab
#
# Usage: fripost_fstab <device> <mountpoint> <type> [<options>]
#
# Note: Don't forget to run 'fripost_mount_partitions' after the last
# call to 'fripost_fstab'.

fripost_fstab() {
    local device="$1" mp="$2" type="$3" options="${4:-defaults}" pass

    # TODO: EFI should be set with pass 1
    if [ "$mp" = / ]; then
        log "Got root device, $device"
        pass=1
    elif [ "$type" = swap ]; then
        log "Got swap device, $device"
        pass=0
    else
        log "Got other device, $device -> $mp"
        pass=2
    fi

    # Add an entry to the fstab
    fripost_fstab_addentry $device $mp $type $options 0 $pass
}

fripost_fstab_addentry() {
    local device=$1 mp=$2 type=$3 options=$4 dump=$5 pass=$6
    local fstab=/tmp/fstab
    local entry="%-41s %-15s %-7s %-31s %-7s %s\n"
    log "Adding fstab entry $device -> $mp"

    if ! [ -s $fstab -a -f $fstab ]; then
        # No fstab has bee created yet. We copy the header from
        # /lib/partman/finish.d/39create_fstab_header
        cat > $fstab <<-EOF
			# /etc/fstab: static file system information.
			#
			# Use 'blkid' to print the universally unique identifier for a device;
			# this may be used with UUID= as a more robust way to name devices that
			# works even if disks are added and removed. See fstab(5).
			#
		EOF
        printf "$entry" '# <file system>' '<mount point>' '<type>' \
                        '<options>' '<dump>' '<pass>' \
            >> $fstab
    fi

    # Always use UUIDs
    local uuid=$(/bin/block-attr --uuid "$device")
    if [ -z "$uuid" ]; then
        log "Warning: Block device $device has no UUID!"
        printf "$entry" "$device" $mp $type $options $dump $pass >> $fstab
    else
        log "Block device $device has UUID $uuid"
        printf "$entry" "UUID=$uuid" $mp $type $options $dump $pass >> $fstab

        # Fix udev's UUID mapping
        if ! [ -h /dev/disk/by-uuid/"$uuid" ]; then
            device=$(/bin/mapdevfs "$device")
            log "Adding a symlink UUID $uuid -> $device"
            ln -s ../../"${device#/dev/}" /dev/disk/by-uuid/"$uuid"
        fi
    fi
}



##############################################################################
# Mount all entries in the fstab
#
# Usage: fripost_mount_partitions

fripost_mount_partitions() {
    local cdrom uuid mp type options _

    # sort(1), to ensure that roots are mounted earlier, regardless of
    # where they appear in the file (eg, /boot can appear before /).
    sed -e 's/#.*//' -e '/^\s*$/d' -e 's/\s\s*/ /g' /tmp/fstab \
    | grep -v '\s0$' \
    | sort -t' ' -k2,2 \
    | while read uuid mp type options _; do
        local dev=$(/sbin/blkid -U "${uuid#UUID=}")
        log "Found fstab entry: $dev -> $mp"
        wait_for_device "$dev"
        if [ "$type" = swap ]; then
            /sbin/swapon "$dev"
        elif [ / = $(expr substr "$mp" 1 1) ]; then
            mkdir "/target$mp"
            /bin/mount -t "$type" -o "$options" "$dev" "/target$mp"
        fi
      done

    # Post-installation scripts may use the install CD as a local mirror
    # for APT packages. TODO: What if it was a USB stick? A netboot?
    cdrom=$( sed -rn 's#^(\S+) /cdrom.*#\1#p' /proc/mounts \
           | grep -m1 -Ev '^/dev/(loop|sd[a-z])' )
    [ -n "$cdrom" ] || fatal "Error: Is /cdrom a mountpoint?"
    fripost_fstab_addentry "$cdrom" /media/cdrom0 udf,iso9660 user,noauto 0 0
    mkdir -p /target/media/cdrom0
    ln -s cdrom0 /target/media/cdrom

    # Move the fstab and crypttab where they belong
    [ -d /target/etc ] || mkdir /target/etc
    mv /tmp/fstab /target/etc/fstab
    [ \! -f /tmp/crypttab ] || mv /tmp/crypttab /target/etc/crypttab
}

getIPv4() {
    /bin/ip addr show "${1:-eth0}" | sed -nr 's/^\s+inet\s([0-9.]{4,32}).*/\1/p'
}