#!/bin/sh -efu
#
# For keys description see:
# http://www.opengroup.org/onlinepubs/009695399/utilities/set.html
#
# Copyright (C) 2007  Andrew V. Stepanov <stanv@altlinux.org>
#
# Get Every Archive from git package Repository.
#
# This file 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.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301, USA.
#


# Notes XXX:
#	* Assume that all device drivers already loaded (not need do modprobe, or init LVM, raid, or anything else)
#	* What will be if LILO installed to /dev/hda1, not at /dev/hda ?

PROG="${0##*/}"
PROG_VERSION='1.0.0'
PARTITIONS="/proc/partitions"
LILO_CONFIG="/etc/lilo.conf"
RM_DIRS=( ) # init empty array, will be contains directory that should be deleted at exit_handler
UMOUNT_DIRS=( ) # init empty array, will be contains directory that should be umounted at exit_handler

LC_ALL=C
export LC_ALL

show_help()
{
    cat <<EOF
    $PROG - restore MBR erased by Microsoft Windows (TM)

    Usage: $PROG [options]

    Options:
    --lilo                    restore LILO MBR;
    --grub                    restore GRUB MBR;
    -e, --etc-dev=DEVICE      device containing /etc directory
    -q, --quiet, --silent     try to be more quiet;
    -v, --verbose             print a message for each action;
    -V, --version             print program version and exit;
    -h, --help                show this text and exit.

    Report bugs to http://bugs.altlinux.ru/

EOF
    exit
}

print_version()
{
    cat <<EOF
    $PROG version $PROG_VERSION
    Written by Andrew V. Stepanov <stanv@altlinux.org>

    Copyright (C) 2007  Andrew V. Stepanov <stanv@altlinux.org>
    This is free software; see the source for copying conditions.  There is NO
    warranty; not even for MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
EOF
    exit
}

msg_info()
{
    printf '%s\n' "$PROG: $*" >&2
}

fatal()
{
    msg_info "$@"
    exit 1
}

show_usage()
{
    [ -z "$*" ] || msg_info "$*"
    echo "Try \`$PROG --help' for more information." >&2
    exit 1
}

verbose()
{
    [ -n "$verbose" ] || return 0
    msg_info "$@"
}

uuid2dev()
{
    local uuid=$1 && shift # in form UUID=e280469a-d06f-4c0b-b068-44f3b576029e or LABEL=root
    local dev=$uuid
    dev=$(blkid -o device -t "$uuid" 2>/dev/null) | sed -e '1!d' ||  dev=$uuid
    verbose "uuid2dev: transform $uuid to $dev"
    echo $dev
    
}

is_root()
{
    local disk=$1 && shift
    verbose "is_root() $disk"
    [ -d "$disk/etc" ] &&
    [ -d "$disk/boot" ] &&
    [ -d "$disk/home" ] &&
    [ -d "$disk/var" ] &&
    [ -d "$disk/bin" ] &&
    [ -d "$disk/lib" ] &&
    [ -d "$disk/proc" ] && return 0
    return 1
}

is_etc()
{
    local dir=$1 && shift
    [ -f "$dir/inittab" ] &&
    [ -f "$dir/fstab" ] &&
    [ -f "$dir/group" ] &&
    [ -f "$dir/passwd" ] &&
    [ -f "$dir/mtab" ] && return 0
    return 1
}

good_lilo()
{
    boot_str="`cat $disk/etc/lilo.conf | grep -v '^#' | grep "^boot"`"
    boot_str1=$( echo `expr match "$boot_str" '.*\([s,h]d[a-f][ 1-9]*\)'`  )
    boot_str2=$( echo `expr match "$boot_str" '.*\([s,h]d[a-f][ ]*\)'`  )
    if [ -n "$boot_str2" ] && [ "$boot_str1" = "$boot_str2" ];then
	verbose good boot_str1=$boot_str1
	return 0
    else 
	verbose "bad boot_str1=$boot_str1 boot_str2=$boot_str2"
	return 1
    fi
}

good_grub()
{
    return 0
}

is_mbr()
{
  verbose "is_mbr()"

        if [ -f "$disk/etc/lilo.conf" ] ; then
            verbose "found LILO boot loader"
            bootloader="lilo"
            verbose $bootloader
            good_lilo 
            itog=$?
            verbose "itog=$itog"
                if [ "$itog" -eq 0 ]; then
        	    verbose  "good_lilo"
        	    return 0
        	else 
        	    verbose "bad_lilo"
        	    return 1
    	    fi
        elif [ -f "$disk/boot/grub/menu.lst" ]; then
            verbose "found GRUB boot loader"
            bootloader="grub"
            verbose $bootloader
            good_grub 
            itog=$?
            verbose "itog=$itog"
            if [ -z "$itog" ]; then 
        	    return 0 
        	else 
        	    return 1
    	    fi
    fi
}

is_mbr_boot()
{
  local disk=$1 && shift
  is_root "$disk"  && is_mbr "$disk" && echo "mbr $disk good" && return 0
  verbose "mbr $disk is bad" && return 1
}

find_etc()
{
    local active_root=
    local active_etc=

    verbose "try to find devices with / or /etc directories, have to examine folowing file systems:"
    verbose "FS: $(cat "$tmpdir/available.disks" | paste -s -d ' ')"

    verbose "check already mounted file systems"
    cat "$tmpdir/available.disks" |
    {
        while read disk; do
            if (mount | grep -q "$disk"); then
                sed -i "\@$disk@d" "$tmpdir/available.disks"
                local mounted_at="$( mount | grep "$disk" | cut -d ' ' -f 3 | sed 1!d )" || mounted_at=none
                verbose "check $mounted_at mount point it is a root partition"
                if is_root "$mounted_at" && is_etc "$mounted_at/etc" && is_mbr_boot "$mounted_at" ;then
                    verbose "$disk mounted at $mounted_at is root partition, continue..."
                    echo "$disk" >> $tmpdir/found.disks
                else
                    verbose "$disk isn't a root partition, skip"
                fi
            fi
        done
    }

    verbose "check not mounted file systems, by mounting them to $mnt_point"
    verbose "umounted FS: $(cat "$tmpdir/available.disks" | paste -s -d ' ')"
    cat "$tmpdir/available.disks" |
    {
        while read disk; do
            verbose "mount $disk to $mnt_point"
            mount -o ro "$disk" "$mnt_point" || 
            {
                verbose "can't mount $disk to $mnt_point"
                continue
            }
            verbose "mount done"
            if is_root "$mnt_point" && is_etc "$mnt_point"/etc && is_mbr_boot "$mnt_point"; then
                verbose "temporary mounted $disk at $mnt_point is root or /etc partition, continue..."
                echo "$disk" >> "$tmpdir/found.disks"
                verbose "disk=$disk"
            else
                verbose "$disk is not a root or etc partition, skip"
    		verbose "bad disk = $disk"
        
            fi

            verbose "umount $disk from $mnt_point"
            umount "$disk" || fatal "can't umount $disk from $mnt_point"
        done
    }
}

restore_mbr()
{
    local root=$1 && shift
    local bootloader=$1 && shift
    local device=$1 && shift
     
     verbose restore_mbr $root $device 
    if [ ! -d "$root" ]; then
        fail "bad root $root"
    fi

#    ls $root
    verbose "chroot $root mount -a"

    chroot "$root" mount -a || 

    if [ "$bootloader" = "unknown" ]; then
        verbose "try to detect current boot loader"
        if [ -f "$root/etc/lilo.conf" ] && chroot "$root" lilo -V >/dev/null 2>&1; then
            verbose "found LILO boot loader"
            bootloader="lilo"
        elif [ -f "$root/boot/grub/menu.lst" ] && chroot "$root" which grub-install >/dev/null 2>&1; then
            verbose "found GRUB boot loader"
            bootloader="grub"
        else
            fatal "can't guess what boot loader should be used"
        fi
    fi

    verbose "Update $root/etc/mtab"
    chroot "$root" grep -v rootfs /proc/mounts > "$root/etc/mtab" || fatal "can't update $root/etc/mtab"

    case  "$bootloader"  in
        lilo)
        {
            verbose "runing LILO in chroot $root"
            chroot "$root" lilo -V >/dev/null 2>&1 || fatal "can't find LILO on the chroot $root, maybe try to restore GRUB MBR instead?"
            if [ -n "$verbose" ]; then
                chroot "$root" lilo "$verbose"
            else
                chroot "$root" lilo
            fi
        }
        ;;
        grub)
        {
            verbose "runing GRUB in chroot $root for device $device"
            chroot "$root" which grub-install >/dev/null 2>&1 || fatal "can't find GRUB on the chroot $root, maybe try to restore LILO MBR instead?"
            chroot "$root" grub-install "$device"
        }
        ;;
        *)
        {
            fatal "unsupported bootloader: $bootloader"
        }
        ;;
    esac
}

exit_handler()
{
    local rc=$?
    trap - EXIT

    verbose "call exit_handler with \`$rc' exit code"

    verbose "remove temporary directory $tmpdir"
    [ -z "$tmpdir" ] ||
    rm -rf -- "$tmpdir"
    sync

    verbose "umount mounted filesystem under $mnt_point"
    cat "/proc/mounts" | awk '{print $2}' | grep "^${mnt_point}" | sort -r |
    while read dir;
    do
        verbose "umount $dir"
        umount "$dir" || :
    done

    local rm_num=${#RM_DIRS[@]}
    local umount_num=${#UMOUNT_DIRS[@]}

    while [ "$rm_num" -ne "0" ]; do
        local index=$((rm_num - 1))
        local dir="${RM_DIRS[$index]}"
        verbose "remove directory $dir"
        rmdir $dir || :
        rm_num=$((rm_num - 1))
    done

    verbose "done removing directories"

    while [ "$umount_num" -ne "0" ]; do
        local index=$((umount_num - 1))
        local dir="${UMOUNT_DIRS[$index]}"
        verbose "umount directory $dir"
        umount $dir || :
        umount_num=$((umount_num - 1))
    done

    verbose "done umount directories"

    if [ "$rc" -eq "0" ]; then
        printf "MBR successfully restored!\n"
    fi

    exit $rc
}

scan_disks()
{
    local uuid=
    local disk=
    verbose "remove duplicates disks (same disk may have several names, eg dm-0 or hda)"
    blkid | grep -o 'UUID=[^[:space:]]*' | sort | uniq |
    {
        while read uuid; do
            disk=$(blkid -o device -t "$uuid" 2>/dev/null | sed -e '1!d')
            echo "$disk" >> $tmpdir/available.disks
            verbose "add $disk to further look what it is"
        done
    }
}

# =======================
# Program start from here
# =======================

# Parse command line options

OPT=`getopt -n $PROG -o e:,h,q,v,V -l lilo,grub,etc-dev:,quiet,silent,verbose,version,help -- "$@"` ||
show_usage
eval set -- "$OPT"

# Options
lilo=
grub=
mnt_point=
root_dev=
boot_dev=
etc_dev=
do_nothing=
quiet=
verbose=
# Other variables
tmpdir=
del_mnt_point=

while :; do
    case $1 in
        --) shift; break
        ;;
        --lilo)
        {
            lilo="lilo"
            [ -z "$grub" ] ||
            show_usage 'Options --lilo, --grub are mutually exclusive.'
        }
        ;;
        --grub)
        {
            grub="grub"
            [ -z "$lilo" ] ||
            show_usage 'Options --lilo, --grub are mutually exclusive.'
        }
        ;;
        -e|--etc-dev)
        {
            shift; etc_dev=$1
        }
        ;;
        -q|--quiet|--silent)
        {
            quiet=yes
        }
        ;;
        -v|--verbose)
        {
            verbose=-v
        }
        ;;
        -V|--version)
        {
            print_version
        }
        ;;
        -h|--help)
        {
            show_help
        }
        ;;
        *)
        {
            echo "unrecognized argument: " $1
            show_usage
        }
        ;;
    esac

    shift
done

if [ "$UID" != "0" ] ; then
    fatal "$PROG requires root privileges"
fi

which blkid > /dev/null 2>&1 || fatal "$PROG requires blkid from e2fsprogs package"

if [ -n "$quiet" ]; then
    exec > /dev/null 2>&1
fi

mnt_point="$( (unset TMPDIR; mktemp -dp /mnt/) )"
[ -d "$mnt_point" ] || fatal "Could not create temporary directory for mounting file systems"
RM_DIRS[${#RM_DIRS[*]}]="$mnt_point"

tmpdir="$(mktemp -dt "$PROG.XXXXXXXX")" || fatal "Could not create temporary directory"
verbose "tmpdir=$tmpdir"
verbose "$tmpdir directory created for store temporary files"

trap exit_handler EXIT SIGHUP SIGINT SIGQUIT SIGTERM

scan_disks # result in $tmpdir/available.disks

# Skip active root
active_root=$(grep </etc/fstab -v "^[ \t]*#" | grep "[[:space:]]/[[:space:]]" | sed -e "s/^[ \t]*//" -e "s/[ \t].*//")
if [ -n "$active_root" ]; then
    active_root=$(uuid2dev "$active_root")
else
    active_root=`mount | grep " on / type" | sed "s/ .*//"`
fi
if [ -n "$active_root" ]; then
    verbose "skip active root $active_root from concerned disks"
    sed -i "\@$active_root@d" "$tmpdir/available.disks"
else
    verbose "can't find active mounted root partition, it's ok, continue..." # for example squashfs
fi

# Skip active /etc
active_etc=$(grep </etc/fstab -v "^[ \t]*#" | grep "[[:space:]]/etc[[:space:]]" | sed -e "s/^[ \t]*//" -e "s/[ \t].*//")
if [ -n "$active_etc" ]; then
    active_etc=$(uuid2dev "$active_etc")
else
    active_etc=`mount | grep " on /etc type" | sed "s/ .*//"`
fi
if [ -n "$active_etc" ]; then
    verbose "skip active /etc $active_etc from concerned disks"
    sed -i "\@$active_etc@d" "$tmpdir/available.disks"
fi

verbose "skip non supported file systems. Keep only: ext2, ext3, reiserfs, xfs, jfs"
cat "$tmpdir/available.disks" |
{
    while read disk; do
        blkid "$disk" | grep -iq -e ext2 -e ext3 -e reiser -e xfs -e jfs ||
        {
            verbose "disk $disk has not supported file system, remove him from concerned disks"
            sed -i "\@$disk@d" "$tmpdir/available.disks"
        }
    done
}

if [ -n "$etc_dev" ]; then
    grep -q "$etc_dev" "$tmpdir/available.disks" || fatal "can't operate on $etc_dev device"
    [ -b "$etc_dev" ] || fatal "bad etc-dev \`$etc_dev', it is not a block device"
else
    verbose "no device containing /etc files specified, try to find it by self"
    verbose "create temporary directory $mnt_point for mounting unknown file systems"
    mkdir -p "$mnt_point"
    find_etc # result in $tmpdir/found.disks
    [ -f "$tmpdir/found.disks" ] || show_usage "can't find any disk contain /etc files, try to specify it with --etc-dev=DEVICE option"
    [ "$(wc -l "$tmpdir/found.disks" | cut -f 1 -d ' ')" = "1" ] || show_usage "Found to many devices with /etc directories, try to specify one with --etc-dev=DEVICE option"
    etc_dev=$(head -n 1 "$tmpdir/found.disks")
fi

# At this stage $etc_dev may be mounted or unmounted
if (mount | grep -q "$etc_dev"); then
    # Already mounted
    mounted_at=$(mount | grep "$etc_dev" | cut -f 3 -d ' ' | sed 1!d )
    if ! is_root "$mounted_at" && is_etc "$mounted_at/etc" && is_mbr_boot "$mounted_at"; then
        # maybe user via --etc-dev give us real etc device, not root device
        mounted_at=$(realpath "$mounted_at/..")
        if [ ! is_root "$mnt_point" || ! is_etc "$mnt_point/etc" || ! is_mbr_boot "$mounted_at" ]; then
            fatal "disk $etc_dev doesn't contains /etc files"
        fi
    fi
    verbose "target device $etc_dev already mounted at $mounted_at"
    boot_dev=$(grep <"$mounted_at/etc/fstab" -v "^[ \t]*#" | grep "[[:space:]]/boot[[:space:]]" | sed -e "s/^[ \t]*//" -e "s/[ \t].*//")
    if [ -n "$boot_dev" ] ; then
        boot_dev=$(uuid2dev "$boot_dev")
        if ! mount | grep -q "$boot_dev"; then
            verbose "mount /boot device $boot_dev to $mounted_at/boot"
            mount "$boot_dev" "$mounted_at/boot" || fatal "Could not mount $boot_dev to $mounted_at"
            UMOUNT_DIRS[${#UMOUNT_DIRS[*]}]="$mounted_at/boot"
        fi
    fi


    echo "disk is mounted"

    mount --bind /dev  "$mounted_at/dev" || fatal "Could not mount (bind) /dev to chroot $mounted_at/dev"
    UMOUNT_DIRS[${#UMOUNT_DIRS[*]}]="$mounted_at/dev"
    mount --bind /dev/pts "$mounted_at/dev/pts" || fatal "Could not mount (bind) /dev/pts to chroot $mounted_at/dev/pts"
    UMOUNT_DIRS[${#UMOUNT_DIRS[*]}]="$mounted_at/dev/pts"
    mount --bind /proc "$mounted_at/proc" || fatal "Could not mount (bind) /proc to chroot $mounted_at/proc"
    UMOUNT_DIRS[${#UMOUNT_DIRS[*]}]="$mounted_at/proc"





    restore_mbr "$mounted_at" $([ -n "${lilo}${grub}" ] && echo "${lilo}${grub}" || echo "unknown") "$etc_dev"
else
    # Mount by self
    verbose "mount device $etc_dev to $mnt_point"



    mount "$etc_dev" "$mnt_point" || show_usage "Could not mount $etc_dev to $mnt_point"
    echo "mount $etc_dev $mnt_point"


    if (is_etc "$mnt_point"/etc); then
        root_dev=$(grep <"$mnt_point/etc/fstab" -v "^[ \t]*#" | grep "[[:space:]]/[[:space:]]" | sed -e "s/^[ \t]*//" -e "s/[ \t].*//")
        root_dev=$(uuid2dev "$root_dev")
        verbose "real root is $root_dev, while $etc_dev is disk wich contains /etc directory"
        umount "$mnt_point" || fatal "Could not umount $etc_dev from $mnt_point"
        verbose "mount / device $root_dev to $mnt_point"
        mount "$root_dev" "$mnt_point" || fatal "Could not mount $root_dev to $mnt_point"
        verbose "mount /etc device $etc_dev to $mnt_point"
    #    mount "$etc_dev" "$mnt_point" || fatal "Could not mount $etc_dev to $mnt_point"
    fi
    
    boot_dev=$(grep <"$mnt_point/etc/fstab" -v "^[ \t]*#" | grep "[[:space:]]/boot[[:space:]]" | sed -e "s/^[ \t]*//" -e "s/[ \t].*//")
    if [ -n "$boot_dev" ] ; then
        boot_dev=$(uuid2dev "$boot_dev")
        verbose "mount /boot device $boot_dev to $mnt_point/boot"
        mount "$boot_dev" "$mnt_point/boot" || fatal "Could not mount $boot_dev to $mnt_point"
    fi

    verbose $mnt_point
    verbose "disk is not mounted"

    mount --bind /dev  "$mnt_point/dev" || fatal "Could not mount (bind) /dev to chroot $mnt_point/dev"
    mount --bind /dev/pts "$mnt_point/dev/pts" || fatal "Could not mount (bind) /dev/pts to chroot $mnt_point/dev/pts"
    mount --bind /proc "$mnt_point/proc" || fatal "Could not mount (bind) /proc to chroot $mnt_point/proc"
    restore_mbr "$mnt_point" $([ -n "${lilo}${grub}" ] && echo "${lilo}${grub}" || echo "unknown") "$etc_dev"
fi

exit 0 # Go to exit handler

# vim:ts=4:sw=4:expandtab
