#!/usr/bin/env 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

function get_freebsd_distro_version(){
    uname -a | grep -i -q "freebsd" >& /dev/null
    if [ "$?" -eq 0 ]; then
        distro="freebsd"
    else
        if [ -f "/etc/motd" ]; then
            grep -q "FreeBSD" /det/motd >& /dev/null
            if [ "$?" -eq 0 ]; then
                distro="freebsd"
            fi
        fi
    fi
    if [ "$distro" != "freebsd" ]; then
        echo "Please run the script on FreeBSD system."
        exit
    else
        full_version=$(uname -r | awk -F '-' '{print $1}')
        versioin=$(echo $full_version | awk -F '.' '{print $1}')
    fi
    if [ "$full_version" == "unknown" ]; then
        echo "Can't get the version of FreeBSD"
        exit
    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 ]"
    kldstat -v | grep -i -q "virtio_blk"
    if [ "$?" -ne 0 ]; then
        write_file "\tERROR: the virtio driver doesn't install."
        return $FAILED
    fi
    return $SUCCESS
}

function check_partition_type(){
    write_file $split_line
    write_file "[ check disk partition ]"
    which gpart 1>/dev/null 2>&1
    if [ "$?" -ne 0 ]; then
        write_file "\tERROR: can't find gpart command."
        write_file "\tERROR: please make sure that your system uses MBR partition."
        return $FAILED
    fi
    p_type=$(gpart list | grep scheme | head -n 1 | awk -F ':' '{print $2}' | sed s/[[:space:]]//g)
    if [ "$p_type" != "MBR" ]; then
        write_file "\tERROR: the Partition Table on your disk is $p_type."
        write_file "\tERROR: we recommend using MBR partition."
        return $FAILED
    fi
}

function check_root(){
    write_file $split_line
    write_file "[ check root user ]"
    grep -q "^root:*:" /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_network_config(){
    write_file $split_line
    write_file "[ check network config ]"
    interface=$(grep "ifconfig" /etc/rc.conf | awk -F '=' '{print $1}')
    net_type=$(grep "ifconfig" /etc/rc.conf | awk -F '=' '{print $2}' | sed 's/\"//g')
    if [ "$net_type" != "DHCP" ]; then
        write_file "\tERROR: in /etc/rc.config, $interface is not set to \"DHCP\"."
        return $FAILED
    fi
}

function check_sshd_service(){
    write_file $split_line
    write_file "[ check sshd service ]"

    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

    service sshd status | grep -q "is running"
    if [ "$?" -ne 0 ]; then
        write_file "\tWARNING: the sshd service is running."
        write_file "\tWARNING: please enable sshd service."
        return $WARNING
    fi
    return $SUCCESS
}

function check_firewall_status(){
    write_file $split_line
    write_file "[ check firewall ]"
    cat /etc/rc.conf | grep -iq 'firewall_enable="YES"'
    if [ "$?" -eq 0 ]; then
        write_file "\tERROR: the friewall is enabled on your system."
        return $FAILED
    fi
    cat /etc/rc.conf | grep -iq 'pf_enable="YES"'
    if [ "$?" -eq 0 ]; then
        write_file "\tERROR: the pf is enabled on your system."
        return $FAILED
    fi
    return $SUCCESS
}

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=`gpart list | grep -i "Mediasize" | head -n 1 | cut -d '(' -f2 | cut -d ')' -f1 | tr -cd "[0-9]"`
    size_unit=`gpart list | grep -i "Mediasize" | head -n 1 | cut -d '(' -f2 | cut -d ')' -f1 | tr -d "[0-9]"`
    if [ $size_unit == "T" ]; then
        disk_size=$((disk_size*1024))
    elif [ $size_unit == "M" ]; 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
}

get_freebsd_distro_version

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

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 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 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 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