#!/usr/bin/env bash

#########################################################################################
#### incident-response-kit.sh — Live incident volatile data capture and IOC search   ####
#### Collect processes, network, users, logs into tamper-evident archive             ####
#### Requires: bash 4+, root recommended                                             ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.00                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   sudo ./incident-response-kit.sh --collect                                     ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -uo pipefail
# NOTE: no -e — collection must continue if individual commands fail

#------------------------------------------------------------------------------
# DEFAULTS
#------------------------------------------------------------------------------
OUTPUT_DIR="/var/log/incident-response"
RUN_MODE=""
IOC_PATTERN=""
CASE_ID="$(date +%Y%m%d-%H%M%S)"
VERBOSE="${VERBOSE:-0}"
COLOR="${COLOR:-auto}"
HASH_CMD="sha256sum"

#------------------------------------------------------------------------------
# STATE
#------------------------------------------------------------------------------
SCRIPT_NAME="$(basename "$0")"
readonly SCRIPT_NAME
START_TIME=""
COLLECT_COUNT=0
ERROR_COUNT=0

#------------------------------------------------------------------------------
# COLORS (pre-initialized before setup_colors)
#------------------------------------------------------------------------------
RED=""
GREEN=""
YELLOW=""
BLUE=""
CYAN=""
BOLD=""
DIM=""
RESET=""

setup_colors() {
    if [[ "${COLOR}" == "never" ]]; then
        return
    fi
    if [[ "${COLOR}" == "always" ]] || [[ -t 1 ]]; then
        RED='\033[0;31m'
        GREEN='\033[0;32m'
        YELLOW='\033[1;33m'
        BLUE='\033[0;34m'
        CYAN='\033[0;36m'
        BOLD='\033[1m'
        DIM='\033[2m'
        RESET='\033[0m'
    fi
}

#------------------------------------------------------------------------------
# LOGGING HELPERS
#------------------------------------------------------------------------------
log()   { echo -e "${GREEN}[ir-kit]${RESET} $1"; }
warn()  { echo -e "${YELLOW}[ir-kit]${RESET} $1"; }
err()   { echo -e "${RED}[ir-kit]${RESET} $1" >&2; }
verbose() {
    if [[ "${VERBOSE}" == "1" ]]; then
        echo -e "${DIM}[ir-kit]${RESET} $1"
    fi
}
die() { err "$1"; exit 1; }

section_header() {
    echo -e "\n${CYAN}${BOLD}━━━ $1 ━━━${RESET}"
}

subsection() {
    echo -e "  ${BLUE}▸${RESET} $1"
}

field() {
    printf "  ${BOLD}%-18s${RESET} %s\n" "$1" "$2"
}

field_color() {
    printf "  ${BOLD}%-18s${RESET} %b\n" "$1" "$2"
}

elapsed() {
    local end
    end="$(date +%s)"
    local diff=$(( end - START_TIME ))
    echo "${diff}s"
}

#------------------------------------------------------------------------------
# COLLECT HELPER
#------------------------------------------------------------------------------
collect_cmd() {
    local label="$1" outfile="$2"
    shift 2
    verbose "Collecting: ${label}"
    if "$@" > "$outfile" 2>/dev/null; then
        ((COLLECT_COUNT++))
        echo -e "  ${GREEN}✓${RESET} ${label}"
    else
        ((ERROR_COUNT++))
        echo -e "  ${YELLOW}⊘${RESET} ${label} ${DIM}(skipped)${RESET}"
    fi
}

#------------------------------------------------------------------------------
# MODE: FULL COLLECTION
#------------------------------------------------------------------------------
do_collect() {
    local CASE_DIR="${OUTPUT_DIR}/ir-${CASE_ID}"
    mkdir -p "${CASE_DIR}"

    section_header "Incident Response — Full Collection"
    field "Case ID:" "${CASE_ID}"
    field "Output:" "${CASE_DIR}"
    field "Started:" "$(date)"
    echo ""

    # --- System Info ---
    section_header "System Information"
    collect_cmd "Hostname"          "${CASE_DIR}/hostname.txt"       hostname
    collect_cmd "Date / Time"       "${CASE_DIR}/date.txt"           date --iso-8601=seconds
    collect_cmd "Uptime"            "${CASE_DIR}/uptime.txt"         uptime
    collect_cmd "Kernel (uname)"    "${CASE_DIR}/uname.txt"          uname -a
    collect_cmd "OS Release"        "${CASE_DIR}/os-release.txt"     cat /etc/os-release

    # --- Volatile Data ---
    section_header "Volatile Data"
    collect_cmd "Processes (full)"      "${CASE_DIR}/ps-auxwww.txt"     ps auxwww
    collect_cmd "Network listeners"     "${CASE_DIR}/ss-tulnp.txt"      ss -tulnp
    collect_cmd "Network connections"   "${CASE_DIR}/ss-anp.txt"        ss -anp
    collect_cmd "Routing table (v4)"    "${CASE_DIR}/ip-route.txt"      ip route
    collect_cmd "Routing table (v6)"    "${CASE_DIR}/ip6-route.txt"     ip -6 route
    collect_cmd "ARP / Neighbors"       "${CASE_DIR}/ip-neigh.txt"      ip neigh
    collect_cmd "DNS resolv.conf"       "${CASE_DIR}/resolv-conf.txt"   cat /etc/resolv.conf
    collect_cmd "Open files (summary)"  "${CASE_DIR}/lsof-summary.txt"  lsof -n -P +c 15
    collect_cmd "Loaded kernel modules" "${CASE_DIR}/lsmod.txt"         lsmod
    collect_cmd "Mount points"          "${CASE_DIR}/mount.txt"         mount

    # Environment variables of running processes
    verbose "Collecting: Process environments"
    if ls /proc/[0-9]*/environ > /dev/null 2>&1; then
        (
            for f in /proc/[0-9]*/environ; do
                local_pid="${f#/proc/}"
                local_pid="${local_pid%/environ}"
                printf "=== PID %s ===\n" "${local_pid}"
                tr '\0' '\n' < "${f}" 2>/dev/null || true
                echo ""
            done
        ) > "${CASE_DIR}/proc-environ.txt" 2>/dev/null
        ((COLLECT_COUNT++))
        echo -e "  ${GREEN}✓${RESET} Process environments"
    else
        ((ERROR_COUNT++))
        echo -e "  ${YELLOW}⊘${RESET} Process environments ${DIM}(skipped)${RESET}"
    fi

    # --- User Data ---
    section_header "User Activity"
    collect_cmd "Logged-in users (who)"  "${CASE_DIR}/who.txt"        who
    collect_cmd "User activity (w)"      "${CASE_DIR}/w.txt"          w
    collect_cmd "Login history (last)"   "${CASE_DIR}/last.txt"       last -25
    collect_cmd "Last login (lastlog)"   "${CASE_DIR}/lastlog.txt"    lastlog
    # Users with shells
    verbose "Collecting: Users with login shells"
    if grep -vE '(nologin|false|sync|halt|shutdown)$' /etc/passwd > "${CASE_DIR}/users-with-shells.txt" 2>/dev/null; then
        ((COLLECT_COUNT++))
        echo -e "  ${GREEN}✓${RESET} Users with login shells"
    else
        ((ERROR_COUNT++))
        echo -e "  ${YELLOW}⊘${RESET} Users with login shells ${DIM}(skipped)${RESET}"
    fi
    # Sudo log
    verbose "Collecting: Sudo log"
    if grep -i sudo /var/log/auth.log > "${CASE_DIR}/sudo-log.txt" 2>/dev/null || \
       grep -i sudo /var/log/secure > "${CASE_DIR}/sudo-log.txt" 2>/dev/null; then
        ((COLLECT_COUNT++))
        echo -e "  ${GREEN}✓${RESET} Sudo log"
    else
        ((ERROR_COUNT++))
        echo -e "  ${YELLOW}⊘${RESET} Sudo log ${DIM}(skipped)${RESET}"
    fi

    # --- Scheduled Tasks ---
    section_header "Scheduled Tasks"
    collect_cmd "System crontab"         "${CASE_DIR}/crontab-system.txt"  cat /etc/crontab
    # User crontabs
    verbose "Collecting: User crontabs"
    (
        while IFS=: read -r user _; do
            local_cron="$(crontab -l -u "${user}" 2>/dev/null)" || true
            if [[ -n "${local_cron}" ]]; then
                printf "=== %s ===\n%s\n\n" "${user}" "${local_cron}"
            fi
        done < /etc/passwd
    ) > "${CASE_DIR}/crontab-users.txt" 2>/dev/null
    ((COLLECT_COUNT++))
    echo -e "  ${GREEN}✓${RESET} User crontabs"

    collect_cmd "Systemd timers"         "${CASE_DIR}/systemd-timers.txt"  systemctl list-timers --all --no-pager

    # --- Containers ---
    section_header "Container Info"
    if command -v docker &>/dev/null; then
        collect_cmd "Docker containers"  "${CASE_DIR}/docker-ps.txt"  docker ps -a --no-trunc
    else
        echo -e "  ${DIM}⊘ docker not installed${RESET}"
    fi
    if command -v podman &>/dev/null; then
        collect_cmd "Podman containers"  "${CASE_DIR}/podman-ps.txt"  podman ps -a --no-trunc
    else
        echo -e "  ${DIM}⊘ podman not installed${RESET}"
    fi

    # --- Recent Logs ---
    section_header "Recent Logs"
    # auth.log or secure
    verbose "Collecting: Auth log"
    if tail -500 /var/log/auth.log > "${CASE_DIR}/auth-log.txt" 2>/dev/null || \
       tail -500 /var/log/secure > "${CASE_DIR}/auth-log.txt" 2>/dev/null; then
        ((COLLECT_COUNT++))
        echo -e "  ${GREEN}✓${RESET} Auth log (last 500 lines)"
    else
        ((ERROR_COUNT++))
        echo -e "  ${YELLOW}⊘${RESET} Auth log ${DIM}(skipped)${RESET}"
    fi
    # syslog or messages
    verbose "Collecting: Syslog"
    if tail -500 /var/log/syslog > "${CASE_DIR}/syslog.txt" 2>/dev/null || \
       tail -500 /var/log/messages > "${CASE_DIR}/syslog.txt" 2>/dev/null; then
        ((COLLECT_COUNT++))
        echo -e "  ${GREEN}✓${RESET} Syslog (last 500 lines)"
    else
        ((ERROR_COUNT++))
        echo -e "  ${YELLOW}⊘${RESET} Syslog ${DIM}(skipped)${RESET}"
    fi
    # kern.log
    verbose "Collecting: Kern log"
    if tail -500 /var/log/kern.log > "${CASE_DIR}/kern-log.txt" 2>/dev/null; then
        ((COLLECT_COUNT++))
        echo -e "  ${GREEN}✓${RESET} Kern log (last 500 lines)"
    else
        ((ERROR_COUNT++))
        echo -e "  ${YELLOW}⊘${RESET} Kern log ${DIM}(skipped)${RESET}"
    fi
    collect_cmd "Journal (last hour)"    "${CASE_DIR}/journalctl-1h.txt"  journalctl --since "1 hour ago" --no-pager

    # --- Firewall ---
    section_header "Firewall Rules"
    collect_cmd "iptables rules"   "${CASE_DIR}/iptables-save.txt"  iptables-save
    collect_cmd "nftables rules"   "${CASE_DIR}/nftables.txt"       nft list ruleset
    collect_cmd "UFW status"       "${CASE_DIR}/ufw-status.txt"     ufw status verbose

    # --- Manifest & Archive ---
    section_header "Creating Tamper-Evident Archive"

    local manifest="${CASE_DIR}/manifest-sha256.txt"
    verbose "Building SHA-256 manifest"
    : > "${manifest}"
    while IFS= read -r -d '' file; do
        ${HASH_CMD} "${file}" >> "${manifest}" 2>/dev/null || true
    done < <(find "${CASE_DIR}" -type f ! -name "manifest-sha256.txt" -print0 2>/dev/null)
    echo -e "  ${GREEN}✓${RESET} Manifest created ($(wc -l < "${manifest}") files)"

    local archive_path="${OUTPUT_DIR}/ir-${CASE_ID}.tar.gz"
    tar -czf "${archive_path}" -C "$(dirname "${CASE_DIR}")" "$(basename "${CASE_DIR}")"
    echo -e "  ${GREEN}✓${RESET} Archive: ${archive_path}"

    local archive_hash
    archive_hash="$(${HASH_CMD} "${archive_path}" | awk '{print $1}')"
    echo "${archive_hash}  $(basename "${archive_path}")" > "${archive_path}.sha256"
    echo -e "  ${GREEN}✓${RESET} Archive hash: ${archive_hash}"

    # --- Summary ---
    section_header "Collection Summary"
    field "Case ID:"        "${CASE_ID}"
    field "Files collected:" "${COLLECT_COUNT}"
    field_color "Errors:"   "${YELLOW}${ERROR_COUNT}${RESET}"
    field "Elapsed:"        "$(elapsed)"
    field "Archive:"        "${archive_path}"
    field "SHA-256:"        "${archive_hash}"
    echo ""
}

#------------------------------------------------------------------------------
# MODE: QUICK SNAPSHOT
#------------------------------------------------------------------------------
do_quick() {
    local CASE_DIR="${OUTPUT_DIR}/ir-${CASE_ID}"
    mkdir -p "${CASE_DIR}"
    local snapshot="${CASE_DIR}/quick-snapshot.txt"

    section_header "Incident Response — 30-Second Snapshot"
    {
        echo "=== Quick Incident Snapshot ==="
        echo "Date:     $(date)"
        echo "Hostname: $(hostname)"
        echo ""

        echo "--- Top 20 Processes by CPU ---"
        ps aux --sort=-%cpu | head -21 2>/dev/null || true
        echo ""

        echo "--- Top 20 Processes by Memory ---"
        ps aux --sort=-%mem | head -21 2>/dev/null || true
        echo ""

        echo "--- Network Listeners ---"
        ss -tulnp 2>/dev/null || true
        echo ""

        echo "--- Logged-in Users ---"
        who 2>/dev/null || true
        echo ""

        echo "--- Last 10 Auth Failures ---"
        grep -i "failed\|failure" /var/log/auth.log 2>/dev/null | tail -10 || \
        grep -i "failed\|failure" /var/log/secure 2>/dev/null | tail -10 || \
        journalctl -p err --since "1 hour ago" --no-pager 2>/dev/null | grep -i "auth\|login\|failed" | tail -10 || \
        echo "(no auth failure data found)"
        echo ""
    } | tee "${snapshot}"

    log "Snapshot saved: ${snapshot}"
}

#------------------------------------------------------------------------------
# MODE: IOC SEARCH
#------------------------------------------------------------------------------
do_ioc() {
    if [[ -z "${IOC_PATTERN}" ]]; then
        die "IOC pattern required: --ioc <IP|domain|hash|filename>"
    fi

    section_header "IOC Search: ${IOC_PATTERN}"
    subsection "Scanning live system for indicator matches"
    local found=0

    local sources=(
        "Process list:ps auxwww"
        "Network connections:ss -anp"
        "Open files:lsof -n -P"
    )
    for entry in "${sources[@]}"; do
        local label="${entry%%:*}"
        local cmd="${entry#*:}"
        verbose "Searching ${label}"
        local matches
        matches="$(eval "${cmd}" 2>/dev/null | grep -i -- "${IOC_PATTERN}" 2>/dev/null)" || true
        if [[ -n "${matches}" ]]; then
            echo -e "  ${RED}⚑${RESET} ${BOLD}${label}${RESET}"
            echo "${matches}" | while IFS= read -r line; do
                echo "    ${line}"
            done
            ((found++)) || true
        else
            echo -e "  ${GREEN}✓${RESET} ${label} — clean"
        fi
    done

    local log_files=(
        "/var/log/auth.log"
        "/var/log/syslog"
        "/etc/hosts"
        "/etc/resolv.conf"
    )
    for lf in "${log_files[@]}"; do
        verbose "Searching ${lf}"
        local matches
        matches="$(grep -i -- "${IOC_PATTERN}" "${lf}" 2>/dev/null)" || true
        if [[ -n "${matches}" ]]; then
            echo -e "  ${RED}⚑${RESET} ${BOLD}${lf}${RESET}"
            echo "${matches}" | tail -20 | while IFS= read -r line; do
                echo "    ${line}"
            done
            ((found++)) || true
        else
            echo -e "  ${GREEN}✓${RESET} ${lf} — clean"
        fi
    done

    # Crontabs
    verbose "Searching crontabs"
    local cron_matches
    cron_matches="$(
        cat /etc/crontab 2>/dev/null
        while IFS=: read -r user _; do
            crontab -l -u "${user}" 2>/dev/null || true
        done < /etc/passwd 2>/dev/null
    )" || true
    local cron_hits
    cron_hits="$(echo "${cron_matches}" | grep -i -- "${IOC_PATTERN}" 2>/dev/null)" || true
    if [[ -n "${cron_hits}" ]]; then
        echo -e "  ${RED}⚑${RESET} ${BOLD}Crontabs${RESET}"
        echo "${cron_hits}" | while IFS= read -r line; do
            echo "    ${line}"
        done
        ((found++)) || true
    else
        echo -e "  ${GREEN}✓${RESET} Crontabs — clean"
    fi

    echo ""
    if [[ "${found}" -gt 0 ]]; then
        echo -e "${RED}${BOLD}⚠  IOC found in ${found} source(s)${RESET}"
    else
        echo -e "${GREEN}${BOLD}✓  No matches found for IOC${RESET}"
    fi
}

#------------------------------------------------------------------------------
# MODE: TIMELINE
#------------------------------------------------------------------------------
do_timeline() {
    section_header "Incident Timeline Builder"
    local tmpfile
    tmpfile="$(mktemp /tmp/ir-timeline.XXXXXX)"

    verbose "Extracting events from auth.log"
    if [[ -r /var/log/auth.log ]]; then
        awk '{print $0}' /var/log/auth.log >> "${tmpfile}" 2>/dev/null || true
    elif [[ -r /var/log/secure ]]; then
        awk '{print $0}' /var/log/secure >> "${tmpfile}" 2>/dev/null || true
    fi

    verbose "Extracting events from syslog"
    if [[ -r /var/log/syslog ]]; then
        awk '{print $0}' /var/log/syslog >> "${tmpfile}" 2>/dev/null || true
    elif [[ -r /var/log/messages ]]; then
        awk '{print $0}' /var/log/messages >> "${tmpfile}" 2>/dev/null || true
    fi

    verbose "Extracting events from kern.log"
    if [[ -r /var/log/kern.log ]]; then
        awk '{print $0}' /var/log/kern.log >> "${tmpfile}" 2>/dev/null || true
    fi

    verbose "Extracting events from journalctl"
    journalctl --since "24 hours ago" --no-pager -o short-iso 2>/dev/null >> "${tmpfile}" || true

    local total_events
    total_events="$(wc -l < "${tmpfile}")"

    if [[ "${total_events}" -eq 0 ]]; then
        warn "No log events found — insufficient permissions or empty logs"
        rm -f "${tmpfile}"
        return
    fi

    log "Sorting ${total_events} events chronologically..."
    sort -t' ' -k1,3 "${tmpfile}" | uniq > "${tmpfile}.sorted" 2>/dev/null || \
        sort "${tmpfile}" | uniq > "${tmpfile}.sorted" 2>/dev/null

    local sorted_count
    sorted_count="$(wc -l < "${tmpfile}.sorted")"
    field "Total events:" "${sorted_count} (deduplicated)"
    echo ""

    echo -e "${BOLD}Last 50 events:${RESET}"
    tail -50 "${tmpfile}.sorted"

    echo ""
    field "Full timeline:" "${tmpfile}.sorted"
    rm -f "${tmpfile}"
}

#------------------------------------------------------------------------------
# HELP
#------------------------------------------------------------------------------
show_help() {
    cat <<EOF
${BOLD}${SCRIPT_NAME}${RESET} — Live incident volatile data capture and IOC search

${BOLD}MODES${RESET}
  --collect          Full volatile data collection into tamper-evident archive
  --quick            30-second rapid snapshot (top processes, network, users)
  --ioc INDICATOR    Search for a specific IOC (IP, domain, hash, filename)
  --timeline         Build incident timeline from system logs

${BOLD}OPTIONS${RESET}
  --output DIR       Output directory (default: /var/log/incident-response)
  --case-id ID       Override auto-generated case ID
  --verbose          Enable verbose output
  --no-color         Disable colored output
  --help             Show this help message

${BOLD}ENVIRONMENT VARIABLES${RESET}
  VERBOSE=1          Same as --verbose
  COLOR=never        Same as --no-color
  COLOR=always       Force colors even when not a terminal

${BOLD}EXAMPLES${RESET}
  sudo ./${SCRIPT_NAME} --collect
  sudo ./${SCRIPT_NAME} --quick
  sudo ./${SCRIPT_NAME} --ioc 192.168.1.100
  sudo ./${SCRIPT_NAME} --ioc evil.example.com
  sudo ./${SCRIPT_NAME} --timeline
  sudo ./${SCRIPT_NAME} --collect --output /tmp/cases --case-id CASE-42
  VERBOSE=1 sudo ./${SCRIPT_NAME} --collect
EOF
}

#------------------------------------------------------------------------------
# ARGUMENT PARSING
#------------------------------------------------------------------------------
parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --collect)   RUN_MODE="collect" ;;
            --quick)     RUN_MODE="quick" ;;
            --ioc)
                RUN_MODE="ioc"
                if [[ -z "${2:-}" ]]; then
                    die "--ioc requires an indicator argument"
                fi
                IOC_PATTERN="$2"; shift
                ;;
            --timeline)  RUN_MODE="timeline" ;;
            --output)
                if [[ -z "${2:-}" ]]; then
                    die "--output requires a directory argument"
                fi
                OUTPUT_DIR="$2"; shift
                ;;
            --case-id)
                if [[ -z "${2:-}" ]]; then
                    die "--case-id requires an ID argument"
                fi
                CASE_ID="$2"; shift
                ;;
            --verbose)   VERBOSE=1 ;;
            --no-color)  COLOR="never" ;;
            --help|-h)   show_help; exit 0 ;;
            *)           die "Unknown option: $1 (see --help)" ;;
        esac
        shift
    done

    if [[ -z "${RUN_MODE}" ]]; then err "No mode specified"; echo ""; show_help; exit 1; fi
}

#------------------------------------------------------------------------------
# MAIN
#------------------------------------------------------------------------------
main() {
    parse_args "$@"
    setup_colors
    START_TIME="$(date +%s)"

    if [[ "$(id -u)" -ne 0 ]]; then
        warn "Running without root — some data may be inaccessible. Consider: sudo $0"
    fi

    case "${RUN_MODE}" in
        collect)  do_collect  ;;
        quick)    do_quick    ;;
        ioc)      do_ioc      ;;
        timeline) do_timeline ;;
        *)        die "Unknown mode: ${RUN_MODE}" ;;
    esac
}

main "$@"
