#!/usr/bin/env bash

set -euo pipefail

MALICIOUS_PKGS=("axios@1.14.1" "axios@0.30.4" "plain-crypto-js@4.2.1")
C2_DOMAIN_REGEX="sfrclak[.]com"
C2_IP="142.11.206.73"
MAC_ARTIFACT="/Library/Caches/com.apple.act.mond"
WIN_ARTIFACT="${PROGRAMDATA:-}\\wt.exe"
LIN_ARTIFACT="/tmp/ld.py"
ROOT_DIR="$(pwd)"
MAX_DEPTH=10
RUN_NPM_AUDIT=0
AUTO_REMEDIATE=0
NO_COLOR=0
FORCE_COLOR=0
INTERACTIVE=0
FINDINGS=0
SCANNED=0
LOCK_SCANNED=0
declare -A AXIOS_VERSIONS=()
declare -A AFFECTED_DIRS=()

usage() {
	cat <<'EOF'
Usage: axios-malware-solution.sh [options]

Options:
  --npm-audit          Run npm audit in detected folders
  --remediate          Try dynamic remediation in detected folders
  --max-depth N        Limit traversal depth from current directory (default: 10)
  --interactive        Show interactive options menu
  --no-color           Disable colored output
  --force-color        Force colored output
  -h, --help           Show help
EOF
}

if [ "$#" -eq 0 ]; then
	INTERACTIVE=1
fi

while [ "$#" -gt 0 ]; do
	case "$1" in
		--npm-audit)
			RUN_NPM_AUDIT=1
			;;
		--remediate)
			AUTO_REMEDIATE=1
			;;
		--max-depth)
			shift
			MAX_DEPTH="${1:-10}"
			;;
		--interactive)
			INTERACTIVE=1
			;;
		--no-color)
			NO_COLOR=1
			;;
		--force-color)
			FORCE_COLOR=1
			;;
		-h|--help)
			usage
			exit 0
			;;
		*)
			echo "Unknown option: $1"
			usage
			exit 1
			;;
	esac
	shift
done

if ! [[ "$MAX_DEPTH" =~ ^[0-9]+$ ]]; then
	echo "Invalid --max-depth value: $MAX_DEPTH"
	exit 1
fi

show_interactive_menu() {
	echo ""
	echo "Choose run mode:"
	echo "  1) Scan only"
	echo "  2) Scan + npm audit in detected folders"
	echo "  3) Scan + npm audit + remediation"
	echo "  4) Set max depth (current: ${MAX_DEPTH})"
	echo "  0) Exit"
	printf "Enter option: "
	read -r menu_choice
	case "$menu_choice" in
		1)
			RUN_NPM_AUDIT=0
			AUTO_REMEDIATE=0
			;;
		2)
			RUN_NPM_AUDIT=1
			AUTO_REMEDIATE=0
			;;
		3)
			RUN_NPM_AUDIT=1
			AUTO_REMEDIATE=1
			;;
		4)
			printf "Enter max depth (1-50): "
			read -r new_depth
			if [[ "$new_depth" =~ ^[0-9]+$ ]] && [ "$new_depth" -ge 1 ] && [ "$new_depth" -le 50 ]; then
				MAX_DEPTH="$new_depth"
			else
				echo "Invalid depth. Keeping ${MAX_DEPTH}."
			fi
			show_interactive_menu
			;;
		0)
			echo "Exited."
			exit 0
			;;
		*)
			echo "Invalid option."
			show_interactive_menu
			;;
	esac
}

if [ "$INTERACTIVE" -eq 1 ]; then
	show_interactive_menu
fi

if [ "$NO_COLOR" -eq 0 ] && { [ "$FORCE_COLOR" -eq 1 ] || [ -t 1 ]; } && [ "${TERM:-dumb}" != "dumb" ]; then
	RED='\033[0;31m'
	YLW='\033[0;33m'
	GRN='\033[0;32m'
	CYN='\033[0;36m'
	RST='\033[0m'
	BLD='\033[1m'
else
	RED=''
	YLW=''
	GRN=''
	CYN=''
	RST=''
	BLD=''
fi

detect_os() {
	case "$(uname -s)" in
		Darwin) echo "macos" ;;
		Linux) echo "linux" ;;
		MINGW*|MSYS*|CYGWIN*) echo "windows" ;;
		*) echo "unknown" ;;
	esac
}

detect_shell() {
	local shell_name=""
	if [ -n "${SHELL:-}" ]; then
		shell_name="$(basename "$SHELL")"
	fi
	if [ -z "$shell_name" ]; then
		shell_name="$(ps -p $$ -o comm= 2>/dev/null | tr -d ' ' || true)"
	fi
	if [[ "$shell_name" =~ (pwsh|powershell) ]] || [ -n "${PSModulePath:-}" ]; then
		echo "powershell"
		return
	fi
	case "$shell_name" in
		zsh|bash|fish|sh) echo "$shell_name" ;;
		*) echo "unknown" ;;
	esac
}

detect_hosts_file() {
	case "$OS_TYPE" in
		windows)
			if [ -n "${WINDIR:-}" ]; then
				echo "${WINDIR}\\System32\\drivers\\etc\\hosts"
			elif [ -n "${SystemRoot:-}" ]; then
				echo "${SystemRoot}\\System32\\drivers\\etc\\hosts"
			else
				echo "C:\\Windows\\System32\\drivers\\etc\\hosts"
			fi
			;;
		*)
			echo "/etc/hosts"
			;;
	esac
}

check_active_c2_connections() {
	if command -v ss >/dev/null 2>&1; then
		if ss -tnp 2>/dev/null | grep -q "${C2_IP//./\\.}"; then
			bad "Active connection to C2 IP found via ss"
		else
			ok "No active C2 connection via ss"
		fi
		return
	fi
	if command -v netstat >/dev/null 2>&1; then
		case "$OS_TYPE" in
			windows)
			if netstat -ano 2>/dev/null | grep -q "${C2_IP//./\\.}"; then
				bad "Active connection to C2 IP found via netstat"
			else
				ok "No active C2 connection via netstat"
			fi
			;;
			*)
			if netstat -tn 2>/dev/null | grep -q "${C2_IP//./\\.}"; then
				bad "Active connection to C2 IP found via netstat"
			else
				ok "No active C2 connection via netstat"
			fi
			;;
		esac
		return
	fi
	warn "ss/netstat not available"
}

note() { echo -e "  ${CYN}$1${RST}"; }
ok() { echo -e "  ${GRN}[OK]  $1${RST}"; }
warn() { echo -e "  ${YLW}[--]  $1${RST}"; }
bad() {
	echo -e "  ${RED}${BLD}[!!!] $1${RST}"
	FINDINGS=$((FINDINGS + 1))
}

relpath() {
	local p="$1"
	p="${p#$ROOT_DIR/}"
	if [ -z "$p" ] || [ "$p" = "$1" ]; then
		echo "."
	else
		echo "$p"
	fi
}

mark_affected_dir() {
	local d="$1"
	AFFECTED_DIRS["$d"]=1
}

add_axios_version() {
	local d="$1"
	local v="$2"
	local cur="${AXIOS_VERSIONS[$d]-}"
	if [[ " $cur " != *" $v "* ]]; then
		AXIOS_VERSIONS["$d"]="${cur} ${v}"
	fi
}

get_mtime_epoch() {
	local f="$1"
	if stat -c %Y "$f" >/dev/null 2>&1; then
		stat -c %Y "$f"
	elif stat -f %m "$f" >/dev/null 2>&1; then
		stat -f %m "$f"
	else
		echo 0
	fi
}

is_recent_file() {
	local f="$1"
	local now mtime
	now="$(date +%s)"
	mtime="$(get_mtime_epoch "$f")"
	if [ "$mtime" -eq 0 ]; then
		return 2
	fi
	if [ $((now - mtime)) -le 86400 ]; then
		return 0
	fi
	return 1
}

find_package_jsons() {
	find "$ROOT_DIR" -maxdepth "$MAX_DEPTH" \
		\( \
			-path "*/node_modules" -o \
			-path "*/.next" -o \
			-path "*/dist" -o \
			-path "*/build" -o \
			-path "*/out" -o \
			-path "*/coverage" -o \
			-path "*/.turbo" -o \
			-path "*/.vercel" -o \
			-path "*/.git" -o \
			-path "*/.cache" -o \
			-path "*/.pnpm-store" -o \
			-path "*/target" -o \
			-path "*/tmp" \
		\) -prune -o -name "package.json" -print0
}

collect_axios_versions() {
	local lf="$1"
	local d
	d="$(dirname "$lf")"
	case "$lf" in
		*pnpm-lock.yaml)
			while IFS= read -r v; do add_axios_version "$d" "$v"; done < <(
				grep -E "^[[:space:]]+axios@[0-9]+\.[0-9]+\.[0-9]+:" "$lf" | sed -E 's/.*axios@([0-9]+\.[0-9]+\.[0-9]+):.*/\1/'
			)
			;;
		*yarn.lock)
			while IFS= read -r v; do add_axios_version "$d" "$v"; done < <(
				awk '
					$0 ~ /^"?axios@/ { in_block=1 }
					in_block && ($1 == "version" || $1 == "version:") { gsub(/"/, "", $2); print $2; in_block=0 }
					in_block && $0 ~ /^$/ { in_block=0 }
				' "$lf"
			)
			;;
		*package-lock.json|*npm-shrinkwrap.json)
			while IFS= read -r v; do add_axios_version "$d" "$v"; done < <(
				awk '
					$0 ~ "\"node_modules/axios\"" { in_block=1 }
					in_block && $0 ~ "\"version\":" {
						gsub(/[^0-9.]/, "", $0)
						print $0
						in_block=0
					}
					in_block && $0 ~ /^\s*}\s*,?\s*$/ { in_block=0 }
				' "$lf"
			)
			;;
		*bun.lock)
			while IFS= read -r av; do add_axios_version "$d" "${av#axios@}"; done < <(grep -oE 'axios@[0-9]+\.[0-9]+\.[0-9]+' "$lf")
			;;
		*bun.lockb)
			if command -v strings >/dev/null 2>&1; then
				while IFS= read -r av; do add_axios_version "$d" "${av#axios@}"; done < <(strings "$lf" | grep -oE 'axios@[0-9]+\.[0-9]+\.[0-9]+' || true)
			fi
			;;
	esac
}

has_exact_hit_in_lockfile() {
	local lf="$1"
	local n="$2"
	local v="$3"
	case "$lf" in
		*pnpm-lock.yaml)
			grep -qE "^[[:space:]]+${n}@${v}:" "$lf"
			;;
		*yarn.lock)
			awk -v n="$n" -v v="$v" '
				$0 ~ "^\"?" n "@" { in_block=1 }
				in_block && ($1 == "version" || $1 == "version:") { gsub(/"/, "", $2); if ($2 == v) found=1; in_block=0 }
				in_block && $0 ~ /^$/ { in_block=0 }
				END { exit(found ? 0 : 1) }
			' "$lf" >/dev/null 2>&1
			;;
		*package-lock.json|*npm-shrinkwrap.json)
			awk -v n="$n" -v v="$v" '
				$0 ~ "\"node_modules/" n "\"" { in_block=1 }
				in_block && $0 ~ "\"version\": \"" v "\"" { found=1; in_block=0 }
				in_block && $0 ~ /^\s*}\s*,?\s*$/ { in_block=0 }
				END { exit(found ? 0 : 1) }
			' "$lf" >/dev/null 2>&1
			;;
		*bun.lock)
			grep -qE "${n}@${v}" "$lf"
			;;
		*bun.lockb)
			command -v strings >/dev/null 2>&1 && strings "$lf" | grep -qi "${n}@${v}"
			;;
		*)
			return 1
			;;
	esac
}

scan_package_json() {
	local pkgjson="$1"
	local dir
	dir="$(dirname "$pkgjson")"
	for pkg in "${MALICIOUS_PKGS[@]}"; do
		local name="${pkg%@*}"
		local ver="${pkg##*@}"
		if grep -qE "\"${name}\"[[:space:]]*:[[:space:]]*\"[^\"]*${ver}[^\"]*\"" "$pkgjson" 2>/dev/null; then
			bad "Malicious package declaration found: ${pkg} in $(relpath "$pkgjson")"
			mark_affected_dir "$dir"
		fi
	done
}

scan_lockfile() {
	local lf="$1"
	local dir
	local hits=()
	dir="$(dirname "$lf")"
	collect_axios_versions "$lf"
	for pkg in "${MALICIOUS_PKGS[@]}"; do
		local name="${pkg%@*}"
		local ver="${pkg##*@}"
		if has_exact_hit_in_lockfile "$lf" "$name" "$ver"; then
			hits+=("$pkg")
		fi
	done
	if [ "${#hits[@]}" -eq 0 ]; then
		ok "No known bad versions in $(relpath "$lf")"
		return
	fi
	if is_recent_file "$lf"; then
		bad "Recent lockfile hit (last 24h) in $(relpath "$lf"): ${hits[*]}"
		mark_affected_dir "$dir"
	else
		ok "Known hit exists but lockfile is older than 24h in $(relpath "$lf"): ${hits[*]}"
	fi
}

get_versions_for_pkg() {
	local pkg="$1"
	npm view "$pkg" versions --json 2>/dev/null | tr -d '[]"' | tr ',' '\n' | sed 's/^ *//;s/ *$//' | grep -E '^[0-9]+\.[0-9]+\.[0-9]+$' || true
}

pick_safe_version() {
	local pkg="$1"
	local bad_ver="$2"
	local lines same_branch latest
	lines="$(get_versions_for_pkg "$pkg")"
	if [ -z "$lines" ]; then
		echo ""
		return
	fi
	same_branch=""
	if [[ "$bad_ver" =~ ^[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
		local mm="${bad_ver%.*}"
		same_branch="$(echo "$lines" | grep -E "^${mm//./\\.}\.[0-9]+$" | grep -v "^${bad_ver}$" | sort -V | tail -n 1 || true)"
	fi
	if [ -n "$same_branch" ]; then
		echo "$same_branch"
		return
	fi
	latest="$(echo "$lines" | grep -v "^${bad_ver}$" | sort -V | tail -n 1 || true)"
	echo "$latest"
}

remediate_dir() {
	local dir="$1"
	local changed=0
	for pkg in "${MALICIOUS_PKGS[@]}"; do
		local name="${pkg%@*}"
		local bad_ver="${pkg##*@}"
		local present=0
		if [ -f "$dir/package.json" ] && grep -qE "\"${name}\"[[:space:]]*:[[:space:]]*\"[^\"]*${bad_ver}[^\"]*\"" "$dir/package.json"; then
			present=1
		fi
		if [ "$present" -eq 0 ]; then
			continue
		fi
		local safe_ver
		safe_ver="$(pick_safe_version "$name" "$bad_ver")"
		if [ -n "$safe_ver" ]; then
			note "Remediating $(relpath "$dir"): npm install ${name}@${safe_ver} --save-exact"
			(cd "$dir" && npm install "${name}@${safe_ver}" --save-exact >/dev/null 2>&1 || true)
			changed=1
		else
			if [ "$name" = "plain-crypto-js" ]; then
				note "Remediating $(relpath "$dir"): npm uninstall plain-crypto-js"
				(cd "$dir" && npm uninstall plain-crypto-js >/dev/null 2>&1 || true)
				changed=1
			fi
		fi
	done
	if [ "$changed" -eq 0 ]; then
		warn "No remediation applied in $(relpath "$dir")"
	fi
}

run_audit_for_dir() {
	local dir="$1"
	if ! command -v npm >/dev/null 2>&1; then
		warn "npm is not available; audit skipped"
		return
	fi
	note "Running npm audit in $(relpath "$dir")"
	(cd "$dir" && npm audit --audit-level=high) || warn "npm audit reported issues in $(relpath "$dir")"
}

OS_TYPE="$(detect_os)"
SHELL_TYPE="$(detect_shell)"
HOSTS_FILE="$(detect_hosts_file)"

echo ""
echo -e "${BLD}${CYN}==============================================${RST}"
echo -e "${BLD}${CYN}  Supply Chain Threat Scanner v1.3${RST}"
echo -e "${BLD}${CYN}==============================================${RST}"
echo -e "  Root: $(relpath "$ROOT_DIR")"
echo -e "  Max depth: ${MAX_DEPTH}"
echo -e "  OS: ${OS_TYPE}"
echo -e "  Shell: ${SHELL_TYPE}"
echo ""

echo -e "${BLD}[1/4] Scanning package.json files...${RST}"
echo ""
while IFS= read -r -d '' pkgjson; do
	SCANNED=$((SCANNED + 1))
	note "Checking: $(relpath "$pkgjson")"
	scan_package_json "$pkgjson"
done < <(find_package_jsons)
echo ""
echo -e "  Scanned ${BLD}${SCANNED}${RST} package.json file(s)"

echo ""
echo -e "${BLD}[2/4] Checking for malware artifacts...${RST}"
echo ""
case "$OS_TYPE" in
	macos)
		[ -f "$MAC_ARTIFACT" ] && bad "macOS artifact found: $MAC_ARTIFACT" || ok "macOS artifact not present"
		;;
	linux)
		[ -f "$LIN_ARTIFACT" ] && bad "Linux artifact found: $LIN_ARTIFACT" || ok "Linux artifact not present"
		;;
	windows)
		if [ -n "$WIN_ARTIFACT" ] && [ -f "$WIN_ARTIFACT" ]; then
			bad "Windows artifact found: $WIN_ARTIFACT"
		elif [ -f "/c/ProgramData/wt.exe" ]; then
			bad "Windows artifact found: /c/ProgramData/wt.exe"
		else
			ok "Windows artifact not present"
		fi
		;;
	*)
		warn "Unknown OS; artifact check limited"
		;;
esac

echo ""
echo -e "${BLD}[3/4] Checking network indicators...${RST}"
echo ""
if grep -qE "${C2_DOMAIN_REGEX}|${C2_IP//./\\.}" "$HOSTS_FILE" 2>/dev/null; then
	bad "C2 indicator found in hosts file: $HOSTS_FILE"
else
ok "Hosts file clean: $HOSTS_FILE"
fi
check_active_c2_connections

echo ""
echo -e "${BLD}[4/4] Scanning lockfiles...${RST}"
echo ""
while IFS= read -r -d '' pkgjson; do
	pkgdir="$(dirname "$pkgjson")"
	for lf in \
		"$pkgdir/package-lock.json" \
		"$pkgdir/npm-shrinkwrap.json" \
		"$pkgdir/yarn.lock" \
		"$pkgdir/pnpm-lock.yaml" \
		"$pkgdir/bun.lock" \
		"$pkgdir/bun.lockb"; do
		if [ -f "$lf" ]; then
			LOCK_SCANNED=$((LOCK_SCANNED + 1))
			note "Scanning: $(relpath "$lf")"
			scan_lockfile "$lf"
		fi
	done
done < <(find_package_jsons)

if [ "$LOCK_SCANNED" -eq 0 ]; then
	warn "No lockfiles found"
fi

echo ""
echo -e "${BLD}Axios versions by project:${RST}"
if [ "${#AXIOS_VERSIONS[@]}" -eq 0 ]; then
	warn "No axios versions found in lockfiles"
else
	while IFS= read -r d; do
		versions="${AXIOS_VERSIONS[$d]}"
		versions="$(echo "$versions" | tr ' ' '\n' | sed '/^$/d' | sort -u | tr '\n' ' ')"
		echo -e "  ${CYN}$(relpath "$d"):${RST} ${versions}"
	done < <(printf '%s\n' "${!AXIOS_VERSIONS[@]}" | sort)
fi

if [ "$RUN_NPM_AUDIT" -eq 1 ]; then
	echo ""
	echo -e "${BLD}npm audit in detected folders:${RST}"
	if [ "${#AFFECTED_DIRS[@]}" -eq 0 ]; then
		warn "No detected folders for npm audit"
	else
		while IFS= read -r d; do
			run_audit_for_dir "$d"
		done < <(printf '%s\n' "${!AFFECTED_DIRS[@]}" | sort)
	fi
fi

if [ "$AUTO_REMEDIATE" -eq 1 ]; then
	echo ""
	echo -e "${BLD}Dynamic remediation:${RST}"
	if [ "${#AFFECTED_DIRS[@]}" -eq 0 ]; then
		warn "No detected folders for remediation"
	else
		if ! command -v npm >/dev/null 2>&1; then
			warn "npm is not available; remediation skipped"
		else
			while IFS= read -r d; do
				remediate_dir "$d"
			done < <(printf '%s\n' "${!AFFECTED_DIRS[@]}" | sort)
		fi
	fi
fi

echo ""
echo -e "${BLD}${CYN}==============================================${RST}"
if [ "$FINDINGS" -gt 0 ]; then
	echo -e "${BLD}${RED}RESULT: ${FINDINGS} active indicator(s) found${RST}"
else
	echo -e "${BLD}${GRN}RESULT: No active indicators found${RST}"
fi
echo -e "${BLD}${CYN}==============================================${RST}"

echo ""
echo -e "Run options:"
echo -e "  ./axios-malware-solution.sh --npm-audit"
echo -e "  ./axios-malware-solution.sh --npm-audit --remediate"
echo -e "  ./axios-malware-solution.sh --max-depth 10"
echo -e "  ./axios-malware-solution.sh --force-color"