#!/bin/sh

# Functions library :: for Linux Live Kit scripts
# Author: Tomas M. <http://www.linux-live.org>
#

# =================================================================
# debug and output functions
# =================================================================

debug_start()
{
   if grep -q debug /proc/cmdline; then
      DEBUG_IS_ENABLED=1
   else
      DEBUG_IS_ENABLED=
   fi
}

debug_log()
{
   if [ "$DEBUG_IS_ENABLED" ]; then
      echo "- debug: $*" >&2
      log "- debug: $*"
   fi
}

# header
# $1 = text to show
#
header()
{
   echo """$@"""
}


# echo green star
#
echo_green_star()
{
   echo -ne """* """
}

# log - store given text in /var/log/livedbg
log()
{
   echo "$@" 2>/dev/null >>/var/log/livedbg
}

echolog()
{
   echo "$@"
   log "$@"
}

# show information about the debug shell
show_debug_banner()
{
   echo
   echo "====="
   echo ": Debugging started. Here is the root shell for you."
   echo ": Type your desired commands or hit Ctrl+D to continue booting."
   echo
}

# debug_shell
# executed when debug boot parameter is present
#
debug_shell()
{
   if [ "$DEBUG_IS_ENABLED" ]; then
      show_debug_banner
      setsid sh -c 'exec sh < /dev/tty1 >/dev/tty1 2>&1'
      echo
   fi
}

fatal()
{
   echolog
   header "Fatal error occured - $1"
   echolog "Something went wrong and we can't continue. This should never happen."
   echolog "Please reboot your computer with Ctrl+Alt+Delete ..."
   echolog
   setsid sh -c 'exec sh < /dev/tty1 >/dev/tty1 2>&1'
}

# get value of commandline parameter $1
# $1 = parameter to search for
#
cmdline_value()
{
   cat /proc/cmdline | egrep -o "(^|[[:space:]])$1=[^[:space:]]+" | tr -d " " | cut -d "=" -f 2- | tail -n 1
}


# test if the script is started by root user. If not, exit
#
allow_only_root()
{
  if [ "0$UID" -ne 0 ]; then
     echo "Only root can run $(basename $0)"; exit 1
  fi
}

# Create bundle
# call mksquashfs with apropriate arguments
# $1 = directory which will be compressed to squashfs bundle
# $2 = output file
# $3..$9 = optional arguments like -keep-as-directory or -b 123456789
#
create_bundle()
{
   debug_log "create_module" "$*"
   rm -f "$2" # overwrite, never append to existing file
   mksquashfs "$1" "$2" -comp xz -b 512K $3 $4 $5 $6 $7 $8 $9>/dev/null
}


# Move entire initramfs tree to tmpfs mount.
# It's a bit tricky but is necessray to enable pivot_root
# even for initramfs boot image
#
transfer_initramfs()
{
   if [ ! -r /lib/initramfs_escaped ]; then
      echo "switch root from initramfs to ramfs"
      SWITCH=/m # one letter directory
      mkdir -p $SWITCH
      mount -t tmpfs -o size="100%" tmpfs $SWITCH
      cp -a /??* $SWITCH 2>/dev/null # only copy two-and-more-letter directories
      cd $SWITCH
      echo "This file indicates that we successfully escaped initramfs" > $SWITCH/lib/initramfs_escaped
      exec switch_root -c /dev/console . $0
   fi
}

# mount virtual filesystems like proc etc
#
init_proc_sysfs()
{
   debug_log "init_proc_sysfs" "$*"
   mount -n -t proc proc /proc
   echo "0" >/proc/sys/kernel/printk
   mount -n -t sysfs sysfs /sys
   mount -n -o remount,rw rootfs /
   ln -sf /proc/mounts /etc/mtab
}

# make sure some devices are there
init_devs()
{
   debug_log "init_devs" "$*"
   echo /sbin/mdev > /proc/sys/kernel/hotplug
   mdev -s
   modprobe zram 2>/dev/null
   modprobe loop 2>/dev/null
   modprobe squashfs 2>/dev/null
   modprobe fuse 2>/dev/null
}

# Activate zram (auto-compression of RAM)
# Compressed RAM consumes 1/2 or even 1/4 of original size
# Setup static size of 500MB
#
init_zram()
{
   debug_log "init_zram" "$*"
   echo_green_star
   echo "Setting dynamic RAM compression using ZRAM"
   echo 536870912 > /sys/block/zram0/disksize # 512MB
   mkswap /dev/zram0 >/dev/null
   swapon /dev/zram0 -p 32767
   echo 100 > /proc/sys/vm/swappiness
}

# load the AUFS kernel module if needed
#
init_aufs()
{
   debug_log "init_aufs" "$*"
   # TODO maybe check here if aufs support is working at all
   # and procude useful error message if user has no aufs
   modprobe aufs 2>/dev/null
}

# Setup empty union
# $1 = changes directory (ramfs or persistent changes)
# $2 = union directory where to mount the union
#
init_union()
{
   debug_log "init_union" "$*"

   echo_green_star
   echo "Setting up union using AUFS 3"
   mkdir -p "$1"
   mkdir -p "$2"
   mount -t aufs -o xino="/.xino",br="$1" aufs "$2"
}

# Return device mounted for given directory
# $1 = directory
#
mounted_device()
{
   debug_log "mounted_device" "$*"

   local MNT TARGET
   MNT="$1"
   while [ "$MNT" != "/" -a "$MNT" != "." -a "$MNT" != "" ]; do
      TARGET="$(grep -F " $MNT " /proc/mounts | cut -d " " -f 1)"
      if [ "$TARGET" != "" ]; then
         echo "$TARGET:$MNT"
         return
      fi
      MNT="$(dirname $MNT)"
   done
}

# Make sure to mount FAT12/16/32 using vfat
# in order to support long filenames
# $1 = device
#
device_bestfs()
{
   debug_log "device_bestfs" "$*"
   local FS

   FS="$(blkid "$1" | sed -r "s/.*TYPE=//" | tr -d '"' | tr [A-Z] [a-z])"
   if [ "$FS" = "msdos" -o "$FS" = "fat" -o "$FS" = "vfat" ]; then
      FS="vfat"
   elif [ "$FS" = "ntfs" ]; then
      FS="ntfs-3g"
   fi
   echo "-t $FS"
}

# Filesystem options for mount
# $1 = filesystem or '-t filesystem'
#
fs_options()
{
   debug_log "fs_options" "$*"

   if [ "$1" = "-t" ]; then
      shift
   fi
   if [ "$1" = "vfat" ]; then
      echo "-o check=s,shortname=mixed,iocharset=utf8"
   fi
}


# Find LIVEKIT data by mounting all devices
# If found, keep mounted, else unmount
# $1 = data directory target (mount here)
# $2 = data directory which contains compressed bundles
#
find_data_try()
{
   debug_log "find_data_try" "$*"

   local DEVICE FS MNT OPTIONS

   mkdir -p "$1"
   blkid | sort | cut -d: -f 1 | grep -E -v "/loop|/ram|/zram" | while read DEVICE; do
      FS="$(device_bestfs "$DEVICE")"
      OPTIONS="$(fs_options $FS)"
      mount -r "$DEVICE" "$1" $FS $OPTIONS 2>/dev/null
      if [ "$(find "$1/$2" -maxdepth 1 -name "*.$BEXT" 2>/dev/null)" != "" ]; then
         # we found at least one bundle/module here
         mount -o remount,rw "$DEVICE" "$1" 2>/dev/null
         echo "$1/$2" | tr -s "/"
         return
      fi
      umount "$1" 2>/dev/null
   done
}

# Retry finding LIVEKIT data several times,
# until timeouted or until data is found
# $1 = timeout
# $2 = data directory target (mount here)
#
find_data()
{
   debug_log "find_data" "$*"

   local DATA FROM

   FROM="$(cmdline_value from)"
   if [ "$FROM" = "" ]; then FROM="$LIVEKITNAME"; fi

   echo_green_star >&2
   echo -n "Looking for $LIVEKITNAME data in /$FROM/ .." | tr -s "/" >&2
   for timeout in $(seq 1 $1); do
      echo -n "." >&2
      DATA="$(find_data_try "$2" "$FROM")"
      if [ "$DATA" != "" ]; then
         echo "" >&2
         echo "* Found on $(mounted_device "$DATA" | cut -d : -f 1)" >&2
         echo "$DATA"
         return
      fi
      sleep 1
   done
   echo "" >&2

   if [ "$DATA" = "" ]; then
      fatal "$LIVEKITNAME data not found"
   fi

}

# Activate persistent changes
# $1 = data directory
# $2 = target changes directory
#
persistent_changes()
{
   debug_log "persistent_changes" "$*"

   local CHANGES T1 T2

   CHANGES="$1/$(basename "$2")"
   T1="$CHANGES/.empty"
   T2="$T1"2

   # Setup the directory anyway, it will be used in all cases
   mkdir -p "$2"

   # If persistent changes are not requested, end here
   if grep -vq perch /proc/cmdline; then
      return
   fi

   # check if changes directory exists and is writable
   touch "$T1" 2>/dev/null && rm -f "$T1" 2>/dev/null

   # if not, simply return back
   if [ $? -ne 0 ]; then
      echo "* Persistent changes not writable or not used"
      return
   fi

   echo_green_star
   echo "Testing persistent changes for posix compatibility"
   touch "$T1" && ln -sf "$T1" "$T2" 2>/dev/null && \
   chmod +x "$T1" 2>/dev/null && test -x "$T1" && \
   chmod -x "$T1" 2>/dev/null && test ! -x "$T1" && \
   rm "$T1" "$T2" 2>/dev/null

   if [ $? -ne 0 ]; then
      echo "* Activating dynamic sized storage for persistent changes"
      rm "$T1" "$T2" 2>/dev/null

      mount.dynfilefs "$CHANGES/changes.dat" 4000 "$2"
      if [ "$(device_bestfs "$2/loop.fs" | tr -d " ")" = "-t" ]; then
         mke2fs -F "$2/loop.fs" >/dev/null
      fi
      mount -o loop,sync "$2/loop.fs" "$2"
      rmdir "$2/lost+found"
   else
      echo "* Activating native persistent changes"
      mount --bind "$CHANGES" "$2"
   fi
}

# Copy data to RAM if requested
# $1 = live data directory
# $2 = changes directory
#
copy_to_ram()
{
   debug_log "copy_to_ram" "$*"

   local DM RAM CHANGES

   if grep -vq toram /proc/cmdline; then
      echo "$1"
      return
   fi

   CHANGES="$(basename $2)"
   DM="$(mounted_device "$1" | cut -d : -f 2-)"
   RAM="$DM.ram"

   echo "* Copying $LIVEKITNAME data to RAM..." >&2
   mkdir  -p "$RAM"
   cp -a "$DM/$LIVEKITNAME" "$RAM"
   echo "$RAM/$LIVEKITNAME"

   if grep -q perch /proc/cmdline; then
      umount "$2" 2>/dev/null
      umount "$DM/$LIVEKITNAME/$CHANGES" 2>/dev/null
      umount "$DM/$LIVEKITNAME/$CHANGES" 2>/dev/null
      mount --bind "$RAM/$LIVEKITNAME/$CHANGES" "$2"
   fi

   umount "$DM"
}

# Mount squashfs filesystem bundles
# and add them to union
# $1 = directory where to search for bundles
# $2 = directory where to mount bundles
# $3 = directory where union is mounted
#
union_append_bundles()
{
   debug_log "union_append_bundles" "$*"
   echo_green_star
   echo "Adding bundles to union"
   ls -1 "$1" | grep '.'$BEXT'$' | sort | while read BUNDLE; do
      echo "* $BUNDLE"
      mkdir -p "$2/$BUNDLE"
      mount -o loop -t squashfs "$1/$BUNDLE" "$2/$BUNDLE"
      mount -o remount,add:1:"$2/$BUNDLE" aufs "$3"
   done
}

# Create empty fstab properly
# $1 = root directory
#
fstab_create()
{
   debug_log "fstab_create" "$*"

   local FSTAB
   FSTAB="$1/etc/fstab"
   echo aufs / aufs defaults 0 0 > $FSTAB
   echo proc /proc proc defaults 0 0 >> $FSTAB
   echo sysfs /sys sysfs defaults 0 0 >> $FSTAB
   echo devpts /dev/pts devpts gid=5,mode=620 0 0 >> $FSTAB
   echo tmpfs /dev/shm tmpfs defaults 0 0 >> $FSTAB
}


# Change root and execute init
# $1 = where to change root
#
change_root()
{
   debug_log "change_root" "$*"

   umount /proc
   umount /sys
   rm -Rf /lib/modules # this will no longer be needed at all

   cd "$1"

   # make sure important device files and directories are in union
   mkdir -p boot dev proc sys tmp mnt run
   if [ ! -e dev/console ]; then mknod dev/console c 5 1; fi
   if [ ! -e dev/null ]; then mknod dev/null c 1 3; fi
   if [ ! -e sbin/fsck.aufs ]; then ln -s /bin/true sbin/fsck.aufs; fi

   # find chroot and init
   if [ -x bin/chroot ]; then  CHROOT=bin/chroot; fi
   if [ -x sbin/chroot ]; then  CHROOT=sbin/chroot; fi
   if [ -x usr/bin/chroot ]; then  CHROOT=usr/bin/chroot; fi
   if [ -x usr/sbin/chroot ]; then CHROOT=usr/sbin/chroot; fi
   if [ "$CHROOT" = "" ]; then fatal "Can't find executable chroot command"; fi

   if [ -x bin/init ]; then INIT=bin/init; fi
   if [ -x sbin/init ]; then INIT=sbin/init; fi
   if [ "$INIT" = "" ]; then fatal "Can't find executable init command"; fi

   mkdir -p mnt/live
   mount -n -o remount,ro aufs .
   pivot_root . mnt/live
   exec $CHROOT . $INIT < dev/console > dev/console 2>&1
}