#!/usr/bin/env bash

#########################################################################################
#### vpc-audit.sh — Audit AWS VPC configuration, security groups, and networking     ####
#### Reports on open security groups, missing flow logs, peering, subnet utilization  ####
#### Requires: bash 4+, aws-cli v2, jq                                              ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.01                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   export AWS_PROFILE="production"                                               ####
####   ./vpc-audit.sh --full                                                         ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
AWS_REGION="${AWS_REGION:-}"
VPC_ID="${VPC_ID:-}"
OUTPUT_FORMAT="${OUTPUT_FORMAT:-text}"
VERBOSE="${VERBOSE:-false}"
COLOR="${COLOR:-auto}"

# Section flags
RUN_FULL="false"
RUN_SGS="false"
RUN_FLOW="false"
RUN_PEER="false"
RUN_SUBNETS="false"
RUN_ENDPOINTS="false"
RUN_NACLS="false"

# ── State ─────────────────────────────────────────────────────────────
SCRIPT_NAME="$(basename "$0")"
readonly SCRIPT_NAME
START_TIME=""
PASS_COUNT=0
WARN_COUNT=0
CRIT_COUNT=0

readonly DANGEROUS_PORTS="22 3389 3306 5432 6379 27017 9200 9300 11211"

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

# ── Logging ───────────────────────────────────────────────────────────
log()     { echo -e "${BLUE}[INFO]${RESET} $*"; }
warn()    { echo -e "${YELLOW}[WARN]${RESET} $*" >&2; ((WARN_COUNT++)) || true; }
err()     { echo -e "${RED}[ERROR]${RESET} $*" >&2; }
verbose() { if [[ "$VERBOSE" == "true" ]]; then echo -e "${DIM}[DEBUG]${RESET} $*"; fi; }

die() { err "$*"; exit 1; }

# ── AWS CLI wrapper ───────────────────────────────────────────────────
aws_cmd() {
    local args=("$@")
    [[ -n "$AWS_REGION" ]] && args+=(--region "$AWS_REGION")
    verbose "aws ${args[*]}"
    aws "${args[@]}"
}

# ── Dependency check ──────────────────────────────────────────────────
check_deps() {
    for cmd in aws jq; do
        if ! command -v "$cmd" &>/dev/null; then
            die "${cmd} is required but not installed"
        fi
    done

    if ! aws sts get-caller-identity &>/dev/null; then
        die "AWS credentials not configured or expired"
    fi

    if [[ -z "$AWS_REGION" ]]; then
        AWS_REGION=$(aws configure get region 2>/dev/null || echo "")
        if [[ -z "$AWS_REGION" ]]; then
            die "AWS_REGION is required"
        fi
    fi
    verbose "Using region: ${AWS_REGION}"
}

should_show() {
    [[ "$RUN_FULL" == "true" ]] && return 0
    case "$1" in
        sgs)       [[ "$RUN_SGS" == "true" ]] ;;
        flow)      [[ "$RUN_FLOW" == "true" ]] ;;
        peer)      [[ "$RUN_PEER" == "true" ]] ;;
        subnets)   [[ "$RUN_SUBNETS" == "true" ]] ;;
        endpoints) [[ "$RUN_ENDPOINTS" == "true" ]] ;;
        nacls)     [[ "$RUN_NACLS" == "true" ]] ;;
        *)         return 1 ;;
    esac
}

# ── Get VPC list ──────────────────────────────────────────────────────
get_vpcs() {
    if [[ -n "$VPC_ID" ]]; then
        echo "$VPC_ID"
    else
        aws_cmd ec2 describe-vpcs \
            --query 'Vpcs[*].VpcId' \
            --output text 2>/dev/null | tr '\t' '\n'
    fi
}

get_vpc_name() {
    local vid="$1"
    # shellcheck disable=SC2016
    aws_cmd ec2 describe-vpcs \
        --vpc-ids "$vid" \
        --query 'Vpcs[0].Tags[?Key==`Name`].Value | [0]' \
        --output text 2>/dev/null || echo "N/A"
}

# ══════════════════════════════════════════════════════════════════════
#                       SECURITY GROUPS
# ══════════════════════════════════════════════════════════════════════

audit_security_groups() {
    log "Auditing security groups..."

    local sg_filter=()
    [[ -n "$VPC_ID" ]] && sg_filter=(--filters "Name=vpc-id,Values=${VPC_ID}")

    local sgs_json
    sgs_json=$(aws_cmd ec2 describe-security-groups \
        "${sg_filter[@]}" \
        --query 'SecurityGroups[*].{Id:GroupId,Name:GroupName,VpcId:VpcId,Ingress:IpPermissions}' \
        --output json 2>/dev/null)

    local total_sgs
    total_sgs=$(echo "$sgs_json" | jq 'length')
    log "Found ${total_sgs} security group(s)"

    # Check for 0.0.0.0/0 ingress
    echo ""
    echo -e "  ${BOLD}Open Ingress Rules (0.0.0.0/0)${RESET}"
    printf "  ${BOLD}%-14s %-24s %-10s %-8s %s${RESET}\n" "SG ID" "NAME" "PORT" "PROTO" "SEVERITY"
    printf "  %s\n" "$(printf '%.0s─' {1..70})"

    local open_count=0
    echo "$sgs_json" | jq -c '.[]' | while IFS= read -r sg; do
        local sg_id sg_name
        sg_id=$(echo "$sg" | jq -r '.Id')
        sg_name=$(echo "$sg" | jq -r '.Name')

        echo "$sg" | jq -c '.Ingress[]?' | while IFS= read -r rule; do
            local has_open
            has_open=$(echo "$rule" | jq '[.IpRanges[]? | select(.CidrIp == "0.0.0.0/0")] | length')

            if [[ "$has_open" -gt 0 ]]; then
                local proto from_port to_port
                proto=$(echo "$rule" | jq -r '.IpProtocol')
                from_port=$(echo "$rule" | jq -r '.FromPort // "all"')
                to_port=$(echo "$rule" | jq -r '.ToPort // "all"')

                local port_display="$from_port"
                [[ "$from_port" != "$to_port" ]] && port_display="${from_port}-${to_port}"
                [[ "$proto" == "-1" ]] && { proto="all"; port_display="all"; }

                local severity="WARN"
                local sev_color="$YELLOW"
                for dp in $DANGEROUS_PORTS; do
                    if [[ "$from_port" == "$dp" || "$port_display" == "all" ]]; then
                        severity="CRITICAL"
                        sev_color="$RED"
                        ((CRIT_COUNT++)) || true
                        break
                    fi
                done
                [[ "$severity" == "WARN" ]] && { ((WARN_COUNT++)) || true; }

                printf "  %-14s %-24s %-10s %-8s ${sev_color}%s${RESET}\n" \
                    "$sg_id" "${sg_name:0:24}" "$port_display" "$proto" "$severity"
                ((open_count++)) || true
            fi
        done
    done

    if [[ "$open_count" -eq 0 ]]; then
        echo -e "  ${GREEN}No open ingress rules found ✓${RESET}"
        ((PASS_COUNT++)) || true
    fi

    # Unused security groups
    echo ""
    echo -e "  ${BOLD}Unused Security Groups${RESET}"
    local unused=0
    echo "$sgs_json" | jq -r '.[].Id' | while IFS= read -r sg_id; do
        local eni_count
        eni_count=$(aws_cmd ec2 describe-network-interfaces \
            --filters "Name=group-id,Values=${sg_id}" \
            --query 'NetworkInterfaces | length(@)' \
            --output text 2>/dev/null) || continue

        if [[ "$eni_count" -eq 0 ]]; then
            local sg_name
            sg_name=$(echo "$sgs_json" | jq -r ".[] | select(.Id == \"${sg_id}\") | .Name")
            [[ "$sg_name" == "default" ]] && continue
            echo -e "    ${YELLOW}⊘${RESET} ${sg_id} (${sg_name}) — not attached to any ENI"
            ((unused++)) || true
        fi
    done

    if [[ "$unused" -eq 0 ]]; then
        echo -e "    ${GREEN}No unused security groups ✓${RESET}"
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                       FLOW LOGS
# ══════════════════════════════════════════════════════════════════════

audit_flow_logs() {
    log "Checking VPC flow logs..."

    echo ""
    echo -e "  ${BOLD}VPC Flow Log Status${RESET}"
    printf "  ${BOLD}%-22s %-24s %s${RESET}\n" "VPC ID" "NAME" "FLOW LOGS"
    printf "  %s\n" "$(printf '%.0s─' {1..60})"

    local vpcs
    vpcs=$(get_vpcs)

    while IFS= read -r vpc_id; do
        [[ -z "$vpc_id" ]] && continue
        local vpc_name
        vpc_name=$(get_vpc_name "$vpc_id")
        [[ "$vpc_name" == "None" ]] && vpc_name="N/A"

        local flow_count
        flow_count=$(aws_cmd ec2 describe-flow-logs \
            --filter "Name=resource-id,Values=${vpc_id}" \
            --query 'FlowLogs | length(@)' \
            --output text 2>/dev/null) || flow_count=0

        if [[ "$flow_count" -gt 0 ]]; then
            printf "  %-22s %-24s ${GREEN}enabled (%s)${RESET}\n" "$vpc_id" "${vpc_name:0:24}" "$flow_count"
            ((PASS_COUNT++)) || true
        else
            printf "  %-22s %-24s ${RED}NONE${RESET}\n" "$vpc_id" "${vpc_name:0:24}"
            ((CRIT_COUNT++)) || true
        fi
    done <<< "$vpcs"
}

# ══════════════════════════════════════════════════════════════════════
#                       VPC PEERING
# ══════════════════════════════════════════════════════════════════════

audit_peering() {
    log "Checking VPC peering connections..."

    local peering_json
    peering_json=$(aws_cmd ec2 describe-vpc-peering-connections \
        --query 'VpcPeeringConnections[*].{Id:VpcPeeringConnectionId,Status:Status.Code,Requester:RequesterVpcInfo.VpcId,Accepter:AccepterVpcInfo.VpcId,RequesterCIDR:RequesterVpcInfo.CidrBlock,AccepterCIDR:AccepterVpcInfo.CidrBlock}' \
        --output json 2>/dev/null)

    local total
    total=$(echo "$peering_json" | jq 'length')

    echo ""
    echo -e "  ${BOLD}VPC Peering Connections${RESET}"

    if [[ "$total" -eq 0 ]]; then
        echo "    No peering connections found"
        return
    fi

    printf "  ${BOLD}%-26s %-14s %-14s %s${RESET}\n" "PEERING ID" "REQUESTER" "ACCEPTER" "STATUS"
    printf "  %s\n" "$(printf '%.0s─' {1..65})"

    echo "$peering_json" | jq -c '.[]' | while IFS= read -r peer; do
        local peer_id status requester accepter
        peer_id=$(echo "$peer" | jq -r '.Id')
        status=$(echo "$peer" | jq -r '.Status')
        requester=$(echo "$peer" | jq -r '.Requester')
        accepter=$(echo "$peer" | jq -r '.Accepter')

        local color="$GREEN"
        [[ "$status" != "active" ]] && color="$YELLOW"

        printf "  %-26s %-14s %-14s ${color}%s${RESET}\n" "$peer_id" "$requester" "$accepter" "$status"
    done
}

# ══════════════════════════════════════════════════════════════════════
#                       SUBNET UTILIZATION
# ══════════════════════════════════════════════════════════════════════

audit_subnets() {
    log "Checking subnet utilization..."

    local subnet_filter=()
    [[ -n "$VPC_ID" ]] && subnet_filter=(--filters "Name=vpc-id,Values=${VPC_ID}")

    local subnets_json
    subnets_json=$(aws_cmd ec2 describe-subnets \
        "${subnet_filter[@]}" \
        --query 'Subnets[*].{Id:SubnetId,VpcId:VpcId,AZ:AvailabilityZone,CIDR:CidrBlock,Available:AvailableIpAddressCount}' \
        --output json 2>/dev/null)

    echo ""
    echo -e "  ${BOLD}Subnet IP Utilization${RESET}"
    printf "  ${BOLD}%-18s %-14s %-16s %-20s %8s %s${RESET}\n" "SUBNET" "VPC" "AZ" "CIDR" "AVAIL" "STATUS"
    printf "  %s\n" "$(printf '%.0s─' {1..90})"

    echo "$subnets_json" | jq -c '.[]' | while IFS= read -r subnet; do
        local sid vpc_id az cidr avail
        sid=$(echo "$subnet" | jq -r '.Id')
        vpc_id=$(echo "$subnet" | jq -r '.VpcId')
        az=$(echo "$subnet" | jq -r '.AZ')
        cidr=$(echo "$subnet" | jq -r '.CIDR')
        avail=$(echo "$subnet" | jq -r '.Available')

        # Calculate total IPs from CIDR
        local prefix
        prefix=$(echo "$cidr" | cut -d/ -f2)
        local total_ips=$(( (1 << (32 - prefix)) - 5 ))
        [[ "$total_ips" -lt 0 ]] && total_ips=0

        local pct_free=100
        if [[ "$total_ips" -gt 0 ]]; then
            pct_free=$(( avail * 100 / total_ips ))
        fi

        local status="OK"
        local color="$GREEN"
        if [[ "$pct_free" -lt 10 ]]; then
            status="CRITICAL"
            color="$RED"
            ((CRIT_COUNT++)) || true
        elif [[ "$pct_free" -lt 20 ]]; then
            status="LOW"
            color="$YELLOW"
            ((WARN_COUNT++)) || true
        else
            ((PASS_COUNT++)) || true
        fi

        printf "  %-18s %-14s %-16s %-20s %8s ${color}%s${RESET}\n" \
            "$sid" "${vpc_id:0:14}" "$az" "$cidr" "$avail" "$status"
    done
}

# ══════════════════════════════════════════════════════════════════════
#                       VPC ENDPOINTS
# ══════════════════════════════════════════════════════════════════════

audit_endpoints() {
    log "Listing VPC endpoints..."

    local ep_filter=()
    [[ -n "$VPC_ID" ]] && ep_filter=(--filters "Name=vpc-id,Values=${VPC_ID}")

    local ep_json
    ep_json=$(aws_cmd ec2 describe-vpc-endpoints \
        "${ep_filter[@]}" \
        --query 'VpcEndpoints[*].{Id:VpcEndpointId,Service:ServiceName,Type:VpcEndpointType,VpcId:VpcId,State:State}' \
        --output json 2>/dev/null)

    local total
    total=$(echo "$ep_json" | jq 'length')

    echo ""
    echo -e "  ${BOLD}VPC Endpoints${RESET}"

    if [[ "$total" -eq 0 ]]; then
        echo "    No VPC endpoints found"
        return
    fi

    printf "  ${BOLD}%-26s %-40s %-12s %s${RESET}\n" "ENDPOINT ID" "SERVICE" "TYPE" "STATE"
    printf "  %s\n" "$(printf '%.0s─' {1..85})"

    echo "$ep_json" | jq -c '.[]' | while IFS= read -r ep; do
        local ep_id service ep_type state
        ep_id=$(echo "$ep" | jq -r '.Id')
        service=$(echo "$ep" | jq -r '.Service')
        ep_type=$(echo "$ep" | jq -r '.Type')
        state=$(echo "$ep" | jq -r '.State')

        local color="$GREEN"
        [[ "$state" != "available" ]] && color="$YELLOW"

        printf "  %-26s %-40s %-12s ${color}%s${RESET}\n" "$ep_id" "${service:0:40}" "$ep_type" "$state"
    done
}

# ══════════════════════════════════════════════════════════════════════
#                       NETWORK ACLS
# ══════════════════════════════════════════════════════════════════════

audit_nacls() {
    log "Auditing network ACLs..."

    local nacl_filter=()
    [[ -n "$VPC_ID" ]] && nacl_filter=(--filters "Name=vpc-id,Values=${VPC_ID}")

    local nacls_json
    nacls_json=$(aws_cmd ec2 describe-network-acls \
        "${nacl_filter[@]}" \
        --output json 2>/dev/null)

    echo ""
    echo -e "  ${BOLD}Overly Permissive Network ACL Rules${RESET}"

    local found=0
    echo "$nacls_json" | jq -c '.NetworkAcls[]' | while IFS= read -r nacl; do
        local nacl_id
        nacl_id=$(echo "$nacl" | jq -r '.NetworkAclId')

        echo "$nacl" | jq -c '.Entries[]' | while IFS= read -r entry; do
            local cidr rule_action egress protocol rule_num
            cidr=$(echo "$entry" | jq -r '.CidrBlock // ""')
            rule_action=$(echo "$entry" | jq -r '.RuleAction')
            egress=$(echo "$entry" | jq -r '.Egress')
            protocol=$(echo "$entry" | jq -r '.Protocol')
            rule_num=$(echo "$entry" | jq -r '.RuleNumber')

            # Skip default deny-all rules (rule 32767)
            [[ "$rule_num" == "32767" ]] && continue

            if [[ "$cidr" == "0.0.0.0/0" && "$rule_action" == "allow" && "$protocol" == "-1" ]]; then
                local direction="ingress"
                [[ "$egress" == "true" ]] && direction="egress"

                echo -e "    ${YELLOW}!${RESET} ${nacl_id} — rule ${rule_num}: allow all ${direction} from 0.0.0.0/0"
                found=1
                ((WARN_COUNT++)) || true
            fi
        done
    done

    if [[ "$found" -eq 0 ]]; then
        echo -e "    ${GREEN}No overly permissive NACL rules found ✓${RESET}"
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                       SUMMARY
# ══════════════════════════════════════════════════════════════════════

print_summary() {
    echo ""
    echo -e "  ${BOLD}═══════════════════════════════════════${RESET}"
    echo -e "  ${BOLD}VPC Audit Summary${RESET}"
    echo -e "  ${BOLD}═══════════════════════════════════════${RESET}"
    echo -e "    ${GREEN}PASS:${RESET}      ${PASS_COUNT}"
    echo -e "    ${YELLOW}WARN:${RESET}      ${WARN_COUNT}"
    echo -e "    ${RED}CRITICAL:${RESET}  ${CRIT_COUNT}"
    echo "  ───────────────────────────────────────"

    local end_time
    end_time=$(date +%s)
    echo "    Completed in $(( end_time - START_TIME ))s"

    if [[ "$CRIT_COUNT" -gt 0 ]]; then
        echo ""
        echo -e "  ${RED}Action required: ${CRIT_COUNT} critical finding(s)${RESET}"
        exit 2
    elif [[ "$WARN_COUNT" -gt 0 ]]; then
        echo ""
        echo -e "  ${YELLOW}Review recommended: ${WARN_COUNT} warning(s)${RESET}"
        exit 1
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                       HELP
# ══════════════════════════════════════════════════════════════════════

show_help() {
    cat <<EOF
Usage: $SCRIPT_NAME [MODE] [OPTIONS]

Audit AWS VPC configuration and security.

MODES:
    --full                  Run all audit checks
    --security-groups       Audit security group rules
    --flow-logs             Check VPC flow log status
    --peering               List VPC peering connections
    --subnets               Check subnet IP utilization
    --endpoints             List VPC endpoints
    --nacls                 Audit network ACL rules

OPTIONS:
    --vpc-id VPC_ID         Scope audit to a specific VPC
    --format FORMAT         Output: text (default), json, csv
    --verbose               Debug output
    --no-color              Disable colored output
    --help, -h              Show this help

EXIT CODES:
    0   All checks passed
    1   Warnings found
    2   Critical findings

ENVIRONMENT VARIABLES:
    AWS_PROFILE             AWS CLI profile
    AWS_REGION              AWS region
    VPC_ID                  Target VPC ID
    OUTPUT_FORMAT           Output format (default: text)
    VERBOSE                 Debug output (default: false)
    COLOR                   Color mode: auto, always, never

EXAMPLES:
    # Full VPC audit
    ./$SCRIPT_NAME --full

    # Audit specific VPC
    ./$SCRIPT_NAME --full --vpc-id vpc-0abc123

    # Security groups only
    ./$SCRIPT_NAME --security-groups

    # Subnet utilization
    ./$SCRIPT_NAME --subnets
EOF
}

# ══════════════════════════════════════════════════════════════════════
#                       MAIN
# ══════════════════════════════════════════════════════════════════════

main() {
    local mode_set="false"

    while [[ $# -gt 0 ]]; do
        case "$1" in
            --full)             RUN_FULL="true"; mode_set="true"; shift ;;
            --security-groups)  RUN_SGS="true"; mode_set="true"; shift ;;
            --flow-logs)        RUN_FLOW="true"; mode_set="true"; shift ;;
            --peering)          RUN_PEER="true"; mode_set="true"; shift ;;
            --subnets)          RUN_SUBNETS="true"; mode_set="true"; shift ;;
            --endpoints)        RUN_ENDPOINTS="true"; mode_set="true"; shift ;;
            --nacls)            RUN_NACLS="true"; mode_set="true"; shift ;;
            --vpc-id)           VPC_ID="$2"; shift 2 ;;
            --format)           OUTPUT_FORMAT="$2"; shift 2 ;;
            --verbose)          VERBOSE="true"; shift ;;
            --no-color)         COLOR="never"; shift ;;
            --help|-h)          show_help; exit 0 ;;
            *)                  die "Unknown option: $1 (see --help)" ;;
        esac
    done

    setup_colors

    if [[ "$mode_set" == "false" ]]; then err "No mode specified"; echo ""; show_help; exit 1; fi

    START_TIME=$(date +%s)

    echo ""
    echo -e "${BOLD}VPC Audit${RESET}"
    echo "Region: ${AWS_REGION:-$(aws configure get region 2>/dev/null || echo 'default')}"
    [[ -n "$VPC_ID" ]] && echo "VPC:    ${VPC_ID}"
    echo "Time:   $(date -u +%Y-%m-%dT%H:%M:%SZ)"
    echo ""

    check_deps

    should_show "sgs"       && audit_security_groups
    should_show "flow"      && audit_flow_logs
    should_show "peer"      && audit_peering
    should_show "subnets"   && audit_subnets
    should_show "endpoints" && audit_endpoints
    should_show "nacls"     && audit_nacls

    print_summary
}

main "$@"