#!/bin/bash
# Copyright (C) 2016 Ian Kelling
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

script=$(readlink -f -- "$BASH_SOURCE")
cd /
[[ $EUID == 0 ]] || exec sudo -E "$script" "$@"

source /usr/local/lib/err

usage() {
  cat <<EOF
Usage: ${0##*/} [OPTIONS]

-h|--help  Print help and exit.


Note, at source location, intentionally not executable, run and read
install-my-scripts.

Note: Uses util-linux getopt option parsing: spaces between args and
options, short options can be combined, options before args.
EOF
  exit $1
}


tu() {
  while read -r line; do
    file="$1"
    grep -xFq "$line" "$file" || tee -a "$file"<<<"$line"
  done
}
m() {
  if $verbose; then
    printf "%s\n" "$*"
  fi
  "$@"
}
x() {
  printf "%s\n" "$*"
  "$@"
}

mnt() {
  dir=$1
  if ! mountpoint -q $dir; then
    mkdir -p $dir
    m mount $dir
  fi
}
fstab() {
  while read -r start mpoint end; do
    l="$start $mpoint $end"
    # kill off any lines that duplicate the mount point.
    sed --follow-symlinks -ri "\%$l%b;\%^\s*\S+\s+$mpoint\s%d" /etc/fstab
    tu /etc/fstab <<<"$l"
  done
}
pid-check() {
  for p in ${pids}; do
    for m in ${my_pids[@]}; do
      if (( p == m )); then
        echo "$0: error: pids to kill includes our pid or a parent. ps output:" >&2
        ps -f -p $p
        exit 1
      fi
    done
  done
}
kill-dir() {
  for sig; do
    echo kill-dir $sig
    found_pids=false
    if pids=$(timeout 4 lsof -t $dir); then
      found_pids=true
      timeout 4 lsof -w $dir
      pid-check
      kill -$sig $pids
    fi
    # fuser will find open sockets that lsof won't, for example from gpg-agent.
    # note: -v shows kernel processes, which then doesn't return true when we want
    if pids=$(timeout 4 fuser -m $dir 2>/dev/null); then
      pid-check
      found_pids=true
      fuser -$sig -mvk $dir
    fi
    sleep .5
    if ! $found_pids; then
      return 0
    fi
  done
  return 1
}

##### begin command line parsing ########

# you can remove this if you do not have options which can have args with spaces or empty.

verbose=false
force=false
temp=$(getopt -l help,force,verbose hfv "$@") || usage 1
eval set -- "$temp"
while true; do
  case $1 in
    -f|--force) force=true ;;
    -v|--verbose) verbose=true ;;
    -h|--help) usage ;;
    --) shift; break ;;
    *) echo "$0: unexpected args: $*" >&2 ; usage 1 ;;
  esac
  shift
done

##### end command line parsing ########

ret=0

##### begin setup fstab for subvols we care about ######
root_dev=$(awk '$2 == "/" {print $1}' /etc/mtab)
if [[ $root_dev == /dev/dm-* ]]; then
  for d in /dev/mapper/*; do
    if [[ $(readlink -f $d) == "$root_dev" ]]; then
      root_dev=$d
      break
    fi
  done
fi

if cryptsetup status $root_dev &>/dev/null; then
  crypt_dev=$root_dev
else # if we are in a recovery boot, find the next best crypt device
  noauto=,noauto
  for dev in $(dmsetup ls --target crypt | awk '{print $1}'); do
    dev=/dev/mapper/$dev
    if awk '{print $1}' /etc/mtab | grep -Fx $dev &>/dev/null; then
      crypt_dev=$dev
      break
    fi
  done
fi


fstab <<EOF
$crypt_dev  /a  btrfs  noatime,subvol=a$noauto  0 0
EOF

shopt -s nullglob

# ssh and probably some other things care about parent directory
# ownership, and ssh doesn\'t allow any group writable parent
# directories, so we are forced to use a directory structure similar
# to home directories
f=(/mnt/root/btrbk/q.*); f=${f[0]}
if [[ -e $f ]]; then
  fstab <<EOF
$crypt_dev  /q  btrfs  noatime,subvol=q,gid=1000$noauto  0 0
/q/p  /p  none  bind$noauto  0 0
EOF
fi

f=(/mnt/root/btrbk/o.*); f=${f[0]}
if [[ -e $f ]]; then
  fstab <<EOF
$crypt_dev  /o  btrfs  noatime,subvol=o$noauto  0 0
/o/m  /m  none  bind$noauto  0 0
EOF
fi

if [[ $HOSTNAME == frodo ]]; then
  fstab <<EOF
$crypt_dev  /i  btrfs  noatime,subvol=i$noauto  0 0
EOF
fi
##### end setup fstab for subvols we care about ######

# get pids that this program depends on so we dont kill them
my_pids=($$ $PPID)
loop_limit=30
count=0
while [[ ${my_pids[-1]} != 1 && ${my_pids[-1]} != ${my_pids[-2]} && $count -lt $loop_limit  ]]; do
  count=$((count + 1))
  p=$(ps -p ${my_pids[-1]} -o ppid=)
  if [[ $p == 0 || ! $p ]]; then
    break
  fi
  my_pids+=($p)
done


for vol in q a o i; do
  d=/$vol
  if ! awk '{print $2}' /etc/fstab | grep -xF $d &>/dev/null; then
    continue
  fi


  ##### begin building up list of bind mounts  ######
  binds=() # list of bind mounts
  roots=($d) # list of bind mounts, plus the original mount
  while true; do
    new_roots=()
    for r in ${roots[@]}; do
      # eg. when r=/q/p, for lines like
      # /q/p  /p  none  bind  0 0
      # output /p
      new_roots+=($(sed -rn "s#^$r/\S+\s+(\S+)\s+none\s+bind\s.*#\1#p" /etc/fstab))
    done
    (( ${#new_roots} )) || break
    binds+=(${new_roots[@]})
    roots=( ${new_roots[@]} )
  done
  ##### end building up list of bind mounts  ######


  # if latest is already mounted, make sure binds are mounted and move on
  m check-subvol-stale $d
  # populated by check-subvol-stale if stale
  if ! fresh_snap=$(cat /nocow/btrfs-stale/$vol 2>/dev/null); then
    mnt $d
    for b in ${binds[@]}; do
      mnt $b
    done
    continue
  fi

  umount_ret=true
  unmounted=()
  for dir in $(echo $d ${binds[*]}\ |tac -s\ ); do
    if mountpoint -q $dir; then
      if m umount -R $dir; then
        unmounted+=($dir)
      else
        if ! kill-dir TERM TERM TERM INT INT HUP HUP; then
          if $force; then  kill-dir KILL; fi
        fi

        if m umount -R $dir; then
          unmounted+=($dir)
        else
          echo "$0: failed to umount $dir"
          umount_ret=false
          ret=1
          continue
        fi
      fi
    fi
  done

  # if we unmounted some but not all, restore them and move on
  if ! $umount_ret; then
    for dir in ${unmounted[@]}; do
      mnt $dir
    done
    continue
  fi

  #### begin dealing with leaf vols ####
  # todo: decipher /mnt/root, like we do in check-subvol-stale
  cd /mnt/root
  if [[ -e $vol ]]; then
    leaf=$vol.leaf.$(date +%Y-%m-%dT%H:%M:%S%z)
    m mv $vol $leaf
    m btrfs property set -ts $leaf ro true

    ### begin check if leaf is different, delete it if not ###
    if [[ -e /a/opt/btrfs-snapshots-diff/btrfs-snapshots-diff.py ]]; then
      source /a/bin/distro-functions/src/package-manager-abstractions
      pi python-jmespath # dependency
      parentid=$(btrfs sub show $leaf | awk '$1 == "Parent" && $2 == "UUID:" {print $3}')
      bsubs=(/mnt/root/btrbk/$vol.*)
      bsub=
      # go in reverse order as its more likely to be at the end
      for ((i=${#bsubs[@]}-1; i>=0; i--)); do
        if [[ $parentid == $(btrfs sub show ${bsubs[i]} | awk '$1 == "UUID:" {print $2}') ]]; then
          bsub=${bsubs[i]}
          break
        fi
      done
      if [[ $bsub ]]; then
        tmp=$(mktemp)
        # in testing, same subvol is 136 bytes. allow some overhead
        btrfs send --no-data -p $bsub $leaf | head -c 1000 > $tmp || [[ $? == 141 ]]
        if (( $(stat -c%s $tmp) < 1000)); then
          # example output for an empty diff:
          # Found a valid Btrfs stream header, version 1
          # o.leaf.2019-05-15T14:00:50-0400;snapshot: uuid=ba045ea30737dd449003f1ee40ec12d0, ctrasid=109533, clone_uuid=3c7e3544e486834aa71d89e5b8f30056, clone_ctransid=109533
          lines=$(/a/opt/btrfs-snapshots-diff/btrfs-snapshots-diff.py -s -f $tmp | \
                    grep -vxF "Found a valid Btrfs stream header, version 1" | \
                    grep -cv "^[^;]*;snapshot: ") ||:
          if [[ $lines == 0 ]]; then
            x btrfs sub del $leaf
          fi
        fi
      fi
    fi
    ### end check if leaf is different, delete it if not ###

    ## begin expire leaf vols ##
    leaf_vols=($vol.leaf.*)
    count=1
    for leaf in ${leaf_vols[@]}; do
      leaf_secs=$(date -d ${leaf#$vol.leaf.} +%s)
      if (( $(date +%s) - 60*60*24*60 > leaf_secs || count > 200 )); then # 60 days
        x btrfs sub del $leaf
      fi
      count=$((count+1))
    done
    ## end expire leaf vols ##
  fi
  #### end dealing with leaf vols ####

  # Note, we make a few assumptions in this script, like
  # $d was not a different subvol id than $vol, and
  # things otherwise didn't get mounted very strangely.
  m btrfs sub snapshot $fresh_snap $vol
  for dir in $d ${binds[@]}; do
    m mnt $dir
  done
  stale_dir=/nocow/btrfs-stale
  rm -f $stale_dir/$d

done

### disabled
if [[ $HOSTNAME == kdxxxxxxxxx ]]; then
  # partitioned it with fai partitioner outside of fai,
  # because it\'s worth it to have 1% space reserved for boot and
  # swap partitions in case I ever want to boot off those drives.
  # as root:
  # . /a/bin/fai/fai-wrapper
  # eval-fai-classfile /a/bin/fai/fai/config/class/51-multi-boot
  # fai-setclass ROTATIONAL
  # export LUKS_DIR=/q/root/luks/
  # # because the partition nums existed already
  # fai-setclass REPARTITION
  # /a/bin/fai/fai/config/hooks/partition.DEFAULT

  devs=(
    ata-TOSHIBA_MD04ACA500_84REK6NTFS9A-part1
    ata-TOSHIBA_MD04ACA500_84R2K773FS9A-part1
    ata-TOSHIBA_MD04ACA500_8471K430FS9A-part1
    ata-TOSHIBA_MD04ACA500_8481K493FS9A-part1
  )
  first=true
  for dev in ${devs[@]}; do
    if $first; then
      first=false
      tu /etc/fstab <<EOF
/dev/mapper/crypt_dev_$dev /i btrfs  noatime,subvol=i,noauto  0 0
/dev/mapper/crypt_dev_$dev /mnt/iroot btrfs  noatime,subvolid=0,noauto  0 0
EOF
    fi
    tu /etc/crypttab <<EOF
crypt_dev_$dev  /dev/disk/by-id/$dev  /q/root/luks/host-kd  discard,luks
EOF
    if [[ ! -e /dev/mapper/crypt_dev_$dev ]]; then
      cryptdisks_start crypt_dev_$dev
    fi
  done
  # note, could do an else here and have some kind of mount for /i
  # on other hosts.
fi

exit $ret