#!/bin/bash
#
# vmware-mkinitrd
#
# Adapted from 'mkinitrd' by Erik Troan and others (see mkinitrd). 
#
# Portions copyright 2004 VMware, Inc.  
# 
#  This program is free software; you can redistribute it and/or
#  modify it under the terms of the GNU General Public License
#  as published by the Free Software Foundation; either version 2
#  of the License, or (at your option) any later version.


PATH=/sbin:/usr/sbin:/usr/bin:$PATH
export PATH

VERSION=3.5.13
IMAGESIZE=96000

IMAGE=
MNTPOINT=

compress=1
target=""
kernel=""
force=""
verbose=""
showrc=""
rc=0
fstab="/etc/fstab"

msgfile="log/messages"
mountroot="bin/mountroot" 
shellrc="bin/shellrc"
vmsupport="bin/vm-support"

PREPSCRDIR=/etc/vmware/init

UTILDRVS="usbcore.o usb-uhci.o usb-ohci.o usb-storage.o input.o \
          hid.o keybdev.o fat.o vfat.o cramfs.o loop.o"

usage() {
/bin/cat << EOF  
usage: `basename $0` [options] <initrd-image> <kernel-version> 
 
-v               Be verbose. (-vv: Show resulting rc files too.) 
-f               Force; overwrite <initrd-image> if it exists. 
-r               Refresh; update <initrd-image>, don't create. 
--version        Print vmware-mkinitrd version. 
--fstab=<fstab>  Use <fstab> instead of the default /etc/fstab. 
 
(e.g. `basename $0` /boot/initrd-2.2.5-15.img 2.2.5-15) 
EOF
      exit 1
}

inst() {
    if [ "$#" != "2" ];then
        echo "usage: inst <file> <destination>"
        exit 1 
    fi 
    if [ ! -e "$1" ]; then
       echo "$1 not found."
       exit 1
    fi
    cp $verbose $1 $2
}

main () 
{
   local refresh=

   if [ $(id -u) != 0 ]; then
      echo "vmware-mkinitrd must be run as root."
      exit 1
   fi

   # Arg parsing.
   while [ $# -gt 0 ]; do
      case $1 in
         --fstab*)
               if echo $1 | grep -q '=' ; then
                  fstab=`echo $1 | sed 's/^--fstab=//'`
               else
                  fstab=$2
                  shift
               fi		    
               ;;

         --version)
               echo "vmware-mkinitrd: version $VERSION"
               exit 0
               ;;

         -v)
               verbose=-v
               ;;
         -vv)
               verbose=-v
               showrc="yes"
               ;;

         --nocompress)
               compress=""
               ;;
         -r)
               refresh=1
               ;;
         -f)
               force=1
               ;;
         *)
               if [ -z "$target" ]; then
                  target=$1
               elif [ -z "$kernel" ]; then
                  kernel=$1
               else
                  usage
               fi
               ;;
      esac

      shift
   done

   if [ -z "$target" -o -z "$kernel" ]; then
      usage
   fi

   if [ -z "$force" -a -f "$target" ]; then
      echo "$target exists." >&2
      exit 1
   fi

   # find a temporary directory which doesn't use tmpfs
   TMPDIR=""
   for t in /tmp /var/tmp /root ${PWD}; do
      if [ ! -d $t ]; then continue; fi
      if ! echo access -w $t | /sbin/nash --quiet; then continue; fi

      fs=$(df -T $t 2>/dev/null | tail -1 | awk '{printf $2}')
      if [ "$fs" != "tmpfs" ]; then 
         TMPDIR=$t
         break
      fi
   done

   if [ -z "$TMPDIR" ]; then
      echo "No temporary directory could be found." >&2
      exit 1
   fi

   if [ $TMPDIR = "/root" -o $TMPDIR = "${PWD}" ]; then 
      echo "WARNING: using $TMPDIR for temporary files" >&2
   fi

   IMAGE=`mktemp ${TMPDIR}/initrd.img.XXXXXX`
   MNTPOINT=`mktemp -d ${TMPDIR}/initrd.mnt.XXXXXX`
   trap "umount $MNTPOINT >&/dev/null; rm -rf $MNTPOINT $IMAGE >&/dev/null" EXIT TERM KILL

   if [ -z "$IMAGE" -o -z "$MNTPOINT" ]; then
      echo "Error creating temporaries. Try again." >&2
      exit 1
   fi

   if [ -n "$refresh" ]; then
      gunzip -c $target > $IMAGE || exit 1
   else
      dd if=/dev/zero of=$IMAGE bs=1k count=$IMAGESIZE 2> /dev/null || exit 1
      # echo y to confirm that $IMAGE is not a block device.
      echo y | mke2fs $IMAGE $IMAGESIZE >/dev/null 2>/dev/null || exit 1
      tune2fs -c0 -i0 $IMAGE >/dev/null
   fi

   mount -t ext2 -o loop $IMAGE $MNTPOINT || {
      echo "Couldn't mount loopback device." >&2
      exit 1
   }

   if [ -n "$refresh" ]; then
      refresh_image || exit 1 
   else
      create_image || exit 1 
   fi

   umount $MNTPOINT

   if [ -n "$compress" ]; then
      gzip --fast < $IMAGE > $target || rc=1
   else
      cp -a $IMAGE $target || rc=1
   fi

   rm -rf $MNTPOINT $IMAGE

   trap - EXIT TERM KILL
   exit $rc
}

#
# Overwrites those parts of the initrd that are affected by config changes.
#
update_changeable ()
{
   local mode=$1
   local targdir=$2
   local kernel=$3
   local options=

   if [ "$mode" == "refresh" ]; then
      options="-r"
   fi
  
   # Call all prep scripts to do their refresh actions.
   for prep in `ls $PREPSCRDIR`; do
      [ -n "$verbose" ] && echo "Calling prep script: $prep $options $targdir $kernel"
      $PREPSCRDIR/$prep "$options" "$targdir" "$kernel" 
      if [ $? -ne 0 ]; then
         echo "$prep failed."
         return 1
      fi
   done

   return 0
}

#
# Refresh ramdisk image "$target" for "$kernel".
#
refresh_image ()
{
   update_changeable "refresh" "$MNTPOINT" "$kernel" || return 1
   return 0
}

#
# Create a new ramdisk image "$target" for "$kernel"
#
create_image ()
{
   KMODDIR=/lib/modules/$kernel
   if [ ! -d $KMODDIR ]; then
      echo "$KMODDIR is not a directory." >&2
      return 1
   fi

   RCFILE=$MNTPOINT/linuxrc
   MOUNTRC=$MNTPOINT/$mountroot 
   SHELLRC=$MNTPOINT/$shellrc
   ROOTRC=$MNTPOINT/init/100.rootfs

   rootfs=$(awk '{ if ($1 !~ /^[ \t]*#/ && $2 == "/") { print $3; }}' $fstab)
   rootopts=$(awk '{ if ($1 !~ /^[ \t]*#/ && $2 == "/") { print $4; }}' $fstab)

   # We don't need this directory, so let's save space
   rm -rf $MNTPOINT/lost+found

   mkdir -p $MNTPOINT/lib
   mkdir -p $MNTPOINT/bin
   mkdir -p $MNTPOINT/etc
   mkdir -p $MNTPOINT/init
   mkdir -p $MNTPOINT/log
   mkdir -p $MNTPOINT/dev/pts
   mkdir -p $MNTPOINT/loopfs
   mkdir -p $MNTPOINT/proc
   mkdir -p $MNTPOINT/sysroot
   ln -s bin $MNTPOINT/sbin

   inst /sbin/nash "$MNTPOINT/bin/nash"
   inst /sbin/insmod.static "$MNTPOINT/bin/insmod"
   inst /sbin/busybox "$MNTPOINT/bin/ash"
   inst /sbin/dumpilogs "$MNTPOINT/bin/vm-support"
   inst /usr/bin/md5sum "$MNTPOINT/bin/md5sum"
   ln -sf ash $MNTPOINT/bin/sh  

   # mknod'ing the devices instead of copying them works both with and
   # without devfs...
   [ -n "$verbose" ] && echo "Creating device nodes"
   mknod $MNTPOINT/dev/console c 5 1
   mknod $MNTPOINT/dev/null c 1 3
   mknod $MNTPOINT/dev/ram b 1 1
   mknod $MNTPOINT/dev/systty c 4 0
   for i in 1 2 3 4 5 6 7; do
      mknod $MNTPOINT/dev/tty$i c 4 $i
      mknod $MNTPOINT/dev/loop$i b 7 $i
   done
   mknod $MNTPOINT/dev/mem c 1 1
   mknod $MNTPOINT/dev/psaux c 10 1
   mknod $MNTPOINT/dev/ptmx c 5 2
   mknod $MNTPOINT/dev/zero c 1 5
   mknod $MNTPOINT/dev/vmkcall c 10 167

   # Call all prep scripts to populate the rd image.
   update_changeable "new" "$MNTPOINT" "$kernel" || return 1

   # ... and do a bit of our own prepping.
   # Init script to mount root and optionally call fallback script.
   # Written first so that it's picked up along with the other init scripts
   # during linuxrc creation. 
    
   [ -n "$verbose" ] && echo "Writing $ROOTRC"
   /bin/cat > $ROOTRC << EOF
#!/bin/ash 

echo -n "Mounting root..." 
$mountroot >>$msgfile 2>&1
if [ \$? -eq 0 ]; then
   echo "ok." 
else
   echo "failed." 
   $shellrc 'last'
fi
EOF

   [ $? -ne 0 ] && return 1 
   chmod +x $ROOTRC
   [ -n "$showrc" ] && echo "" && /bin/cat $ROOTRC && echo ""

   # Linuxrc, the main initrd script, simply calls all other scripts in order.

   [ -n "$verbose" ] && echo "Writing linuxrc"
   initscripts=`ls $MNTPOINT/init | grep -e '^[0-9]\+\.' | sort -n | xargs`

   /bin/cat > $RCFILE << EOF
#!/bin/ash

echo "Mounting /proc filesystem"
mount -t proc /proc /proc

grep -q "failshell" /proc/cmdline 
if [ \$? -eq 0 ]; then
   export failshell="$shellrc 'inter'"
fi

\$failshell
for f in $initscripts; do
   /init/\$f
   if [ \$? -eq 0 ]; then
      echo "\$f succeeded." 
   else
      echo "\$f failed."
      \$failshell
   fi
done

# Root must have been mounted at this point.
/bin/umount /initrd/proc
EOF

   [ $? -ne 0 ] && return 1 
   chmod +x $RCFILE
   [ -n "$showrc" ] && echo "" && /bin/cat $RCFILE && echo ""

   # Nash script to mount root.
   # Separate script because nash doesn't accept cmdline commands, and only it 
   # has mkrootdev. 

   [ -n "$verbose" ] && echo "Writing mountroot"
   /bin/cat > $MOUNTRC << EOF
#!/bin/nash --force

echo Creating block devices
mkdevices /dev
echo Creating root device
mkrootdev /dev/root
echo 0x0100 > /proc/sys/kernel/real-root-dev
echo Mounting root filesystem
mount -o $rootopts --ro -t $rootfs /dev/root /sysroot
pivot_root /sysroot /sysroot/initrd
EOF

   [ $? -ne 0 ] && return 1 
   chmod +x $MOUNTRC
   [ -n "$showrc" ] && echo "" && /bin/cat $MOUNTRC && echo ""

   # Util modules for shell. 

   for f in $UTILDRVS; do
      mod=`find $KMODDIR -type f -name $f` 
      if [ ! -n "$mod" ]; then
         echo "$f not found in $KMODDIR." && return 1
      fi
      inst "$mod" "$MNTPOINT/lib"
   done

   # This script drops you into a shell. It loads USB and other utility 
   # drivers that may be needed for dumping logs and general noodling. 

   [ -n "$verbose" ] && echo "Writing $SHELLRC"
   /bin/cat > $SHELLRC << EOF
#!/bin/ash

last=
if [ "\$1" = "last" ]; then
   last=yes
fi

if [ ! -e /log/util ]; then
   for f in $UTILDRVS; do
      echo "Loading console \$f"
      insmod /lib/\$f >>$msgfile 2>&1
      if [ \$? -ne 0 ]; then
         echo "\$f failed to load in Service Console."
      fi
   done
   touch /log/util
fi

if [ -n "\$last" ]; then
   echo ""
   echo "Mounting root failed. Dropping into basic maintenance shell."
   echo "To collect logs for VMware, connect a USB storage device and" 
   echo "run '$vmsupport <devicename>'." 
   echo "Machine will be rebooted when you exit from this shell."
fi

/bin/ash

[ -n "\$last" ] && reboot
EOF

   [ $? -ne 0 ] && return 1 
   chmod +x $SHELLRC
   [ -n "$showrc" ] && echo "" && /bin/cat $SHELLRC && echo ""
   
   return 0
}

# Start here.
main $*
