#!/bin/bash
# author: Saiwa
# date: 2024-04-06

convert_bytes() {
    bytes=$1
    if [[ $bytes -lt 1024 ]]; then
        echo "${bytes} B"
    elif [[ $bytes -lt $((1024**2)) ]]; then
        printf "%.1f KB\n" $(awk "BEGIN { printf \"%.1f\", $bytes / 1024 }")
    elif [[ $bytes -lt $((1024**3)) ]]; then
        printf "%.1f MB\n" $(awk "BEGIN { printf \"%.1f\", $bytes / (1024^2) }")
    elif [[ $bytes -lt $((1024**4)) ]]; then
        printf "%.1f GB\n" $(awk "BEGIN { printf \"%.1f\", $bytes / (1024^3) }")
    else
        printf "%.1f TB\n" $(awk "BEGIN { printf \"%.1f\", $bytes / (1024^4) }")
    fi
}

get_cpu_info() {
    # 获取 CPU 型号和核心数
    cpu_name=$(lscpu | awk -F: '/Model name/ {gsub(/^[ \t]+|[ \t]+$/, "", $2); print $2}')
    cpu_socket=$(lscpu | awk -F: '/^Socket\(s\)/ {gsub(/^[ \t]+|[ \t]+$/, "", $2); print $2}')
    cpu_core_per_socket=$(lscpu | awk -F: '/Core\(s\) per socket/ {gsub(/^[ \t]+|[ \t]+$/, "", $2); print $2}')
    cpu_thread_per_core=$(lscpu | awk -F: '/Thread\(s\) per core/ {gsub(/^[ \t]+|[ \t]+$/, "", $2); print $2}')
    cpu_cores=$(lscpu | awk -F: '/^CPU\(s\)/ {gsub(/^[ \t]+|[ \t]+$/, "", $2); print $2}')

    echo "Model name,Socket(s),Core(s)/socket,Thread(s)/core,CPU(s)"
    echo "$cpu_name,$cpu_socket,$cpu_core_per_socket,$cpu_thread_per_core,$cpu_cores"
}

get_mem_info() {
    # 执行 dmidecode 命令获取内存设备信息，并用 awk 处理
    memory_info=$(dmidecode --type memory | awk '/Memory Device$/,/^$/')

    # 输出表头，# 处理每个内存设备的信息，并输出结果
    echo "Locator,Size,Speed,Type,Serial Number,Manufacturer"
    echo "$memory_info" | awk 'BEGIN {FS="\n"; RS=""} {
        gsub(/^.*: /, "", $9);
        gsub(/^.*: /, "", $6);
        gsub(/^.*: /, "", $13);
        gsub(/^.*: /, "", $11);
        gsub(/^.*: /, "", $15);
        gsub(/^.*: /, "", $14);
        print $9 "," $6 "," $13 "," $11 "," $15 "," $14
    }'
  }

get_disk_info() {
    # 获取系统中所有硬盘设备的列表
    disks_name=$(lsblk -do NAME | sed 1d)                           # 未使用
    disks_list=$(smartctl --scan -j | jq -r '.devices[].name')      # 未使用
    disks_type=$(smartctl --scan-open -j | jq -r '.devices[] | select(has("open_error") | not) | (.type + " " + .name)')

    # 遍历每个硬盘设备
    echo -e "DISK,SIZE,MODEL_NAME,SERIAL_NUMBER,POWER_TIME,ERROR_LOG"

    while IFS= read -r line; do
        disk_name=$(echo $line | cut -d' ' -f2)
        disk_type=$(echo $line | cut -d' ' -f1)
        # 使用 smartctl 获取硬盘信息并转换为 JSON 格式
        info=$(smartctl -a -d $disk_type $disk_name -j)

        # 提取型号、序列号、容量和健康状态
        model=$(echo "$info" | jq -r '.model_name')
        serial=$(echo "$info" | jq -r '.serial_number')
        health=$(echo "$info" | jq -r '.smart_status.passed')
        [[ $health == "true" ]] && health="OK"
        # 判断硬盘类型，提取启动时间
        if [[ $disk_name == /dev/nvme* ]]; then
            capacity=$(echo "$info" | jq -r '.nvme_total_capacity')
            power_on_hours=$(echo "$info" | jq -r '.nvme_smart_health_information_log.power_on_hours')
            error_log=$(echo "$info" | jq -r '.nvme_smart_health_information_log.media_errors')
        else
            capacity=$(echo "$info" | jq -r '.user_capacity.bytes')
            error_log=$(echo "$info" | jq -r '.ata_smart_error_log.summary.count')
            if [[ $disk_type == "scsi" ]]; then
                power_on_hours="null"
            else
                power_on_hours=$(echo "$info" | jq -r '.power_on_time.hours')
            fi
        fi
        # 将硬盘大小转换易读格式
        disk_size=$(convert_bytes $capacity)
        # 将启动时间转换为天数和小时
        if [[ -n $power_on_hours ]] && [[ $power_on_hours != "null" ]]; then
            year=$((power_on_hours / 8760))
            day=$((power_on_hours % 8760 / 24))
            hour=$((power_on_hours % 24))
            power_time="${year}年${day}天${hour}小时"
        else
            power_time="null"
        fi
        # 打印硬盘信息
        echo -e "$disk_name,$disk_size,$model,$serial,$power_time,$error_log"
    done <<< "$disks_type"
}

# check required commands
check_commands() {
    if ! whereis lscpu | grep "/lscpu" &> /dev/null; then
        yum install util-linux -y &> /dev/null || apt install util-linux -y &> /dev/null
    fi

    if ! whereis dmidecode | grep "/dmidecode" &> /dev/null; then
        yum install dmidecode -y &> /dev/null || apt install dmidecode -y &> /dev/null
    fi

    if ! whereis smartctl | grep "/smartctl" &> /dev/null; then
        yum install smartmontools -y &> /dev/null || apt install smartmontools -y &> /dev/null
    fi
    if ! whereis jq | grep "/jq" &> /dev/null; then
        yum install jq -y &> /dev/null || apt install jq -y &> /dev/null || \
        (
        wget -O /usr/bin/jq https://github.com/jqlang/jq/releases/download/jq-1.7/jq-linux-amd64 &> /dev/null || \
        curl -o /usr/bin/jq https://github.com/jqlang/jq/releases/download/jq-1.7/jq-linux-amd64 &> /dev/null;
        chmod +x /usr/bin/jq;
        )
    fi

    if ! whereis lspci | grep "/lspci" &> /dev/null; then
        yum install pciutils -y &> /dev/null || apt install pciutils -y &> /dev/null
    fi
}


# try to get cpu info
try_get_cpu_info() {
    if whereis lscpu | grep "/lscpu" &> /dev/null; then
        echo -e "\n\033[0;34m===============  CPU Info at $(date "+%Y-%m-%d %H:%M:%S")  =================\033[0m" | \
        tee -a /tmp/hardware_info.log
        get_cpu_info | column -t -s ',' | tee -a /tmp/hardware_info.log
    else
        echo -e "\033[33m lscpu command not found, can not get CPU info \033[0m"
    fi
}

# try to get memory info
try_get_memory_info() {
    if whereis dmidecode | grep "/dmidecode" &> /dev/null; then
        echo -e "\n\033[0;34m===============  Memory Info at $(date "+%Y-%m-%d %H:%M:%S")  ==============\033[0m" | \
        tee -a /tmp/hardware_info.log
        get_mem_info | column -t -s ',' | grep -v "No Module" | tee -a /tmp/hardware_info.log
        echo -e "\033[1;36m$(lsmem | grep -E "Total online memory.*$")\033[0m" | tee -a /tmp/hardware_info.log
    else
        echo -e "\033[33m dmidecode command not found, can not get Memory info \033[0m"
    fi
}

# try to get disk info
try_get_disk_info() {
    if whereis smartctl | grep "/smartctl" &> /dev/null && whereis jq | grep "/jq" &> /dev/null; then
        echo -e "\n\033[0;34m===============  Disk Info at $(date "+%Y-%m-%d %H:%M:%S")  ================\033[0m" | \
        tee -a /tmp/hardware_info.log
        get_disk_info | column -t -s ',' | tee -a /tmp/hardware_info.log
    else
        echo -e "\033[33m smartctl or jq command not found, can not get Disk info \033[0m"
    fi
}

# try to get disk info list
try_get_disk_info_list() {
    if whereis smartctl | grep "/smartctl" &> /dev/null && whereis jq | grep "/jq" &> /dev/null; then
        echo -e "\n\033[0;34m===============  Disk Info at $(date "+%Y-%m-%d %H:%M:%S")  ================\033[0m" | \
        tee -a /tmp/hardware_info.log
        get_disk_info
    else
        echo -e "\033[33m smartctl or jq command not found, can not get Disk info \033[0m"
    fi
}

# try to get gpu info
try_get_gpu_info() {
    if whereis lspci | grep "/lspci" &> /dev/null; then
        echo -e "\n\033[0;34m===============  GPU Info at $(date "+%Y-%m-%d %H:%M:%S")  =================\033[0m" | \
        tee -a /tmp/hardware_info.log
        lspci | egrep -i "VGA|NVIDIA" | awk -F: '{print $3}' | tee -a /tmp/hardware_info.log
    else
        echo -e "\033[33m lspci command not found, can not get GPU info \033[0m"
    fi
}

# try to get network info
try_get_nic_info() {
    echo -e "\n\033[0;34m===============  NIC Info at $(date "+%Y-%m-%d %H:%M:%S")  =================\033[0m" | \
    tee -a /tmp/hardware_info.log
    if [[ $(uname -r) =~ "5.4.166-1.el7.elrepo.x86_64" ]]; then
        ip -c addr show type bond_slave | tee -a /tmp/hardware_info.log;
        ip -c addr show type bond | tee -a /tmp/hardware_info.log;
    else
        ip -c addr show | tee -a /tmp/hardware_info.log;
    fi
}

# try to get stress test info
try_get_stress_test_info() {
    if [[ -f /root/stressapptest.log ]]; then
        echo -e "\n\033[0;34m===============  Stress test Info at $(date "+%Y-%m-%d %H:%M:%S")  =========\033[0m" | \
        tee -a /tmp/hardware_info.log
        grep -v remaining /root/stressapptest.log  | tee -a /tmp/hardware_info.log
    fi
}

# Start, first check command, then get hardware info
if [[ $EUID -ne 0 ]]; then
    echo -e "\033[33m This script must be run as root \033[0m"
    exit 1
else
    check_commands
fi

if [[ $# -eq 0 ]]; then
    try_get_cpu_info
    try_get_memory_info
    try_get_disk_info
    try_get_gpu_info
elif [[ $# -eq 1 ]]; then
    case $1 in
        cpu|-cpu|--cpu|-c)
            try_get_cpu_info
            ;;
        mem|memory|-mem|--mem|-m)
            try_get_memory_info
            ;;
        disk|-disk|--disk|-d)
            try_get_disk_info
            ;;
        disk_list)
            try_get_disk_info_list
            ;;
        gpu|-gpu|--gpu|-g)
            try_get_gpu_info
            ;;
        nic|-nic|--nic|network|ip|-n)
            try_get_nic_info
            ;;
        test|-test|--test|stress|-t)
            try_get_stress_test_info
            ;;
        all|-all|--all|-a)
            try_get_cpu_info
            try_get_memory_info
            try_get_disk_info
            try_get_gpu_info
            try_get_nic_info
            try_get_stress_test_info
            ;;
        *)
            echo "Usage: $0 {cpu|mem|disk|nic|test|all}"
            ;;
    esac
else
    echo "Usage: $0 {cpu|mem|disk|nic|test|all}"
    exit 1
fi