#!/bin/sh
#
# linbo_create_image
# thomas@linuxmuster.net
# 20260518
#


#### functions begin ####

# usage info
usage(){
  local RC="$1"
  echo
  echo "Creates a base or differential image (option \"qdiff\") of an operating system"
  echo "either defined by start.conf postition number or by root partition."
  echo
  echo "Usage:"
  echo "  linbo_create_image <#> [qdiff] | [help]"
  echo "  linbo_create_image <root> [qdiff] | [help]"
  echo
  echo "For compatibility reasons legacy options are also accepted. In this case"
  echo "the image type is defined by the filename:"
  echo "  linbo_create_image <cache> <imagefilename> <root>"
  echo
  exit "$RC"
}


# save_efi_bcd targetdir efipart
# saves the windows efi file to os partition
save_efi_bcd(){
  local targetdir="$1"
  local efipart="$2"
  local efimnt="/cache/boot/efi"
  mkdir -p "$efimnt"
  mount "$efipart" "$efimnt" || return 1
  local sourcedir="$(ls -d "$efimnt"/[Ee][Ff][Ii]/[Mm][Ii][Cc][Rr][Oo][Ss][Oo][Ff][Tt]/[Bb][Oo][Oo][Tt] 2> /dev/null)"
  if [ -z "$sourcedir" ]; then
    echo "$sourcedir not found. Cannot copy windows efi bootfiles to $targetdir."
    umount "$efimnt" || umount -l "$efimnt"
    return 1
  fi
  echo "Copying windows efi bootfiles from $sourcedir to $targetdir."
  mkdir -p "$targetdir"
  local RC=0
   rsync --skip-compress="$RSYNC_SKIP_COMPRESS" -r "$sourcedir/" "$targetdir/" || RC="1"
  umount "$efimnt" || umount -l "$efimnt"
  [ "$RC" = "1" ] && echo "Copying of windows efi bootfiles failed!"
  return "$RC"
}


# prepare_fs directory partition
# Removes all files from /etc/rsync.exclude and saves win boot configuration in
# the root directory of the os.
prepare_fs(){
  (
    # remove excluded files
    cd "$1" || return 1
    local item
    grep -v ^# /etc/rsync.exclude | while read item; do
      if ls $item &> /dev/null; then
        echo -n "Removing $item ... "
        rm -rf $item
        echo "done."
      fi
    done
    # save win bcd & mbr
    local targetdir
    # in case of efi save the windows efi files
    local efipart="$(print_efipart)"
    if [ -n "$efipart" ]; then
      # save partition & disk uuids
      echo "Saving partition guids."
      print_guid "$efipart" > /mnt/.guid.efi
      print_guid "$2" > /mnt/.guid.part
      local disk="$(get_disk_from_partition "$2")"
      print_guid "$disk" > /mnt/.guid.disk
      if [ "$(linbo_fstype $2)" = "ntfs" ]; then
        targetdir="$(ls -d [Ee][Ff][Ii]/[Mm][Ii][Cc][Rr][Oo][Ss][Oo][Ff][Tt]/[[Bb][Oo][Oo][Tt] 2> /dev/null)"
        [ -z "$targetdir" ] && targetdir="EFI/Microsoft/Boot"
        save_efi_bcd "$targetdir" "$efipart"
      fi
    else
      targetdir="$(ls -d [Bb][Oo][Oo][Tt] 2> /dev/null)"
    fi
    if [ -n "$targetdir" ]; then
      local bcd="$(ls $targetdir/[Bb][Cc][Dd] 2> /dev/null)"
      if [ -n "$bcd" -a -n "$HOSTGROUP" ]; then
        echo "Saving windows bcd for group $group."
        # BCD group specific and partition specific on efi systems
        if [ -n "$efipart" ]; then
          cp -f "$bcd" "$bcd"."$HOSTGROUP"."$(basename "$2")"
        else
          cp -f "$bcd" "$bcd"."$HOSTGROUP"
        fi
        # boot sector backup group specific
        # delete obsolete mbr backups
        rm -f "$targetdir/winmbr.$HOSTGROUP" "$targetdir/win7mbr.$HOSTGROUP" "$targetdir/winmbr446.$HOSTGROUP" "$targetdir/bsvbr.$HOSTGROUP"
        local disk="$(get_disk_from_partition "$2")"
        local bsmbr="$targetdir/bsmbr.$HOSTGROUP"
        # save bootloader sectors
        echo "Saving bootloader sectors for group $HOSTGROUP."
        dd if="$disk" of="$bsmbr" bs=446 count=1
        # ntfs partition id
        echo "Saving ntfs id."
        local ntfsid="$targetdir/ntfs.id"
        dd if="$2" of="$ntfsid" bs=8 count=1 skip=9
      fi
    fi
  )
}


# mk_baseimage rootdev imagefile
mk_baseimage(){
  local rootdev="$1"
  local imagefile="$2"
  echo "Filling up partition with zeroes..."
  freespace=$(df -kP $rootdev | awk 'NR==2 {print $4}')k
  interruptible pv -s $freespace -tprebf /dev/zero | dd of=/mnt/zero.tmp bs=1024k


  # Create nulled files of size 1GB, should work on any FS.
  sync ; sync ; sync
  rm -f /mnt/zero.tmp
  umount /mnt || umount -l /mnt
  echo "Starting compression of $rootdev -> $imagefile (full partition, ${size}K)."
  echo "qemu-img convert -p -c -f raw -O qcow2 $rootdev $imagefile"
  interruptible qemu-img convert -p -c -f raw -O qcow2 "$rootdev" "$imagefile" ; RC="$?"
  if [ "$RC" != "0" ]; then
    echo "Creation of $imagefile failed. :("
    return "$RC"
  fi
  # remove diffimage files in case of new baseimage
  rm -f "${imagefile/.qcow2/.qdiff}"*
}


# mk_diffimage rootdev imagefile
mk_diffimage(){
  local rootdev="$1"
  local imagefile="$2"
  local baseimage="${imagefile/.qdiff/.qcow2}"
  mkdir -p /image
  # test baseimage
  qemu-nbd --disconnect /dev/nbd0
  if ! qemu-nbd --connect /dev/nbd0 "$baseimage"; then
    echo "$baseimage is not connectable!"
    return 1
  fi
  sleep 3
  if ! mount /dev/nbd0 /image; then
    echo "$baseimage is not mountable!"
    qemu-nbd --disconnect /dev/nbd0
    return 1
  fi
  umount /image
  qemu-nbd --disconnect /dev/nbd0
  # create new diffimage
  if [ ! -e "$imagefile" ]; then
    local new="yes"
    if ! qemu-img create -f qcow2 -b "$baseimage" -F qcow2 "$imagefile"; then
      echo "Creation of $imagefile failed!"
      return 1
    fi
  fi
  # connect imagefile to nbd device
  if ! qemu-nbd --connect /dev/nbd0 "$imagefile"; then
    echo "$imagefile is not connectable!"
    [ -n "$new" ] && rm -f "$imagefile"
    return 1
  fi
  sleep 3
  # mount image filesystem
  if ! mount /dev/nbd0 /image; then
    qemu-nbd --disconnect /dev/nbd0
    echo "Fatal! $imagefile is not mountable"
    [ -n "$new" ] && rm -f "$imagefile"
    return 1
  fi
  echo "Start syncing differences of $rootdev -> $imagefile."
  interruptible rsync --skip-compress="$RSYNC_SKIP_COMPRESS" -HXAa --exclude="/.linbo" --exclude-from="/etc/rsync.exclude" --delete --delete-excluded  /mnt/ /image 2>&1 ; RC="$?"
  umount /image || umount -l /image
  qemu-nbd --disconnect /dev/nbd0
  [ "$RC" = "0" ] && return
  echo "Creation of $imagefile failed. :("
  rm -f "$imagefile"*
  return "$RC"
}


# mk_qcow2 rootdev imagefile
mk_qcow2(){
  local rootdev="$1"
  local imagefile="$2"
  local imageext="${imagefile/*.}"
  [ "$(pwd)" = "/cache" ] || cd /cache
  echo "## $(date) : Starting creation of $imagefile."
  # kill torrent seeder for this image
  linbo_kill_seeder "$imagefile.torrent"
  # remove torrent files
  [ -e "/cache/$imagefile.torrent" ] && rm -f "/cache/$imagefile.torrent"
  [ -e "/cache/$imagefile.complete" ] && rm -f "/cache/$imagefile.complete"
  local RC=1
  local size="$(get_partition_size $rootdev)"
  if [ -z "$size" ]; then
    echo "Cannot get size of $rootdev!"
    return 1
  fi
  linbo_mount "$rootdev" /mnt ; RC="$?"
  if [ "$RC" != "0" ]; then
    echo "Cannot mount $rootdev!"
    return "$RC"
  fi
  echo "Preparing partition $rootdev (size=${size}K) for compression..."
  prepare_fs /mnt "$rootdev" ; RC="$?"
  if [ "$RC" != "0" ]; then
    echo "Cannot prepare $rootdev!"
    return "$RC"
  fi
  case "$imageext" in
    qcow2) mk_baseimage "$rootdev" "$imagefile" ; RC="$?" ;;
    qdiff) mk_diffimage "$rootdev" "$imagefile" ; RC="$?" ;;
    *) echo "Unknown image type!" ; RC="1" ;;
  esac
  if [ "$RC" != "0" ]; then
    umount /mnt || umount -l /mnt
    return "$RC"
  fi
  # mount rootdevice after image creation
  linbo_mount "$rootdev" /mnt ; RC="$?"
  if [ "$RC" != "0" ]; then
    echo "Mounting of $rootdev failed. :("
    return "$RC"
  fi
  # save imagename on partition
  echo "${imagefile%.*}" > /mnt/.linbo
  umount /mnt || umount -l /mnt
  # create info file
  local imgsize="$(get_filesize $imagefile)"
  linbo_mkinfo "$imagefile" "$rootdev"
  echo "Done."
  ls -l "$rootdev"
  touch "$imagefile".complete
  # create torrent file
  linbo_mktorrent "$imagefile" || RC="1"
  echo "## $(date) : Creation of $imagefile finished."
  return "$RC"
}

#### functions end ####


# print help
[ -z "$1" -o "$1" = "help" ] && usage 0

# get environment
source /usr/share/linbo/shell_functions
echo "### $timestamp $(basename "$0") $@ ###"

# start.conf is needed
if [ ! -s /start.conf ]; then
  echo "Fatal: start.conf not found!"
  exit 1
fi

# check os number option
if isinteger "$1"; then
  osnr="$1"

# check rootdev option
elif [ -b "$1" -a "$1" != "$cache" ]; then
  osnr="$(osnr_startconf "$1")"

# check legacy options (cache imagefile rootdev)
else
  for item in $@; do
    [ "$item" = "$cache" ] && continue
    case "$item" in *.qcow2|*.qdiff) imagefile="$item" ; continue ;; esac
    [ -b "$item" ] && root="$item"
  done
  # check items
  [ -z "$imagefile" ] && usage 1
  [ -b "$root" ] || usage 1
fi

# read os config parameters if osnr is known
if [ -n "$osnr" ]; then
  source /conf/os.$osnr &> /dev/null || usage 1
  # check differential image
  if [ "$2" = "qdiff" ]; then
    imagefile="${baseimage/.qcow2/.qdiff}"
  else
    imagefile="$baseimage"
  fi
fi

linbo_mountcache &> /dev/null || exit 1

cd /cache

echo "Creating image $imagefile of partition $root ..."
mk_qcow2 "$root" "$imagefile" ; RC="$?"

if [ "$RC" = "0" ]; then
  echo "Done."
else
  echo "Failed."
fi

sendlog

cd /

exit "$RC"
