#!/bin/bash

# Incident Response Phase 2: Containment & Eradication (Remote Execution)
# Usage: ./phase2.sh <target_host>
# This script runs from operator machine and executes commands remotely via SSH

set -euo pipefail

# Configuration
SCRIPT_VERSION="1.0"
PHASE="phase2-containment"
SSH_USER="ir"

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Check if target host parameter is provided
if [[ $# -ne 1 ]]; then
    echo -e "${RED}Usage: $0 <target_host>${NC}"
    echo -e "${YELLOW}Example: $0 webserver01.company.com${NC}"
    echo -e "${YELLOW}Note: Ensure you have SSH access as 'ir' user to the target host${NC}"
    exit 1
fi

TARGET_HOST="$1"
TIMESTAMP=$(date '+%Y%m%d_%H%M%S')
BASE_DIR="./incident_response"
HOST_DIR="$BASE_DIR/$TARGET_HOST"
PHASE_DIR="$HOST_DIR/$PHASE"
LOG_FILE="$PHASE_DIR/${PHASE}_${TIMESTAMP}.log"
SUMMARY_FILE="$PHASE_DIR/${PHASE}_summary_${TIMESTAMP}.txt"

# Create local directory structure on operator machine
mkdir -p "$PHASE_DIR"

# Logging functions
log() {
    echo -e "${GREEN}[$(date '+%Y-%m-%d %H:%M:%S')] $1${NC}" | tee -a "$LOG_FILE"
}

error() {
    echo -e "${RED}[ERROR] $1${NC}" | tee -a "$LOG_FILE"
    exit 1
}

warn() {
    echo -e "${YELLOW}[WARNING] $1${NC}" | tee -a "$LOG_FILE"
}

info() {
    echo -e "${BLUE}[INFO] $1${NC}" | tee -a "$LOG_FILE"
}

critical() {
    echo -e "${RED}[CRITICAL] $1${NC}" | tee -a "$LOG_FILE"
}

# Function to execute remote command via SSH
remote_cmd() {
    local cmd="$1"
    local description="$2"
    local critical_action="$3"
    
    if [[ "$critical_action" == "true" ]]; then
        critical "EXECUTING: $description"
    else
        log "Running: $description"
    fi
    
    echo "=== $description ===" >> "$LOG_FILE"
    echo "Target: $TARGET_HOST" >> "$LOG_FILE"
    echo "Command: $cmd" >> "$LOG_FILE"
    echo "Output:" >> "$LOG_FILE"
    
    if ssh -o ConnectTimeout=10 -o StrictHostKeyChecking=no "$SSH_USER@$TARGET_HOST" "$cmd" >> "$LOG_FILE" 2>&1; then
        if [[ "$critical_action" == "true" ]]; then
            critical "✓ $description completed successfully"
        else
            info "✓ $description completed"
        fi
    else
        warn "✗ $description failed or returned non-zero exit code (may be expected)"
    fi
    echo "" >> "$LOG_FILE"
}

# Function to safely run containment commands
safe_containment() {
    local cmd="$1"
    local description="$2"
    
    log "CONTAINMENT: $description"
    echo "=== CONTAINMENT: $description ===" >> "$LOG_FILE"
    echo "Target: $TARGET_HOST" >> "$LOG_FILE"
    echo "Command: $cmd" >> "$LOG_FILE"
    echo "Output:" >> "$LOG_FILE"
    
    if timeout 30 ssh -o ConnectTimeout=10 -o StrictHostKeyChecking=no "$SSH_USER@$TARGET_HOST" "$cmd" >> "$LOG_FILE" 2>&1; then
        critical "✓ CONTAINMENT SUCCESS: $description"
    else
        warn "✗ CONTAINMENT WARNING: $description (may be expected if no threats found)"
    fi
    echo "" >> "$LOG_FILE"
}

# Confirmation prompt for destructive actions
confirm_action() {
    local action="$1"
    echo -e "${YELLOW}WARNING: About to perform destructive action on $TARGET_HOST: $action${NC}"
    echo -e "${YELLOW}This may disrupt system operations. Continue? [y/N]${NC}"
    read -r response
    if [[ ! "$response" =~ ^[Yy]$ ]]; then
        warn "Action cancelled by user: $action"
        return 1
    fi
    return 0
}

# Test SSH connectivity
test_ssh_connection() {
    log "Testing SSH connectivity to $TARGET_HOST..."
    if ! ssh -o ConnectTimeout=5 -o StrictHostKeyChecking=no "$SSH_USER@$TARGET_HOST" "echo 'SSH connection successful'" >/dev/null 2>&1; then
        error "Cannot connect to $TARGET_HOST as $SSH_USER. Ensure SSH keys are configured and IR lockdown was deployed."
    fi
    info "✓ SSH connection to $TARGET_HOST successful"
}

# Start containment
log "=== INCIDENT RESPONSE PHASE 2: CONTAINMENT & ERADICATION ==="
log "Target Host: $TARGET_HOST"
log "Operator Machine: $(hostname)"
log "Timestamp: $TIMESTAMP"
log "Local Log Directory: $PHASE_DIR"
critical "STARTING IMMEDIATE THREAT CONTAINMENT"

# Test SSH connection first
test_ssh_connection

# Create summary header
cat > "$SUMMARY_FILE" << EOF
=== INCIDENT RESPONSE PHASE 2 SUMMARY ===
Target Host: $TARGET_HOST
Operator Machine: $(hostname)
Timestamp: $TIMESTAMP
Script Version: $SCRIPT_VERSION
Local Log Directory: $PHASE_DIR

EOF

# 1. Immediate Process Termination
log "Phase 2.1: Immediate Threat Process Termination"
safe_containment "pkill -f 'nc -l' || true" "Kill netcat listeners"
safe_containment "pkill -f 'python.*socket' || true" "Kill Python socket processes"
safe_containment "pkill -f 'perl.*socket' || true" "Kill Perl socket processes"
safe_containment "pkill -f '/tmp/' || true" "Kill processes running from /tmp"
safe_containment "killall -9 nc netcat socat 2>/dev/null || true" "Force kill network tools"
safe_containment "fuser -k 4444/tcp 5555/tcp 6666/tcp 7777/tcp 8080/tcp 9999/tcp 31337/tcp 2>/dev/null || true" "Kill processes on backdoor ports"

# 2. Network Isolation - Backup current rules first
log "Phase 2.2: Network Isolation"
remote_cmd "sudo iptables-save > /tmp/iptables_backup_${TIMESTAMP}.rules" "Backup current iptables rules" "false"

# Block backdoor ports
log "Blocking common backdoor ports"
for port in 4444 5555 6666 7777 8080 9999 31337; do
    safe_containment "sudo iptables -I INPUT -p tcp --dport $port -j DROP 2>/dev/null || true" "Block inbound port $port"
    safe_containment "sudo iptables -I OUTPUT -p tcp --dport $port -j DROP 2>/dev/null || true" "Block outbound port $port"
done

# Save firewall rules
remote_cmd "sudo iptables-save > /etc/iptables/rules.v4 2>/dev/null || sudo iptables-save > /tmp/iptables_active_${TIMESTAMP}.rules" "Save firewall rules" "true"

# 3. User Account Security
log "Phase 2.3: User Account Containment"
remote_cmd "cut -d: -f1 /etc/passwd | grep -v -E '^(root|daemon|bin|sys|sync|games|man|lp|mail|news|uucp|proxy|www-data|backup|list|irc|gnats|nobody|systemd|syslog|_apt|tss|strongswan|ntp|redsocks|rwhod|iodine|tcpdump|miredo)$' > /tmp/regular_users.txt" "Identify regular user accounts" "false"

# Lock suspicious service accounts with shells
safe_containment "sudo chsh -s /bin/false www-data 2>/dev/null || true" "Remove shell from www-data"
safe_containment "sudo chsh -s /bin/false mysql 2>/dev/null || true" "Remove shell from mysql"
safe_containment "sudo chsh -s /bin/false apache 2>/dev/null || true" "Remove shell from apache"
safe_containment "sudo chsh -s /bin/false nginx 2>/dev/null || true" "Remove shell from nginx"

# Clear all SSH authorized keys (backup first)
log "Backing up and clearing SSH authorized keys"
remote_cmd "find /home -name 'authorized_keys' -exec cp {} {}.backup_${TIMESTAMP} \\; 2>/dev/null || true" "Backup authorized_keys files" "false"
remote_cmd "find /root -name 'authorized_keys' -exec cp {} {}.backup_${TIMESTAMP} \\; 2>/dev/null || true" "Backup root authorized_keys" "false"

if confirm_action "Clear all SSH authorized_keys files except IR user"; then
    safe_containment "find /home -name 'authorized_keys' ! -path '/home/ir/*' -exec cp /dev/null {} \\; 2>/dev/null || true" "Clear user authorized_keys"
    safe_containment "find /root -name 'authorized_keys' -exec cp /dev/null {} \\; 2>/dev/null || true" "Clear root authorized_keys"
fi

# 4. File System Eradication
log "Phase 2.4: File System Eradication"

# Find and log webshells before removal
remote_cmd "find /var/www -name '*.php' -exec grep -l 'eval\\|system\\|exec\\|shell_exec\\|passthru\\|base64_decode' {} \\; 2>/dev/null | head -20 > /tmp/webshells_found.txt || echo 'No webshells found' > /tmp/webshells_found.txt" "Identify potential webshells" "false"

if confirm_action "Remove identified webshells and suspicious files"; then
    safe_containment "find /var/www -name '*.php' -exec grep -l 'eval\\|system\\|exec\\|shell_exec\\|passthru\\|base64_decode' {} \\; 2>/dev/null | head -20 | xargs rm -f 2>/dev/null || true" "Remove PHP webshells"
    safe_containment "find /tmp -type f -name '.*' -delete 2>/dev/null || true" "Remove hidden files from /tmp"
    safe_containment "find /tmp -type f -executable -newer /var/log/wtmp -delete 2>/dev/null || true" "Remove recent executables from /tmp"
    safe_containment "find /dev/shm -type f -delete 2>/dev/null || true" "Clean /dev/shm"
    safe_containment "find /var/tmp -type f -executable -delete 2>/dev/null || true" "Remove executables from /var/tmp"
fi

# 5. Service Hardening
log "Phase 2.5: Service Hardening"
safe_containment "sudo systemctl stop telnet 2>/dev/null || true" "Stop telnet service"
safe_containment "sudo systemctl disable telnet 2>/dev/null || true" "Disable telnet service"
safe_containment "sudo systemctl stop rsh 2>/dev/null || true" "Stop rsh service"
safe_containment "sudo systemctl disable rsh 2>/dev/null || true" "Disable rsh service"

# Kill all user sessions except current
remote_cmd "who | grep -v \$(whoami) | awk '{print \$2}' > /tmp/active_sessions.txt" "List active user sessions" "false"
if confirm_action "Kill all other user sessions"; then
    safe_containment "who | grep -v \$(whoami) | awk '{print \$2}' | xargs -I {} pkill -t {} 2>/dev/null || true" "Terminate other user sessions"
fi

# 6. SSH Hardening
log "Phase 2.6: SSH Configuration Hardening"
remote_cmd "sudo cp /etc/ssh/sshd_config /etc/ssh/sshd_config.backup_${TIMESTAMP}" "Backup SSH configuration" "false"

# Apply restrictive SSH settings while preserving IR access
if confirm_action "Apply hardened SSH configuration (will restrict access to IR user only)"; then
    safe_containment 'sudo bash -c "cat > /etc/ssh/sshd_config << EOF
Port 22
Protocol 2
PermitRootLogin no
PasswordAuthentication yes
PubkeyAuthentication yes
ChallengeResponseAuthentication no
UsePAM yes
X11Forwarding no
PrintMotd no
MaxAuthTries 3
ClientAliveInterval 300
ClientAliveCountMax 2
AllowUsers ir
EOF"' "Apply hardened SSH config"
    safe_containment "sudo systemctl restart ssh 2>/dev/null || sudo systemctl restart sshd 2>/dev/null" "Restart SSH service"
fi

# 7. Validation
log "Phase 2.7: Post-Containment Validation"
remote_cmd "ps aux | grep -E '(nc|netcat|python.*socket|perl.*socket|bash.*tcp)' | grep -v grep > /tmp/remaining_suspicious_processes.txt || echo 'No suspicious processes found' > /tmp/remaining_suspicious_processes.txt" "Check for remaining suspicious processes" "false"
remote_cmd "ss -tulpn 2>/dev/null | grep -E ':(4444|5555|6666|7777|8080|9999|31337)' > /tmp/remaining_backdoor_ports.txt || echo 'No backdoor ports detected' > /tmp/remaining_backdoor_ports.txt" "Check for remaining backdoor ports" "false"
remote_cmd "sudo iptables -L -n > /tmp/final_iptables_rules.txt" "Document final firewall rules" "false"

# Copy validation files to operator machine
ssh -o ConnectTimeout=10 -o StrictHostKeyChecking=no "$SSH_USER@$TARGET_HOST" "cat /tmp/remaining_suspicious_processes.txt" > "$PHASE_DIR/remaining_suspicious_processes.txt" 2>/dev/null || true
ssh -o ConnectTimeout=10 -o StrictHostKeyChecking=no "$SSH_USER@$TARGET_HOST" "cat /tmp/remaining_backdoor_ports.txt" > "$PHASE_DIR/remaining_backdoor_ports.txt" 2>/dev/null || true
ssh -o ConnectTimeout=10 -o StrictHostKeyChecking=no "$SSH_USER@$TARGET_HOST" "cat /tmp/final_iptables_rules.txt" > "$PHASE_DIR/final_iptables_rules.txt" 2>/dev/null || true

# 8. Generate Summary
log "Phase 2.8: Generating Containment Summary"

{
    echo "=== CONTAINMENT ACTIONS COMPLETED ==="
    echo ""
    
    # Check if any suspicious processes remain
    REMAINING_PROCESSES=$(cat "$PHASE_DIR/remaining_suspicious_processes.txt" 2>/dev/null | grep -v "No suspicious processes" | wc -l || echo "0")
    if [[ $REMAINING_PROCESSES -gt 0 ]]; then
        echo "⚠️  WARNING: $REMAINING_PROCESSES suspicious processes still running"
    else
        echo "✓ No suspicious processes detected"
    fi
    
    # Check if any backdoor ports remain
    REMAINING_PORTS=$(cat "$PHASE_DIR/remaining_backdoor_ports.txt" 2>/dev/null | grep -v "No backdoor ports" | wc -l || echo "0")
    if [[ $REMAINING_PORTS -gt 0 ]]; then
        echo "⚠️  WARNING: $REMAINING_PORTS backdoor ports still active"
    else
        echo "✓ No backdoor ports detected"
    fi
    
    echo ""
    echo "=== ACTIONS TAKEN ==="
    echo "• Terminated suspicious processes"
    echo "• Blocked backdoor ports with iptables"
    echo "• Cleared SSH authorized keys (except IR user)"
    echo "• Removed webshells and suspicious files"
    echo "• Hardened service configurations"
    echo "• Applied restrictive SSH settings"
    
    echo ""
    echo "=== NEXT STEPS ==="
    echo "1. Verify all threats have been contained"
    echo "2. Monitor system for 15-30 minutes"
    echo "3. Begin Phase 3: Recovery operations"
    echo "4. Update IR team on containment status"
    echo ""
    echo "=== FILES GENERATED (LOCAL) ==="
    ls -la "$PHASE_DIR" | grep -v "^total"
    
} >> "$SUMMARY_FILE"

# Final log entries
critical "Phase 2 Containment completed"
log "Summary file: $SUMMARY_FILE"
log "All output files saved locally to: $PHASE_DIR"

# Display summary
echo
echo -e "${GREEN}========================================${NC}"
echo -e "${GREEN}  PHASE 2 CONTAINMENT COMPLETE         ${NC}"
echo -e "${GREEN}========================================${NC}"
echo
echo -e "${YELLOW}Target Host:${NC} $TARGET_HOST"
echo -e "${YELLOW}Local Log Directory:${NC} $PHASE_DIR"
echo -e "${YELLOW}Summary File:${NC} $SUMMARY_FILE"
echo
echo -e "${RED}CRITICAL: System has been contained. Monitor for stability before proceeding.${NC}"
echo
echo -e "${BLUE}Containment Summary:${NC}"
cat "$SUMMARY_FILE" | grep -A 15 "=== CONTAINMENT ACTIONS COMPLETED ==="
echo
echo -e "${YELLOW}Next: Monitor system stability, then proceed to Phase 3 (Recovery)${NC}"
echo -e "${GREEN}Command: ./phase3.sh $TARGET_HOST${NC}"