#!/bin/bash

# error code
SUCCESS=0
FAILED=1
WARNING=2

# counter
c_ok=0
c_err=0
c_wrn=0

# color string setting 
OK="\033[32;40m [  OK  ] \033[0m"
WRN="\033[33;40m [WARING] \033[0m"
ERR="\033[31;40;1m [ERROR] \033[0m"

fix_width="\033[45G"
split_str=$(printf "%*s" "55")
split_line=${split_str// /-}

distro="unknown"
version="unknown"
full_version="unknown"

# is root user
uid=$(id -u)
if [  "$uid" -ne 0 ]; then
    echo "Please run this script as root."
    exit 
fi

# get distro
for file in issue issue.net lsb-release
do
    [ ! -f "/etc/$file" ] && continue
    grep -iq "ubuntu" /etc/$file
    if [ "$?" -eq 0 ]; then
        distro="ubuntu"
    fi
done

if [ "$distro" == "ubuntu" ]; then
    version=$(cat /etc/lsb-release | grep DISTRIB_RELEASE | awk -F '=' '{print $2}' | awk -F '.' '{print $1}')
    full_version=$(cat /etc/lsb-release | grep DISTRIB_RELEASE | grep -oE '[0-9]+\.[0-9]+')
    if [ "$vresion" == "unknown" ]; then
        echo "Can't get the version of distro."
        exit
    fi
else
    echo "Please run the script on the Ubuntu system."
    exit
fi

dtime=`date +%Y%m%d%H%M%S`
msg_file="${distro}_check_result_${dtime}.txt"

s_grub_file=""
if [ -e "/boot/grub/grub.cfg" ]; then
    s_grub_file="/boot/grub/grub.cfg"
elif [ -e "/boot/grub2/grub.cfg" ]; then
    s_grub_file="/boot/grub2/grub.cfg"
else
    s_grub_file="/boot/grub/menu.lst"
fi

function write_file(){
    echo -e "$1" >> $msg_file
}

function print_result(){
    case "$1" in
    0)
        echo -en $fix_width && echo -e $OK
        c_ok=$[c_ok+1]
        ;;
    1)
        echo -en $fix_width && echo -e $ERR
        c_err=$[c_err+1]
        ;;
    2)
        echo -en $fix_width && echo -e $WRN
        c_wrn=$[c_wrn+1]
        ;;
    *)
        echo -en $fix_width && echo -e $ERR
        c_err=$[c_err+1]
    esac
}

# virtio driver check
function check_virtio_driver(){
    write_file $split_line
    write_file "[ check virtio driver ]"

    s_cmd=""
    kernel_list=()

    # lsinitrd command
    which lsinitrd 1>/dev/null 2>&1
    if [ "$?" -eq 0 ]; then
        s_cmd="lsinitrd"
    else
        which lsinitramfs 1>/dev/null 2>&1
        if [ "$?" -eq 0 ]; then
            s_cmd="lsinitramfs"
        fi
    fi
    if [ -z "$s_cmd" ]; then
        #echo "Can't find lsinitrd or lsinitramfs command"
        write_file $split_line
        write_file "\tERROR: can't find lsinitrd or lsinitramfs command"
        return $FAILED
    fi

    # grub config file exists?
    if [ -z "$s_grub_file" ]; then
        write_file $split_line
        write_file "\tERROR: can't find the grub file.\n"
        return $FAILED
    fi
    k_list=($(cat $s_grub_file | grep "/vmlinuz" | grep -v -E "rescue|version" | sed 's/^.*vmlinuz-//' | awk '{print $1}'))
    
    for kernel in ${k_list[*]}
    do
        list_length=${#kernel_list[*]}
        kernel_exist=0
        for kernel_version in ${kernel_list[*]}
        do
            if [[ $kernel_version == "$kernel" ]];then
                kernel_exist=1
                break
            fi
        done
        [ $kernel_exist -eq 0 ] && kernel_list[$list_length]=$kernel
    done
    
    # Check if the current kernel contains virtio drivers
    current_kernel=`uname -r`
    kernel_conf="/boot/config-$current_kernel"
    if [ -r "$kernel_conf" ]; then
        ret=`grep -i CONFIG_VIRTIO_BLK $kernel_conf | awk -F '=' '{print $2}'`
        if [ $ret == 'y' ]; then
            #echo "The virtio driver exists."
            return 0
        elif [ $ret == 'm' ]; then
            initrdfile=$(find /boot \( -name "init*.img-$kernel_version" -o -name "init*-$kernel_version.img" -o -name "init*-$kernel_version" \)) 1>/dev/null 2>&1
            if [ -n $initrdfile ]; then
                $s_lsinitcmd $initrdfile 2>/dev/null | grep -q "virtio_blk.ko" 1>/dev/null 2>&1
            [ $? -eq 0 ] && ret=0
            else
                ls /boot/ | grep $(uname -r) | lsinitrd 2>/dev/null | grep -q "virtio_blk.ko"
                ret=$?
            fi
            if [[ $ret -eq 0 ]]; then
                return 0
            else
                write_file "\tERROR: can't fine the virtio driver."
            fi
        fi
    fi
    return $FAILED
}

function check_partition_type(){
    write_file $split_line
    write_file "[ check disk partition ]"
    which parted 1>/dev/null 2>&1
    if [ "$?" -ne 0 ]; then
        write_file "\tERROR: can't find parted command."
        write_file "\tERROR: please make sure that your system uses MBR partition."
        return $FAILED
    fi
    parted -l 2>/dev/null | grep "Partition Table" | grep -q "msdos"
    if [ "$?" -eq 0 ]; then
        return $SUCCESS
    else
        p_type=`parted -l 2>/dev/null | grep -i "Partition Table" | head -n 1 | awk '{print $3}'` 1>/dev/null 2>&1
        if [ "$?" -eq 0 ]; then
            write_file "\tERROR: the Partition Table on your disk is $p_type."
            write_file "\tERROR: we recommend using MBR partition."
            return $FAILED
        else
            write_file "\tERROR: can't find Partition Table type."
            write_file "\tERROR: please confirm that the partition is MBR."
            return $FAILED
        fi
    fi
}

function check_lvm(){
    write_file $split_line
    write_file "[ check lvm ]"
    disk=`fdisk -l 2>/dev/null | awk '/^Disk.*bytes/&&/\/dev/{print $2}' | awk -F ':' '{print $1}' | head -n 1`
    if [ -n $disk ]; then
        disk=`echo $disk | sed 's#/#\\\\/#g'`
        fdisk -l 2>/dev/null | awk "/^$disk/{ print $1 }" | grep -i "lvm" 1>/dev/null 2>&1
        if [ "$?" -eq 0 ]; then
            write_file "\tERROR: your system uses LVM partition."
            write_file "\tERROR: please use normal disk partion."
            return $FAILED
        else
            return $SUCCESS
        fi
    else
        write_file "\tERROR: can't find the system partition."
        write_file "\tERROR: please confirm that your system is not using LVM."
        return $FAILED
    fi
}

function check_selinux(){
    write_file $split_line
    write_file "[ check selinux ]"
    if [ -e "/etc/selinux/config" ]; then
        cat /etc/selinux/config | grep -v "#" | grep -v "^;"| grep -v "^$" | grep -i "selinux=enforcing" 1>/dev/null 2>&1
        if [ "$?" -eq 0 ]; then
            write_file "\tERROR: the SELinux on your system is enabled."
            write_file "\tERROR: we recommend disabling SELinux on your system."
            return $FAILED
        else
            return $SUCCESS
        fi
    else
        return $SUCCESS
    fi
}

function check_root(){
    write_file $split_line
    write_file "[ check root user ]"
    grep -q "^root:x:0:0" /etc/passwd
    if [ "$?" -eq 0 ]; then
        return $SUCCESS
    else
        write_file "\tERROR: can't find root user."
        return $FAILED
    fi
}

function check_passwd_cmd(){
    write_file $split_line
    write_file "[ check password command ]"
    which passwd &> /dev/null
    if [ "$?" -eq 0 ]; then
        return $SUCCESS
    else
        which chpasswd &> /dev/null
        if [ "$?" -eq 0 ]; then
            return $SUCCESS
        else
            write_file "\tERROR: can't find passwd & chpasswd command."
        fi
    fi
}

function check_shadow(){
    write_file $split_line
    write_file "[ check /etc/shadow ]"
    shadow="/etc/shadow"
    if [ ! -e $shadow ]; then
        write_file "\tERROR: the $shadow file isn't exist."
        return $FAILED
    fi
    if [ -w $shadow ]; then
        lsattr $shadow | grep 'i' 1>/dev/null 2>&1
        if [ "$?" -eq 0 ]; then
            write_file "\tERROR: the $shadow file is not writable."
            write_file "\tERROR: please use the chattr command to make $shadow writable."
            return $FAILED
        fi
        return $SUCCESS
    else
        write_file "\tERROR: the $shadow file is read only."
        return $FAILED
    fi
}

function check_network_config(){
    write_file $split_line
    write_file "[ check network config ]"
    interfaces=($(ls /sys/class/net | grep -v lo))
    if [ -z "$interfaces" ]; then
        write_file "\tERROR: can't find network device."
        return $FAILED
    fi
    
    ret=0
    for interface in ${interfaces[*]}
    do
        cat /etc/network/interfaces | grep -q "iface $interface" &>/dev/null
        if [ "$?" -eq 0 ]; then
            cat /etc/network/interfaces | grep "iface $interface" | grep -q "dhcp" &>/dev/null
            if [ "$?" -ne 0 ]; then
                write_file "\tERROR: the NIC \"$interface\" isn't set to dhcp mode."
                [ $ret -eq 1 ] || ret=1
            fi
        else
            #config_files=($(ls /etc/network/interfaces.d))
            grep -r "iface $interface" /etc/network/interfaces.d/ | grep -v "^#" &> /dev/null
            if [ "$?" -ne 0 ]; then
                write_file "\tERROR: can't find NIC $interface\'s configuration."
                [ $ret -eq 1 ] || ret=1
            else
                grep -r "iface $interface" /etc/network/interfaces.d/ | grep -v "^#" | grep -q "dhcp" &> /dev/null
                if [ "$?" -ne 0 ]; then
                    write_file "\tERROR: the NIC \"$interface\" isn't set to dhcp mode."
                    [ $ret -eq 1 ] || ret=1
                fi
            fi
        fi
    done
    net_rule_file="/etc/udev/rules.d/70-persistent-net.rules"
    if [ -e $net_rule_file ]; then
        result=$(grep -vE '^$|^#' $net_rule_file)
        if [ -n "$result" ]; then
            write_file "\tERROR: please empty or delete the rule file: $net_rule_file"
            [ $ret -eq 1 ] || ret=1
        fi
    fi
    return $ret
}

# file system check
function check_ext4_csum(){
    root_dev=$(cat /etc/fstab | grep -v "#" | grep -v "^;" | grep "ext4" | head -n 1 | awk '{print $1}')
    if [ -z $root_dev ]; then
        return 0
    fi
    ret=$(echo $root_dev | grep "UUID")
    if [ "$ret" != "" ]; then
        dev_uuid=$(echo $root_dev | awk -F "=" '{print $2}')
        root_dev=$(blkid | grep $dev_uuid | awk -F ':' '{print $1}')
    fi
    which dumpe2fs 1>/dev/null 2>&1
    if [ "$?" -eq 0 ]; then
        dumpe2fs $root_dev 2>/dev/null | grep "Filesystem features" | grep "metadata_csum" 1>/dev/null 2>&1
        [ "$?" -eq 0 ] && return 1 || return 0
    fi
}

function get_fs_type(){
    if [ -e /etc/fstab ]; then
        ret=$(cat /etc/fstab | grep -v "#" | grep -v "^;" | grep -v -E "ext2|ext3|ext4|swap|tmpfs|devpts|sysfs|proc|debugfs")
        if [ -n "$ret" ]; then
            return 1
        else
            return 0
        fi
    else
        return 2
    fi
}

function check_fs(){
    write_file $split_line
    write_file "[ check file system ]"
    get_fs_type
    case "$?" in
    0)
        check_ext4_csum
        if [ "$?" -eq 0 ]; then
            return $SUCCESS
        else
            write_file "\tERROR: bcc don't support ext4 with the features of metadata_csum."
            return $FAILED
        fi
        ;;
    1)
        write_file "\tERROR: your device doesn't use ext2/ext3/ext4 as file system."
        return $FAILED
        ;;
    2)
        write_file  "\tERROR: /etc/fstab doesn't exist."
        return $FAILED
        ;;
    *)
        write_file "\tERROR: unknown error."
        return $FAILED
    esac
}

function check_sshd_service(){
    write_file $split_line
    write_file "[ check sshd service ]"
    dpkg -l | grep -q openssh-server
    if [ $? -ne 0 ]; then
        write_file "\tWARNING: the openssh-server isn't installed."
        return $WARNING
    fi

    sshd_cfg="/etc/ssh/sshd_config"
    if [ ! -e "$sshd_cfg" ]; then
        write_file "\tWARNING: the sshd_config file doesn't exist."
        return $WARNING
    fi

    grep -iq "^PermitRootLogin yes" /etc/ssh/sshd_config
    if [ "$?" -ne 0 ]; then
        write_file "\tWARNING: the setting of \"PermitRootLogin\" must be to yes"
        return $WARNING
    fi

    if [ "$version" == "16" ] || [ "$version" = "18" ]; then
        systemctl status sshd | grep -qE "ssh.service; enabled|running"
        if [ "$?" -ne 0 ]; then
            write_file "\tWARNING: the sshd service is not running."
            return $WARNING
        fi
    else
        service ssh status | grep -q "running" 
        if [ "$?" -ne 0 ]; then
            write_file "\tWARNING: the ssh service is not running."
            return $WARNING
        fi
    fi
    return $SUCCESS
}

function check_firewall_status(){
    write_file $split_line
    write_file "[ check firewall ]"
    which ufw &>/dev/null
    if [ "$?" -ne 0 ]; then
        write_file "\tERROR: can't find ufw command."
        return $FAILED
    fi
    ufw status | grep -q "inactive"
    if [ "$?" -eq 0 ]; then
        return $SUCCESS
    else
        write_file "\tERROR: your system's firewall may be enabled."
        return $FAILED
    fi
}

function check_disk_uuid(){
    write_file $split_line
    write_file "[ check disk uuid ]"
    ret=0
    cat /etc/fstab | grep -v "#" | grep -v "^;"| grep -v "^$" | grep -i "dev" | awk '{ print $2 }' | grep "^/$" &> /dev/null
    if [ "$?" -eq 0 ];then
        write_file "\tERROR: in /etc/fstab, your system doesn't use UUID to mount \"/\"."
        ret=1
    fi
    if [ -n "$s_grub_file" ]; then
        cat $s_grub_file | grep -v "^#" | grep `uname -r` | grep "root" | awk -F "=" '{print $2}' |  awk '{print $1}' | grep -q "UUID"
        if [ "$?" -ne 0 ]; then
            write_file "\tERROR: in grub file: $s_grub_file, your system doesn't mount the root partition with UUID." 
            [ $ret -eq 1 ] || ret=1
        fi
    fi
    return $ret
}

function check_disk_size(){
    write_file $split_line
    write_file "[ check disk size ]"
    which fdisk 1>/dev/null 2>&1
    if [ "$?" -ne 0 ]; then
        write_file "\tERROR: can't find the fdisk command on your system."
        wirte_LOG "\tERROR: please confirm that the disk size is not exceed 500GB."
        return $FAILED
    fi
    disk_size=`fdisk -l 2>/dev/null |awk '/^Disk.*bytes/&&/\/dev/{ printf "%d",$3;print "" }' | head -n 1`
    size_unit=`fdisk -l 2>/dev/null |awk '/^Disk.*bytes/&&/\/dev/{ printf $4 ;print "" }' | awk -F ',' '{print $1}' | head -n 1`
    if [ $size_unit == "TB" ]; then
        disk_size=$((disk_size*1024))
    elif [ $size_unit == "MB" ]; then
        disk_size=$((disk_size/1024))
    fi

    if [ "$disk_size" -le 500 ]; then
        return $SUCCESS
    else
        write_file "\tERROR: the disk size is ${disk_size}GB."
        write_file "\tERROR: the size of the disk on your system must be less than 500G."
        return $FAILED
    fi
}

function check_cloud-int(){
    write_file $split_line
    write_file "[ check cloud-init ]"
    which cloud-init &> /dev/null
    if [ "$?" -ne 0 ]; then
        write_file "\tWARNING: the cloud-init is not installed."
        return $WARNING
    fi
    return $SUCCESS
}

echo $split_line
echo "Distro: $distro"
echo "Version: $full_version"
echo $split_line

echo -n "[ check virtio driver ]"
check_virtio_driver
print_result $?

echo -n "[ check /etc/shadow ]"
check_shadow
print_result $?

echo -n "[ check selinux ]"
check_selinux
print_result $?

echo -n "[ check network config ]"
check_network_config
print_result $?

echo -n "[ check sshd service ]"
check_sshd_service
print_result $?

echo -n "[ check firewall ]"
check_firewall_status
print_result $?

echo -n "[ check file system ]"
check_fs
print_result $?

echo -n "[ check root ]"
check_root
print_result $?

echo -n "[ check passwd command ]"
check_passwd_cmd
print_result $?

echo -n "[ check disk partition ]"
check_partition_type
print_result $?

echo -n "[ check LVM ]"
check_lvm
print_result $?

echo -n "[ check cloud-init ]"
check_cloud-int
print_result $?

echo -n "[ check disk uuid ]"
check_disk_uuid
print_result $?

echo -n "[ check disk size ]"
check_disk_size
print_result $?

echo -e $split_line >> $msg_file
echo $split_line
if [ "$c_wrn" -eq 0 ] && [ "$c_err" -eq 0 ]; then
    echo "There is no error/waring."
    echo "The system can be imported into Baidu cloud."
    rm -f $msg_file
else
    echo "There are $c_wrn warning(s), $c_err error(s)."
    echo "The reslult file is: $msg_file"
    echo "Please check result file for the error/warning information."
fi
echo $split_line