#!/usr/bin/env bash
# base64-scan.sh — Detect base64-obfuscated prompt injection in source files
#
# Extracts base64 blobs >= 40 chars, decodes them, and checks decoded content
# against the same injection patterns used by prompt-injection-scan.sh.
#
# Usage:
#   scripts/base64-scan.sh --diff origin/main   # CI mode: scan changed files
#   scripts/base64-scan.sh --file path/to/file   # Scan a single file
#   scripts/base64-scan.sh --dir agents/          # Scan all files in a directory
#
# Exit codes:
#   0 = clean
#   1 = findings detected
#   2 = usage error
set -euo pipefail

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

# ─── Injection Patterns (decoded content) ────────────────────────────────────
# Subset of patterns — if someone base64-encoded something, check for the
# most common injection indicators.
DECODED_PATTERNS=(
  'ignore[[:space:]]+(all[[:space:]]+)?previous[[:space:]]+instructions'
  'you[[:space:]]+are[[:space:]]+now[[:space:]]+'
  'system[[:space:]]+prompt'
  '</?system>'
  '</?assistant>'
  '\[SYSTEM\]'
  '\[INST\]'
  '<<SYS>>'
  'override[[:space:]]+(system|safety|security)'
  'pretend[[:space:]]+(you|to)[[:space:]]'
  'act[[:space:]]+as[[:space:]]+(a|an|if)'
  'jailbreak'
  'bypass[[:space:]]+(safety|content|security)'
  'eval[[:space:]]*\('
  'exec[[:space:]]*\('
  'rm[[:space:]]+-rf'
  'curl[[:space:]].*\|[[:space:]]*sh'
  'wget[[:space:]].*\|[[:space:]]*sh'
)

# ─── Ignorelist ──────────────────────────────────────────────────────────────

IGNOREFILE=".base64scanignore"
IGNORED_PATTERNS=()

load_ignorelist() {
  if [[ -f "$IGNOREFILE" ]]; then
    while IFS= read -r line; do
      # Skip comments and empty lines
      [[ "$line" =~ ^[[:space:]]*# ]] && continue
      [[ -z "${line// }" ]] && continue
      IGNORED_PATTERNS+=("$line")
    done < "$IGNOREFILE"
  fi
}

is_ignored() {
  local blob="$1"
  if [[ ${#IGNORED_PATTERNS[@]} -eq 0 ]]; then
    return 1
  fi
  for pattern in "${IGNORED_PATTERNS[@]}"; do
    if [[ "$blob" == "$pattern" ]]; then
      return 0
    fi
  done
  return 1
}

# ─── Skip Rules ──────────────────────────────────────────────────────────────

should_skip_file() {
  local file="$1"
  # Skip binary files
  case "$file" in
    *.png|*.jpg|*.jpeg|*.gif|*.ico|*.woff|*.woff2|*.ttf|*.eot|*.otf) return 0 ;;
    *.zip|*.tar|*.gz|*.bz2|*.xz|*.7z) return 0 ;;
    *.pdf|*.doc|*.docx|*.xls|*.xlsx) return 0 ;;
  esac
  # Skip lockfiles and node_modules
  # NOTE: patterns use both `*/X` and `X` to match at any depth INCLUDING repo-root
  # (bash case `*/foo` requires a prefix and won't match bare `foo`).
  case "$file" in
    */node_modules/*|node_modules/*) return 0 ;;
    */package-lock.json|package-lock.json) return 0 ;;
    */yarn.lock|yarn.lock) return 0 ;;
    */pnpm-lock.yaml|pnpm-lock.yaml) return 0 ;;
  esac
  # Skip the scan scripts themselves and test files
  case "$file" in
    */base64-scan.sh) return 0 ;;
    */security-scan.test.cjs) return 0 ;;
  esac
  return 1
}

is_data_uri() {
  local context="$1"
  # data:image/png;base64,... or data:application/font-woff;base64,...
  echo "$context" | grep -qE 'data:[a-zA-Z]+/[a-zA-Z0-9.+-]+;base64,' 2>/dev/null
}

# ─── File Collection ─────────────────────────────────────────────────────────

collect_files() {
  local mode="$1"
  shift

  case "$mode" in
    --diff)
      local base="${1:-origin/main}"
      git diff --name-only --diff-filter=ACMR "$base"...HEAD 2>/dev/null \
        | grep -vE '\.(png|jpg|jpeg|gif|ico|woff|woff2|ttf|eot|otf|zip|tar|gz|pdf)$' || true
      ;;
    --file)
      if [[ -f "$1" ]]; then
        echo "$1"
      else
        echo "Error: file not found: $1" >&2
        exit 2
      fi
      ;;
    --dir)
      local dir="$1"
      if [[ ! -d "$dir" ]]; then
        echo "Error: directory not found: $dir" >&2
        exit 2
      fi
      find "$dir" -type f ! -path '*/node_modules/*' ! -path '*/.git/*' ! -path '*/dist/*' \
        ! -name '*.png' ! -name '*.jpg' ! -name '*.gif' ! -name '*.woff*' 2>/dev/null || true
      ;;
    --stdin)
      cat
      ;;
    *)
      echo "Usage: $0 --diff [base] | --file <path> | --dir <path> | --stdin" >&2
      exit 2
      ;;
  esac
}

# ─── Scanner ─────────────────────────────────────────────────────────────────

extract_and_check_blobs() {
  local file="$1"
  local found=0
  local matches

  # Extract candidate blobs with one grep per file. The previous line-by-line
  # implementation forked grep once per line and made large generated files take
  # tens of seconds even when they contained no candidates.
  matches=$(grep -nEo '[A-Za-z0-9+/]{'"$MIN_BLOB_LENGTH"',}={0,3}' "$file" 2>/dev/null || true)
  if [[ -z "$matches" ]]; then
    return 0
  fi

  while IFS= read -r match; do
    [[ -z "$match" ]] && continue

    local line_num="${match%%:*}"
    local blob="${match#*:}"
    [[ -z "$line_num" || -z "$blob" ]] && continue

    # Cheap validity filter: base64 length cannot be 1 modulo 4.
    if [[ $(( ${#blob} % 4 )) -eq 1 ]]; then
      continue
    fi

    # Skip data URIs — legitimate base64 usage.
    local line
    line=$(awk -v n="$line_num" 'NR == n { print; exit }' "$file" 2>/dev/null || true)
    if is_data_uri "$line"; then
      continue
    fi

    # Check ignorelist
    if [[ ${#IGNORED_PATTERNS[@]} -gt 0 ]] && is_ignored "$blob"; then
      continue
    fi

    # Try to decode — if it fails, not valid base64
    local decoded
    decoded=$(printf '%s' "$blob" | base64 -d 2>/dev/null | tr -d '\000' || echo "")

    if [[ -z "$decoded" ]]; then
      continue
    fi

    # Check if decoded content is mostly printable text (not random binary)
    local total_chars=${#decoded}
    if [[ $total_chars -eq 0 ]]; then
      continue
    fi

    # Count printable ASCII characters
    local printable_count
    printable_count=$(printf '%s' "$decoded" | tr -cd '[:print:]' | wc -c | tr -d ' ')
    # Skip if less than 70% printable (likely binary data, not obfuscated text)
    if [[ $((printable_count * 100 / total_chars)) -lt 70 ]]; then
      continue
    fi

    # Scan decoded content against injection patterns
    for pattern in "${DECODED_PATTERNS[@]}"; do
      if printf '%s' "$decoded" | grep -iqE "$pattern" 2>/dev/null; then
        if [[ $found -eq 0 ]]; then
          echo "FAIL: $file"
          found=1
        fi
        echo "  line $line_num: base64 blob decodes to suspicious content"
        echo "    blob: ${blob:0:60}..."
        echo "    decoded: ${decoded:0:120}"
        echo "    matched: $pattern"
        break
      fi
    done
  done <<< "$matches"

  return $found
}

# ─── Main ────────────────────────────────────────────────────────────────────

main() {
  if [[ $# -eq 0 ]]; then
    echo "Usage: $0 --diff [base] | --file <path> | --dir <path>" >&2
    exit 2
  fi

  load_ignorelist

  local mode="$1"
  shift

  local files
  files=$(collect_files "$mode" "$@")

  if [[ -z "$files" ]]; then
    echo "base64-scan: no files to scan"
    exit 0
  fi

  local total=0
  local failed=0

  while IFS= read -r file; do
    [[ -z "$file" ]] && continue
    if should_skip_file "$file"; then
      continue
    fi
    total=$((total + 1))
    if ! extract_and_check_blobs "$file"; then
      failed=$((failed + 1))
    fi
  done <<< "$files"

  echo ""
  echo "base64-scan: scanned $total files, $failed with findings"

  if [[ $failed -gt 0 ]]; then
    exit 1
  fi
  exit 0
}

main "$@"