#!/bin/bash
# Ver: 1.4 by Endial Fang (endial@126.com)
#
# 操作系统控制函数库

# 加载依赖项
source "/usr/local/lib/liblog.sh"
source "/usr/local/lib/libfs.sh"
source "/usr/local/lib/libvalidations.sh"

# 函数列表

# 检测指定用户账户是否存在
# 参数:
#   $1 - 用户账户
# 返回值:
#   0 / 1
is_user_exists() {
    local user="${1:?user is missing}"
    id "$user" >/dev/null 2>&1
}

# 检测指定用户分组是否存在
# 参数:
#   $1 - 用户组
# 返回值:
#   0 / 1
is_group_exists() {
    local group="${1:?group is missing}"
    getent group "$group" >/dev/null 2>&1
}

# 检测当前是否为 root 用户
# 返回值:
#   true / false
is_root() {
    if [[ "$(id -u)" = "0" ]]; then
        debug "Run as root."
        true
    else
        debug "Run as non-root: $(id -u)"
        false
    fi
}

# 确保指定用户组在系统中存在
# 参数:
#   $1 - 用户组
# 标志位:
#   -i|--gid - 用户组 ID
#   -s|--system - 创建系统用户 (uid <= 999)
ensure_group_exists() {
    local group="${1:?group is missing}"
    local gid=""
    local is_system_user=false

    # 检测标志位
    shift 1
    while [ "$#" -gt 0 ]; do
        case "$1" in
        -i | --gid)
            shift
            gid="${1:?missing gid}"
            ;;
        -s | --system)
                is_system_user=true
                ;;
            *)
                echo "Invalid command line flag $1" >&2
                return 1
                ;;
        esac
        shift
    done

    if ! group_exists "$group"; then
        local -a args=("$group")
        if [[ -n "$gid" ]]; then
            if group_exists "$gid"; then
                error "The GID $gid is already in use." >&2
                return 1
            fi
            args+=("--gid" "$gid")
        fi
        $is_system_user && args+=("--system")
        groupadd "${args[@]}" >/dev/null 2>&1
    fi
}

# 确保指定用户在系统中存在
# 参数:
#   $1 - 用户
# 标志位:
#   -i|--uid - 用户 ID
#   -a|--append-groups - 附加用户组 (逗号分隔)
#   -g|--group - 用户组
#   -h|--home - 用户家目录
#   -s|--system - 创建系统用户 (uid <= 999)
ensure_user_exists() {
    local user="${1:?user is missing}"
    local uid=""
    local group=""
    local append_groups=""
    local home=""
    local is_system_user=false

    # Validate arguments
    shift 1
    while [ "$#" -gt 0 ]; do
        case "$1" in
        -i | --uid)
            shift
            uid="${1:?missing uid}"
            ;;
        -g | --group)
                shift
                group="${1:?missing group}"
                ;;
        -a | --append-groups)
            shift
            append_groups="${1:?missing append_groups}"
            ;;
        -h | --home)
                shift
                home="${1:?missing home directory}"
                ;;
        -s | --system)
                is_system_user=true
                ;;
            *)
                echo "Invalid command line flag $1" >&2
                return 1
                ;;
        esac
        shift
    done

    if ! user_exists "$user"; then
        local -a user_args=("-N" "$user")
        if [[ -n "$uid" ]]; then
            if user_exists "$uid"; then
                error "The UID $uid is already in use."
                return 1
            fi
            user_args+=("--uid" "$uid")
        else
        $is_system_user && user_args+=("--system")
        fi
        useradd "${user_args[@]}" >/dev/null 2>&1
    fi

    if [[ -n "$group" ]]; then
        local -a group_args=("$group")
        $is_system_user && group_args+=("--system")
        ensure_group_exists "${group_args[@]}"
        usermod -g "$group" "$user" >/dev/null 2>&1
    fi

    if [[ -n "$append_groups" ]]; then
        local -a groups
        read -ra groups <<<"$(tr ',;' ' ' <<<"$append_groups")"
        for group in "${groups[@]}"; do
            ensure_group_exists "$group"
            usermod -aG "$group" "$user" >/dev/null 2>&1
        done
    fi

    if [[ -n "$home" ]]; then
        mkdir -p "$home"
        usermod -d "$home" "$user" >/dev/null 2>&1
        configure_permissions_ownership "$home" -d "775" -f "664" -u "$user" -g "$group"
    fi
}

# 获取系统可用内存大小(MB)信息
# 返回值:
#   内存大小(兆字节)
get_total_memory() {
    echo $(($(grep MemTotal /proc/meminfo | awk '{print $2}') / 1024))
}

# 获取以内存定量方式描述的机器类型
# 标志位:
#   --memory - 内存大小 (MB，可选)
# 返回值:
#   类型名称
get_machine_size() {
    local memory=""

    # 检测标志位
    while [[ "$#" -gt 0 ]]; do
        case "$1" in
            --memory)
                shift
                memory="${1:?missing memory}"
                ;;
            *)
                echo "Invalid command line flag $1" >&2
                return 1
                ;;
        esac
        shift
    done

    if [[ -z "$memory" ]]; then
        debug "Memory was not specified, detecting available memory automatically"
        memory="$(get_total_memory)"
    fi
    sanitized_memory=$(convert_to_mb "$memory")
    if [[ "$sanitized_memory" -gt 26000 ]]; then
        echo 2xlarge
    elif [[ "$sanitized_memory" -gt 13000 ]]; then
        echo xlarge
    elif [[ "$sanitized_memory" -gt 6000 ]]; then
        echo large
    elif [[ "$sanitized_memory" -gt 3000 ]]; then
        echo medium
    elif [[ "$sanitized_memory" -gt 1500 ]]; then
        echo small
    else
        echo micro
    fi
}

# 获取已定义的所有内存大小描述
# 返回值:
#   描述值列表
get_supported_machine_sizes() {
    echo micro small medium large xlarge 2xlarge
}

# 将以字符串表示的内存大小转换为以MB为单位的内存大小值 (i.e. 2G -> 2048)
# 参数:
#   $1 - 内存大小
# 返回值:
#   转换后的数值
convert_to_mb() {
    local amount="${1:-}"
    if [[ $amount =~ ^([0-9]+)(m|M|g|G) ]]; then
        size="${BASH_REMATCH[1]}"
        unit="${BASH_REMATCH[2]}"
        if [[ "$unit" = "g" || "$unit" = "G" ]]; then
            amount="$((size * 1024))"
        else
            amount="$size"
        fi
    fi
    echo "$amount"
}

# 如果禁用调试模式，将输出信息重定向至 /dev/null
# 参数:
#   $@ - 待执行的命令
debug_execute() {
    if is_boolean_yes "${ENV_DEBUG:-false}"; then
        "$@"
    else
        "$@" >/dev/null 2>&1
    fi
}

# 重试执行命令
# 参数:
#   $1 - 命令 (字符串)
#   $2 - 最大尝试次数. 默认值: 12
#   $3 - 重试前等待时间(秒). 默认值: 5
# 返回值:
#   0 / 1
retry_while() {
    local cmd="${1:?cmd is missing}"
    local retries="${2:-12}"
    local sleep_time="${3:-5}"
    local return_value=1

    read -r -a command <<<"$cmd"
    for ((i = 1; i <= retries; i += 1)); do
        "${command[@]}" && return_value=0 && break
        sleep "$sleep_time"
    done
    return $return_value
}

# 生成随机字符串
# 标志位:
#   -t|--type - 字符串类型 (ascii, alphanumeric, numeric). 默认值: ascii
#   -c|--count - 字符串长度. 默认值: 32
# 返回值:
#   字符串
generate_random_string() {
    local type="ascii"
    local count="32"
    local filter
    local result

    # 检测标志位
    while [[ "$#" -gt 0 ]]; do
        case "$1" in
        -t | --type)
            shift
            type="$1"
            ;;
        -c | --count)
            shift
            count="$1"
            ;;
        *)
            echo "Invalid command line flag $1" >&2
            return 1
            ;;
        esac
        shift
    done

    # 检测类型
    case "$type" in
    ascii)
        filter="[:print:]"
        ;;
    numeric)
        filter="0-9"
        ;;
    alphanumeric)
        filter="a-zA-Z0-9"
        ;;
    alphanumeric+special|special+alphanumeric)
        # Limit variety of special characters, so there is a higher chance of containing more alphanumeric characters
        # Special characters are harder to write, and it could impact the overall UX if most passwords are too complex
        filter='a-zA-Z0-9:@.,/+!='
        ;;
    *)
        echo "Invalid type ${type}" >&2
        return 1
        ;;
    esac
    # Obtain count + 10 lines from /dev/urandom to ensure that the resulting string has the expected size
    # Note there is a very small chance of strings starting with EOL character
    # Therefore, the higher amount of lines read, this will happen less frequently
    result="$(head -n "$((count + 10))" /dev/urandom | tr -dc "$filter" | head -c "$count")"
    echo "$result"
}

# 为指定字符串生成 MD5 值
# 参数:
#   $1 - 字符串
# 返回值:
#   字符串对应的 MD5
generate_md5_hash() {
  local -r str="${1:?missing input string}"
  echo -n "$str" | md5sum | awk '{print $1}'
}

# 为指定字符串生成 SHA1 值
# 参数:
#   $1 - 字符串
#   $2 - 算法 - 1 (default), 224, 256, 384, 512
# 返回值:
#   字符串对应的 SHA1
generate_sha_hash() {
    local -r str="${1:?missing input string}"
    local -r algorithm="${2:-1}"
    echo -n "$str" | "sha${algorithm}sum" | awk '{print $1}'
}

# 将一个字符串转换为其十六进制表示形式
# 参数:
#   $1 - 字符串
# 返回值:
#   十六进制表示形式
convert_to_hex() {
    local -r str=${1:?missing input string}
    local -i iterator
    local char
    for ((iterator = 0; iterator < ${#str}; iterator++)); do
        char=${str:iterator:1}
        printf '%x' "'${char}"
    done
}

# 获取启动时间
# 返回值:
#   启动时间
get_boot_time() {
    stat /proc --format=%Y
}

# 获取机器 ID
# 返回值:
#   机器 ID
get_machine_id() {
    local machine_id
    if [[ -f /etc/machine-id ]]; then
        machine_id="$(cat /etc/machine-id)"
    fi
    if [[ -z "$machine_id" ]]; then
        # Fallback to the boot-time, which will at least ensure a unique ID in the current session
        machine_id="$(get_boot_time)"
    fi
    echo "$machine_id"
}

# 获取根分区的磁盘设备 ID（如： /dev/sda1）
# 返回值:
#   根分区磁盘设备 ID
get_disk_device_id() {
    local device_id=""
    if grep -q ^/dev /proc/mounts; then
        device_id="$(grep ^/dev /proc/mounts | awk '$2 == "/" { print $1 }' | tail -1)"
    fi
    # If it could not be autodetected, fallback to /dev/sda1 as a default
    if [[ -z "$device_id" || ! -b "$device_id" ]]; then
        device_id="/dev/sda1"
    fi
    echo "$device_id"
}

# 获取根磁盘设备 ID（如： /dev/sda）
# 返回值:
#   根磁盘设备 ID
get_root_disk_device_id() {
    get_disk_device_id | sed -E 's/p?[0-9]+$//'
}

# 获取根磁盘大小
# 返回值:
#   根磁盘大小（字节）
get_root_disk_size() {
    fdisk -l "$(get_root_disk_device_id)" | grep 'Disk.*bytes' | sed -E 's/.*, ([0-9]+) bytes,.*/\1/' || true
}

# 以指定用户身份运行命令
# 注意：此函数将使用 chroot 命令来模拟指定用户身份运行命令，
# 参数：
#   $1 - 用户（可选）
#   $2..$n - 命令及命令参数
run_as_user() {
    run_chroot "$@"
}

# 以指定用户身份执行命令
# 注意：此命令将替换当前进程映像，
# 参数：
#   $1 - 用户（可选）
#   $2..$n - 命令及命令参数
exec_as_user() {
    run_chroot --replace-process "$@"
}

# 使用 chroot 命令以指定用户身份运行命令
# 参数：
#   $1 - 用户（可选）
#   $2..$n - 命令及命令参数
# 标志：
#   -r | --replace-process - 替换当前进程映像
run_chroot() {
    local userspec
    local user
    local homedir
    local replace=false
    local -r cwd="$(pwd)"

    # 解析并验证标志位
    while [[ "$#" -gt 0 ]]; do
        case "$1" in
            -r | --replace-process)
                replace=true
                ;;
            --)
                shift
                break
                ;;
            -*)
                stderr_print "unrecognized flag $1"
                return 1
                ;;
            *)
                break
                ;;
        esac
        shift
    done

    # 解析并验证参数
    if [[ "$#" -lt 2 ]]; then
        echo "expected at least 2 arguments"
        return 1
    else
        userspec=$1
        shift

        # 用户名可以包含组，因此我们解析用户
        user=$(echo "$userspec" | cut -d':' -f1)
    fi

    if ! is_root; then
        error "Could not switch to '${userspec}': Operation not permitted"
        return 1
    fi

    # 获取用户的家目录，因为 chroot 不会正确更新此环境变量，并且一些脚本依赖于它
    homedir=$(eval echo "~${user}")
    if [[ ! -d $homedir ]]; then
        homedir="${HOME:-/}"
    fi

    # 取得 "$@" 的值，以便正确支持 shell 参数扩展
    if [[ "$replace" = true ]]; then
        exec chroot --userspec="$userspec" / bash -c "cd ${cwd}; export HOME=${homedir}; exec \"\$@\"" -- "$@"
    else
        chroot --userspec="$userspec" / bash -c "cd ${cwd}; export HOME=${homedir}; exec \"\$@\"" -- "$@"
    fi
}