#!/usr/bin/env bash
set -euo pipefail
IFS=$'\n\t'

# ubuntu-security-audit.sh
# Purpose: Audit a local Ubuntu host for pending security updates.
# Default behavior: read-only, no package changes.
#
# Notes:
# - This script reports on locally available package metadata.
# - It does not install, remove, reboot, or modify packages by default.
# - Security filtering relies on repository metadata and naming conventions.
#   Validate results in your environment before using them as compliance evidence.

SCRIPT_NAME="$(basename "$0")"
OUTPUT_FORMAT="text"
REFRESH_CACHE="false"
ONLY_SECURITY="true"
JSON_PRETTY="false"

usage() {
  cat <<'EOF'
Usage:
  ubuntu-security-audit.sh [--format text|json] [--refresh-cache] [--all-updates] [--pretty-json]

Options:
  --format text|json     Output format. Default: text
  --refresh-cache        Run apt-get update before auditing (requires appropriate permissions)
  --all-updates          Report all pending updates, not only security updates
  --pretty-json          Pretty-print JSON output when --format json is used
  -h, --help             Show this help message

Examples:
  ./ubuntu-security-audit.sh
  ./ubuntu-security-audit.sh --refresh-cache
  ./ubuntu-security-audit.sh --format json --pretty-json
EOF
}

log_err() {
  printf '%s: ERROR: %s\n' "$SCRIPT_NAME" "$*" >&2
}

log_info() {
  printf '%s: %s\n' "$SCRIPT_NAME" "$*"
}

require_cmd() {
  local cmd="$1"
  command -v "$cmd" >/dev/null 2>&1 || {
    log_err "Required command not found: $cmd"
    exit 2
  }
}

parse_args() {
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --format)
        [[ $# -ge 2 ]] || { log_err "--format requires a value"; exit 2; }
        OUTPUT_FORMAT="$2"
        shift 2
        ;;
      --refresh-cache)
        REFRESH_CACHE="true"
        shift
        ;;
      --all-updates)
        ONLY_SECURITY="false"
        shift
        ;;
      --pretty-json)
        JSON_PRETTY="true"
        shift
        ;;
      -h|--help)
        usage
        exit 0
        ;;
      *)
        log_err "Unknown argument: $1"
        usage >&2
        exit 2
        ;;
    esac
  done

  case "$OUTPUT_FORMAT" in
    text|json) ;;
    *) log_err "Invalid --format value: $OUTPUT_FORMAT"; exit 2 ;;
  esac
}

refresh_cache() {
  if [[ "$REFRESH_CACHE" == "true" ]]; then
    log_info "Refreshing package lists"
    sudo -n apt-get update
  fi
}

audit_updates() {
  # apt list --upgradable outputs a warning line and package entries.
  # We filter out non-package lines and optionally narrow to security sources.
  apt list --upgradable 2>/dev/null | awk 'NR>1 {print}'
}

filter_security_updates() {
  # This approach relies on repository naming conventions and metadata.
  # Validate in your environment before using it as compliance evidence.
  awk '
    {
      line = $0
      if (line ~ /security/ || line ~ /Ubuntu-Security/ || line ~ /security.ubuntu.com/) {
        print line
      }
    }
  '
}

collect_results() {
  local all_updates security_updates status count

  all_updates="$(audit_updates || true)"

  if [[ -n "$all_updates" ]]; then
    count=$(printf '%s\n' "$all_updates" | sed '/^$/d' | wc -l | tr -d ' ')
  else
    count=0
  fi

  if [[ "$ONLY_SECURITY" == "true" ]]; then
    security_updates="$(printf '%s\n' "$all_updates" | filter_security_updates || true)"
  else
    security_updates="$all_updates"
  fi

  if [[ -n "$security_updates" ]]; then
    status="updates_available"
  else
    status="clean"
  fi

  if [[ "$OUTPUT_FORMAT" == "json" ]]; then
    if [[ "$JSON_PRETTY" == "true" ]]; then
      python3 - "$status" "$count" <<'PY'
import json, sys
status = sys.argv[1]
count = int(sys.argv[2])
print(json.dumps({
    "status": status,
    "pending_update_count": count,
}, indent=2, sort_keys=True))
PY
    else
      python3 - "$status" "$count" <<'PY'
import json, sys
status = sys.argv[1]
count = int(sys.argv[2])
print(json.dumps({
    "status": status,
    "pending_update_count": count,
}, separators=(',', ':'), sort_keys=True))
PY
    fi
  else
    printf 'Status: %s\n' "$status"
    printf 'Pending updates: %s\n' "$count"
    if [[ -n "$security_updates" ]]; then
      printf '\nAffected packages:\n%s\n' "$security_updates"
    fi
  fi

  # Exit codes:
  # 0 = no pending security updates found
  # 1 = pending security updates found
  # 2 = script usage or runtime error
  if [[ "$status" == "updates_available" ]]; then
    return 1
  fi
  return 0
}

main() {
  parse_args "$@"
  require_cmd apt
  require_cmd apt-get
  require_cmd awk
  require_cmd sed
  require_cmd python3

  refresh_cache
  collect_results
}

main "$@"