#!/bin/sh
# CVE Triage - Analyzer Module
# Uses LocalAI to analyze CVE impact and generate recommendations

# Check LocalAI availability
check_localai() {
	local url="${localai_url:-http://127.0.0.1:8081}"
	wget -q -O /dev/null --timeout=3 "${url}/v1/models" 2>/dev/null
}

# Call LocalAI chat completion
localai_complete() {
	local prompt="$1"
	local max_tokens="${2:-1024}"
	local url="${localai_url:-http://127.0.0.1:8081}"
	local model="${localai_model:-tinyllama-1.1b-chat-v1.0.Q4_K_M}"

	# Escape prompt for JSON
	local escaped_prompt=$(echo "$prompt" | sed 's/\\/\\\\/g; s/"/\\"/g' | tr '\n' ' ')

	local request=$(cat <<EOF
{
  "model": "$model",
  "messages": [
    {"role": "system", "content": "You are a cybersecurity CVE analyst for SecuBox, an OpenWrt-based security appliance running on ARM64. Provide concise, actionable security analysis."},
    {"role": "user", "content": "$escaped_prompt"}
  ],
  "max_tokens": $max_tokens,
  "temperature": 0.3
}
EOF
)

	local response=$(echo "$request" | wget -q -O - --post-data=- \
		--header="Content-Type: application/json" \
		"${url}/v1/chat/completions" 2>/dev/null)

	echo "$response" | jsonfilter -e '@.choices[0].message.content' 2>/dev/null
}

# Analyze a single CVE with LocalAI
analyze_cve() {
	local cve_id="$1"
	local description="$2"
	local cvss="$3"
	local affected_pkg="$4"

	if ! check_localai; then
		log_warn "LocalAI not available, using basic analysis"
		basic_analyze_cve "$cve_id" "$cvss" "$affected_pkg"
		return
	fi

	local prompt="Analyze this CVE for an OpenWrt ARM64 security appliance:

CVE: $cve_id
CVSS Score: $cvss
Affected Package: ${affected_pkg:-unknown}
Description: $description

Provide a JSON response with these fields:
1. impact_level: critical/high/medium/low
2. exploitability: easy/moderate/difficult/theoretical
3. affects_openwrt: true/false (is this relevant to OpenWrt/embedded Linux?)
4. action: patch/mitigate/monitor/ignore
5. mitigation: brief mitigation steps if patch unavailable
6. urgency: immediate/soon/scheduled/none

Respond ONLY with valid JSON, no explanation."

	local analysis=$(localai_complete "$prompt" 512)

	# Validate JSON response
	if echo "$analysis" | jsonfilter -e '@.impact_level' >/dev/null 2>&1; then
		echo "$analysis"
	else
		# Fallback to basic analysis
		basic_analyze_cve "$cve_id" "$cvss" "$affected_pkg"
	fi
}

# Basic CVE analysis without AI
basic_analyze_cve() {
	local cve_id="$1"
	local cvss="$2"
	local affected_pkg="$3"

	local impact="low"
	local action="monitor"
	local urgency="scheduled"

	# Determine impact based on CVSS
	if [ -n "$cvss" ]; then
		local cvss_int=$(echo "$cvss" | cut -d. -f1)
		if [ "$cvss_int" -ge 9 ]; then
			impact="critical"
			action="patch"
			urgency="immediate"
		elif [ "$cvss_int" -ge 7 ]; then
			impact="high"
			action="patch"
			urgency="soon"
		elif [ "$cvss_int" -ge 4 ]; then
			impact="medium"
			action="mitigate"
			urgency="scheduled"
		fi
	fi

	cat <<EOF
{
  "impact_level": "$impact",
  "exploitability": "unknown",
  "affects_openwrt": true,
  "action": "$action",
  "mitigation": "Update ${affected_pkg:-affected package} to latest version",
  "urgency": "$urgency"
}
EOF
}

# Analyze multiple CVEs and generate recommendations
analyze_cves_batch() {
	local cves="$1"
	local min_severity="${min_severity:-high}"

	local results='[]'

	echo "$cves" | jsonfilter -e '@[*]' 2>/dev/null | while read -r cve_data; do
		local cve_id=$(echo "$cve_data" | jsonfilter -e '@.cve' 2>/dev/null)
		local description=$(echo "$cve_data" | jsonfilter -e '@.description' 2>/dev/null)
		local cvss=$(echo "$cve_data" | jsonfilter -e '@.cvss' 2>/dev/null)
		local severity=$(echo "$cve_data" | jsonfilter -e '@.severity' 2>/dev/null)
		local affected_pkg=$(echo "$cve_data" | jsonfilter -e '@.affected_package' 2>/dev/null)

		# Filter by minimum severity
		case "$min_severity" in
			critical)
				[ "$severity" != "critical" ] && continue
				;;
			high)
				[ "$severity" != "critical" ] && [ "$severity" != "high" ] && continue
				;;
			medium)
				[ "$severity" = "low" ] && continue
				;;
		esac

		log_info "Analyzing $cve_id (CVSS: $cvss)..."

		local analysis=$(analyze_cve "$cve_id" "$description" "$cvss" "$affected_pkg")

		# Merge CVE data with analysis
		printf '{"cve":"%s","cvss":%s,"severity":"%s","affected_package":"%s","analysis":%s}\n' \
			"$cve_id" "${cvss:-0}" "$severity" "$affected_pkg" "$analysis"
	done | json_slurp
}

# Generate security posture summary
generate_summary() {
	local analyzed_cves="$1"

	if ! check_localai; then
		echo '{"summary":"LocalAI unavailable. Manual review required.","risk_score":0}'
		return
	fi

	local cve_count=$(echo "$analyzed_cves" | jsonfilter -e '@[*]' 2>/dev/null | wc -l)
	local critical_count=$(echo "$analyzed_cves" | jsonfilter -e '@[*].severity' 2>/dev/null | grep -c "critical")
	local high_count=$(echo "$analyzed_cves" | jsonfilter -e '@[*].severity' 2>/dev/null | grep -c "high")

	local prompt="Summarize the security posture based on these CVE findings:

Total CVEs analyzed: $cve_count
Critical: $critical_count
High: $high_count

CVE Details:
$(echo "$analyzed_cves" | head -c 2000)

Provide a brief (2-3 sentences) executive summary and a risk score from 0-100.
Format as JSON: {\"summary\":\"...\",\"risk_score\":N}"

	local summary=$(localai_complete "$prompt" 256)

	# Validate JSON
	if echo "$summary" | jsonfilter -e '@.summary' >/dev/null 2>&1; then
		echo "$summary"
	else
		cat <<EOF
{
  "summary": "Found $cve_count CVEs ($critical_count critical, $high_count high). Immediate attention required for critical vulnerabilities.",
  "risk_score": $(( (critical_count * 25) + (high_count * 10) ))
}
EOF
	fi
}