#!/usr/bin/env bash
# ---------------------------------------------------------------------------
# triage.sh - Quick suspicious file triage using native tools
#
# Falls back to native CLI tools (file, sha256sum, strings) when Python
# or the triage.py dependencies are not available.
#
# Usage:
#   bash triage.sh <file_path> [--vt-lookup]
# ---------------------------------------------------------------------------

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

usage() {
    echo "Usage: $0 <file_path> [--vt-lookup]"
    echo ""
    echo "Perform quick triage of a suspicious file."
    echo ""
    echo "Options:"
    echo "  --vt-lookup   Query VirusTotal for the file hash (requires VT_API_KEY)"
    exit 1
}

json_escape() {
    # Minimal JSON string escaping
    printf '%s' "$1" | python3 -c 'import json,sys; print(json.dumps(sys.stdin.read()), end="")' 2>/dev/null \
        || printf '"%s"' "$(printf '%s' "$1" | sed 's/\\/\\\\/g; s/"/\\"/g; s/\t/\\t/g; s/\n/\\n/g')"
}

command_exists() {
    command -v "$1" >/dev/null 2>&1
}

# ---------------------------------------------------------------------------
# Try Python first
# ---------------------------------------------------------------------------

try_python() {
    local file_path="$1"
    local vt_flag="$2"

    if command_exists python3; then
        if python3 -c "import hashlib" 2>/dev/null; then
            local extra_args=""
            if [[ "$vt_flag" == "true" ]]; then
                extra_args="--vt-lookup"
            fi
            python3 "${SCRIPT_DIR}/triage.py" --file "$file_path" $extra_args
            return 0
        fi
    fi
    return 1
}

# ---------------------------------------------------------------------------
# Native triage
# ---------------------------------------------------------------------------

compute_hashes() {
    local file_path="$1"

    if command_exists md5sum; then
        MD5=$(md5sum "$file_path" | cut -d' ' -f1)
    elif command_exists md5; then
        MD5=$(md5 -q "$file_path")
    elif command_exists certutil; then
        MD5=$(certutil -hashfile "$file_path" MD5 2>/dev/null | sed -n '2p' | tr -d ' ')
    else
        MD5="unavailable"
    fi

    if command_exists sha1sum; then
        SHA1=$(sha1sum "$file_path" | cut -d' ' -f1)
    elif command_exists shasum; then
        SHA1=$(shasum -a 1 "$file_path" | cut -d' ' -f1)
    else
        SHA1="unavailable"
    fi

    if command_exists sha256sum; then
        SHA256=$(sha256sum "$file_path" | cut -d' ' -f1)
    elif command_exists shasum; then
        SHA256=$(shasum -a 256 "$file_path" | cut -d' ' -f1)
    elif command_exists certutil; then
        SHA256=$(certutil -hashfile "$file_path" SHA256 2>/dev/null | sed -n '2p' | tr -d ' ')
    else
        SHA256="unavailable"
    fi

    echo "\"md5\": \"$MD5\","
    echo "    \"sha1\": \"$SHA1\","
    echo "    \"sha256\": \"$SHA256\""
}

identify_type() {
    local file_path="$1"

    if command_exists file; then
        FILE_TYPE=$(file "$file_path" | cut -d: -f2- | sed 's/^ //')
        MIME_TYPE=$(file --mime-type "$file_path" 2>/dev/null | cut -d: -f2- | sed 's/^ //' || echo "unknown")
    else
        # Fallback: read magic bytes
        MAGIC=$(xxd -l 4 -p "$file_path" 2>/dev/null || od -A n -t x1 -N 4 "$file_path" 2>/dev/null | tr -d ' ')
        case "$MAGIC" in
            4d5a*|4D5A*) FILE_TYPE="PE executable"; MIME_TYPE="application/x-dosexec" ;;
            7f454c46*|7F454C46*) FILE_TYPE="ELF executable"; MIME_TYPE="application/x-elf" ;;
            504b0304*|504B0304*) FILE_TYPE="ZIP archive"; MIME_TYPE="application/zip" ;;
            25504446*|25504446*) FILE_TYPE="PDF document"; MIME_TYPE="application/pdf" ;;
            d0cf11e0*|D0CF11E0*) FILE_TYPE="OLE2 compound file"; MIME_TYPE="application/x-ole-storage" ;;
            *) FILE_TYPE="data"; MIME_TYPE="application/octet-stream" ;;
        esac
    fi

    echo "\"file_type\": $(json_escape "$FILE_TYPE"),"
    echo "  \"mime_type\": $(json_escape "$MIME_TYPE"),"
}

get_file_size() {
    local file_path="$1"
    if [[ "$(uname)" == "Darwin" ]]; then
        stat -f%z "$file_path"
    else
        stat --format=%s "$file_path" 2>/dev/null || wc -c < "$file_path" | tr -d ' '
    fi
}

extract_suspicious_strings() {
    local file_path="$1"
    local indicators=""

    if command_exists strings; then
        local patterns="CreateRemoteThread|VirtualAllocEx|WriteProcessMemory|URLDownloadToFile"
        patterns+="|WScript\.Shell|powershell|cmd\.exe|regsvr32|rundll32|mshta"
        patterns+="|certutil|bitsadmin|schtasks|HKEY_"

        local matches
        matches=$(strings "$file_path" 2>/dev/null | grep -oiE "$patterns" | sort -u | head -20)

        if [[ -n "$matches" ]]; then
            local first=true
            while IFS= read -r match; do
                if [[ "$first" == "true" ]]; then
                    first=false
                else
                    indicators+=", "
                fi
                indicators+="$(json_escape "$match")"
            done <<< "$matches"
        fi
    fi

    echo "[$indicators]"
}

vt_lookup() {
    local sha256="$1"
    local api_key="${VT_API_KEY:-}"

    if [[ -z "$api_key" ]]; then
        echo "\"virustotal\": {\"error\": \"VT_API_KEY environment variable not set\"}"
        return
    fi

    if ! command_exists curl; then
        echo "\"virustotal\": {\"error\": \"curl not available\"}"
        return
    fi

    local response
    response=$(curl -s -w "\n%{http_code}" \
        -H "x-apikey: $api_key" \
        "https://www.virustotal.com/api/v3/files/$sha256" 2>/dev/null)

    local http_code
    http_code=$(echo "$response" | tail -1)
    local body
    body=$(echo "$response" | sed '$d')

    case "$http_code" in
        200)
            if command_exists python3; then
                python3 -c "
import json, sys
data = json.loads(sys.stdin.read())
attrs = data.get('data', {}).get('attributes', {})
stats = attrs.get('last_analysis_stats', {})
total = sum(stats.values())
mal = stats.get('malicious', 0) + stats.get('suspicious', 0)
result = {
    'detected': mal > 0,
    'detections': f'{mal}/{total}',
    'permalink': f'https://www.virustotal.com/gui/file/$sha256'
}
print('\"virustotal\":', json.dumps(result))
" <<< "$body"
            else
                echo "\"virustotal\": {\"raw_status\": 200, \"note\": \"install python3 for parsed results\"}"
            fi
            ;;
        404)
            echo "\"virustotal\": {\"detected\": false, \"message\": \"File not found in VirusTotal\"}"
            ;;
        429)
            echo "\"virustotal\": {\"error\": \"Rate limit exceeded\"}"
            ;;
        *)
            echo "\"virustotal\": {\"error\": \"API returned status $http_code\"}"
            ;;
    esac
}

# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

main() {
    local file_path=""
    local do_vt_lookup="false"

    while [[ $# -gt 0 ]]; do
        case "$1" in
            --vt-lookup) do_vt_lookup="true"; shift ;;
            --help|-h) usage ;;
            *)
                if [[ -z "$file_path" ]]; then
                    file_path="$1"
                else
                    echo "Error: unexpected argument '$1'" >&2
                    usage
                fi
                shift
                ;;
        esac
    done

    if [[ -z "$file_path" ]]; then
        usage
    fi

    if [[ ! -f "$file_path" ]]; then
        echo "Error: File not found: $file_path" >&2
        exit 1
    fi

    # Try Python script first
    if try_python "$file_path" "$do_vt_lookup" 2>/dev/null; then
        exit 0
    fi

    # Fall back to native tools
    echo "{"
    echo "  \"triage_timestamp\": \"$(date -u +%Y-%m-%dT%H:%M:%SZ)\","
    echo "  \"file_path\": $(json_escape "$(cd "$(dirname "$file_path")" && pwd)/$(basename "$file_path")"),"
    echo "  \"file_name\": $(json_escape "$(basename "$file_path")"),"
    echo "  \"file_size\": $(get_file_size "$file_path"),"
    echo -n "  "
    identify_type "$file_path"
    echo "  \"hashes\": {"
    echo "    $(compute_hashes "$file_path")"
    echo "  },"
    echo "  \"quick_indicators\": $(extract_suspicious_strings "$file_path"),"

    # VirusTotal
    if [[ "$do_vt_lookup" == "true" ]]; then
        echo -n "  "
        # Extract SHA256 from already computed hashes
        local sha256
        if command_exists sha256sum; then
            sha256=$(sha256sum "$file_path" | cut -d' ' -f1)
        elif command_exists shasum; then
            sha256=$(shasum -a 256 "$file_path" | cut -d' ' -f1)
        else
            sha256="unavailable"
        fi
        vt_lookup "$sha256"
        echo ","
    fi

    echo "  \"note\": \"Generated with native tools (Python unavailable)\""
    echo "}"
}

main "$@"