#!/bin/bash
#
# Vuls Scanner Update Script
# Reads configuration from /etc/vuls-scanner.conf and updates enabled databases
# Combines all updater functionality into a single script
#

set -euo pipefail

CONFIG_FILE="/etc/vuls-scanner.conf"
VULS_LIBDIR="/var/lib/vuls"
MAX_AGE_HOURS=24
FORCE=0
TMPDIR=""

usage() {
    cat <<EOF
Usage: $0 [OPTIONS]

Update Vuls vulnerability databases based on configuration in $CONFIG_FILE

Options:
    --force          Force update even if databases are less than 24 hours old
    -h, --help       Show this help message

Configuration:
    Edit $CONFIG_FILE to enable/disable specific databases and configure
    additional arguments for each updater.

EOF
    exit 1
}

# Cleanup function for temporary directory
cleanup() {
    if [[ -n "$TMPDIR" && -d "$TMPDIR" ]]; then
        rm -rf "$TMPDIR"
    fi
}

# Set up trap to cleanup on exit
trap cleanup EXIT

# Function to check database age
check_db_age() {
    local db_file="$1"
    local force="$2"
    
    if [[ ! -f "$db_file" ]]; then
        return 0  # File doesn't exist, need to update
    fi
    
    if [[ $force -eq 1 ]]; then
        return 0  # Force update
    fi
    
    # Get file modification time
    local file_age=$(( $(date +%s) - $(stat -c %Y "$db_file") ))
    local max_age_seconds=$(( MAX_AGE_HOURS * 3600 ))
    
    if [[ $file_age -lt $max_age_seconds ]]; then
        local age_hours=$(( file_age / 3600 ))
        echo "Database $db_file is less than $MAX_AGE_HOURS hours old (${age_hours}h). Skipping update."
        echo "Use --force to update anyway."
        return 1
    fi
    
    return 0
}

# Function to update a database
update_database() {
    local name="$1"
    local db_file="$2"
    local updater_cmd="$3"
    local default_args_str="$4"
    local config_args="$5"
    
    # Check if updater is installed
    if ! command -v "$updater_cmd" >/dev/null 2>&1; then
        echo "Warning: $updater_cmd is not installed, skipping $name update" >&2
        return 1
    fi
    
    # Check database age
    if ! check_db_age "$db_file" "$FORCE"; then
        return 0
    fi
    
    # Create temporary file for download
    local db_basename=$(basename "$db_file")
    local tmp_db_file="$TMPDIR/${db_basename}.tmp"
    
    # Build arguments array
    local args=()
    if [[ -n "${config_args:-}" ]]; then
        # Split config_args into array
        read -ra args <<< "$config_args"
    elif [[ -n "${default_args_str:-}" ]]; then
        # Split default_args_str into array
        read -ra args <<< "$default_args_str"
    fi
    
    # Run the updater to temporary file
    echo "Updating $name database..."
    cd "$TMPDIR"
    if [[ ${#args[@]} -gt 0 ]]; then
        "$updater_cmd" fetch "${args[@]}" --dbpath "$tmp_db_file"
    else
        "$updater_cmd" fetch --dbpath "$tmp_db_file"
    fi
    
    # Verify the temporary database file was created and is not empty
    if [[ ! -f "$tmp_db_file" ]] || [[ ! -s "$tmp_db_file" ]]; then
        echo "Error: Failed to create $name database file" >&2
        rm -f "$tmp_db_file"
        return 1
    fi
    
    # Move temporary file to final location atomically
    mv "$tmp_db_file" "$db_file"
    
    # Set correct permissions
    chown vuls:vuls "$db_file" 2>/dev/null || true
    chmod 644 "$db_file" 2>/dev/null || true
    
    echo "$name database update complete."
    return 0
}

# Parse arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --force)
            FORCE=1
            shift
            ;;
        -h|--help)
            usage
            ;;
        *)
            echo "Unknown option: $1"
            usage
            ;;
    esac
done

# Source configuration file
if [[ ! -f "$CONFIG_FILE" ]]; then
    echo "Error: Configuration file not found: $CONFIG_FILE" >&2
    exit 1
fi

source "$CONFIG_FILE"

# Ensure vuls lib directory exists and has correct permissions
if [[ ! -d "$VULS_LIBDIR" ]]; then
    echo "Error: Vuls library directory not found: $VULS_LIBDIR" >&2
    echo "Please ensure the vuls package is installed." >&2
    exit 1
fi

# Create temporary directory for database downloads
TMPDIR=$(mktemp -d -t vuls-update.XXXXXX)
if [[ ! -d "$TMPDIR" ]]; then
    echo "Error: Failed to create temporary directory" >&2
    exit 1
fi

# Update enabled databases
if [[ "${ENABLE_CVE:-no}" == "yes" ]]; then
    update_database "CVE Dictionary" \
        "$VULS_LIBDIR/cve.sqlite3" \
        "go-cve-dictionary" \
        "nvd" \
        "${CVE_ARGS:-}"
fi

if [[ "${ENABLE_OVAL:-no}" == "yes" ]]; then
    update_database "OVAL Dictionary" \
        "$VULS_LIBDIR/oval.sqlite3" \
        "goval-dictionary" \
        "redhat" \
        "${OVAL_ARGS:-}"
fi

if [[ "${ENABLE_GOST:-no}" == "yes" ]]; then
    update_database "GOST" \
        "$VULS_LIBDIR/gost.sqlite3" \
        "gost" \
        "debian redhat microsoft ubuntu" \
        "${GOST_ARGS:-}"
fi

if [[ "${ENABLE_EXPLOITDB:-no}" == "yes" ]]; then
    update_database "Exploit Database" \
        "$VULS_LIBDIR/go-exploitdb.sqlite3" \
        "go-exploitdb" \
        "" \
        "${EXPLOITDB_ARGS:-}"
fi

if [[ "${ENABLE_MSFDB:-no}" == "yes" ]]; then
    update_database "Metasploit Database" \
        "$VULS_LIBDIR/go-msfdb.sqlite3" \
        "go-msfdb" \
        "" \
        "${MSFDB_ARGS:-}"
fi

if [[ "${ENABLE_KEV:-no}" == "yes" ]]; then
    update_database "KEV" \
        "$VULS_LIBDIR/go-kev.sqlite3" \
        "go-kev" \
        "" \
        "${KEV_ARGS:-}"
fi

if [[ "${ENABLE_CTI:-no}" == "yes" ]]; then
    update_database "CTI" \
        "$VULS_LIBDIR/go-cti.sqlite3" \
        "go-cti" \
        "" \
        "${CTI_ARGS:-}"
fi

# Update Windows WSUS CAB (wsusscn2.cab) if enabled
update_wsusscn2_cab() {
    local cab_file="$VULS_LIBDIR/wsusscn2.cab"
    local tmp_file="$TMPDIR/wsusscn2.cab.tmp"
    local min_size=$(( 1024 * 1024 ))  # 1MB minimum
    local url="${WSUSCAB_URL:-http://download.windowsupdate.com/microsoftupdate/v6/wsusscan/wsusscn2.cab}"
    local url_alt="https://catalog.update.microsoft.com/v7/site/wsusscn2.cab"

    if ! check_db_age "$cab_file" "$FORCE"; then
        return 0
    fi

    if ! command -v wget >/dev/null 2>&1 && ! command -v curl >/dev/null 2>&1; then
        echo "Warning: neither wget nor curl is available, skipping wsusscn2.cab update" >&2
        return 1
    fi

    echo "Updating Windows WSUS CAB (wsusscn2.cab)..."
    local last_err=""
    for attempt in 1 2; do
        local try_url="$url"
        [[ $attempt -gt 1 && -z "${WSUSCAB_URL:-}" ]] && try_url="$url_alt"
        if command -v wget >/dev/null 2>&1; then
            if wget -q -O "$tmp_file" "$try_url" 2>/dev/null; then
                break
            fi
            last_err="wget failed for $try_url"
        else
            if curl -sSfL -o "$tmp_file" "$try_url" 2>/dev/null; then
                break
            fi
            last_err="curl failed for $try_url"
        fi
        rm -f "$tmp_file"
        [[ $attempt -lt 2 ]] && sleep 5
    done

    if [[ ! -f "$tmp_file" ]] || [[ $(stat -c %s "$tmp_file" 2>/dev/null || echo 0) -lt $min_size ]]; then
        echo "Error: Failed to download wsusscn2.cab (${last_err:-file too small})" >&2
        rm -f "$tmp_file"
        return 1
    fi

    mv "$tmp_file" "$cab_file"
    chown vuls:vuls "$cab_file" 2>/dev/null || true
    chmod 644 "$cab_file" 2>/dev/null || true
    echo "Windows WSUS CAB update complete."
    return 0
}

if [[ "${ENABLE_WSUSCAB:-no}" == "yes" ]]; then
    update_wsusscn2_cab || true
fi

echo "Database update complete."
