#!/usr/bin/env bash
set -Eeuo pipefail

LAB="Mission Tech Lab 02"
STATE_DIR="/run/mls1-lab02"
SCENARIO_FILE="$STATE_DIR/scenario"
LOG_DIR="$STATE_DIR/logs"

BR_CORE="mls1-core"
BR_EAST="mls1-east"
BR_WEST="mls1-west"
BRIDGES=("$BR_CORE" "$BR_EAST" "$BR_WEST")

NS_A="mls1-host-a"
NS_B="mls1-host-b"
NS_APP="mls1-app"
NAMESPACES=("$NS_A" "$NS_B" "$NS_APP")

TRUNK_ROOTS=(mls1-ce-c mls1-cw-c mls1-ew-e)

C_RESET='\033[0m'
C_BOLD='\033[1m'
C_PURPLE='\033[38;5;141m'
C_GREEN='\033[38;5;114m'
C_ORANGE='\033[38;5;208m'
C_RED='\033[38;5;203m'

log()  { printf '%b[%s]%b %s\n' "$C_PURPLE" "$LAB" "$C_RESET" "$*"; }
ok()   { printf '%b[PASS]%b %s\n' "$C_GREEN" "$C_RESET" "$*"; }
warn() { printf '%b[CHECK]%b %s\n' "$C_ORANGE" "$C_RESET" "$*"; }
fail() { printf '%b[FAIL]%b %s\n' "$C_RED" "$C_RESET" "$*" >&2; }
die()  { fail "$*"; exit 1; }

need_root() {
  [[ ${EUID:-$(id -u)} -eq 0 ]] || die "Run this command with sudo."
}

have() { command -v "$1" >/dev/null 2>&1; }

ensure_state() {
  mkdir -p "$LOG_DIR"
  chmod 700 "$STATE_DIR"
}

is_wsl2() {
  grep -qi microsoft /proc/sys/kernel/osrelease 2>/dev/null || grep -qi microsoft /proc/version 2>/dev/null
}

banner() {
  cat <<'EOF'

MISSION TECH
LAB 02: SWITCHING, BOOM!
Spanning Tree, VLANs, trunks, FDB learning, failover and the routing handoff

EOF
}

usage() {
  cat <<'EOF'
Usage:
  sudo ./mission-switching-lab-v2.sh <command>

Start here:
  install                 Install required Ubuntu packages
  doctor                  Check WSL2, sudo, tools and bridge support
  build                   Create the complete Layer 2 lab
  verify                  Run a short baseline verification

Learn and inspect:
  topology                Show the lab map
  physical                Show link state and counters
  stp                     Show root election and port states
  vlan                    Show access VLANs and trunks
  fdb                     Show learned MAC addresses
  checks                  Run Physical, STP, FDB and VLAN checks

Scenarios:
  scenario failover       Cut an active trunk and watch STP recover
  scenario vlan-outage    Remove VLAN 110 from an active trunk
  scenario mac-move       Move Host B from the West switch to East
  scenario routing        Prove Layer 2 works but routing is missing

Recovery:
  fix                     Repair the current scenario
  reset                   Rebuild the clean baseline
  destroy                 Remove every lab object
  help                    Show this help

This version intentionally uses the Linux kernel Spanning Tree implementation.
It does not depend on mstpd, systemd, a vendor image or a cloud account.
EOF
}

required=(ip bridge brctl tcpdump python3 ping curl tc)

install_packages() {
  need_root
  banner
  log "Installing the open-source lab tools"
  export DEBIAN_FRONTEND=noninteractive
  apt-get update
  apt-get install -y iproute2 bridge-utils tcpdump python3 curl iputils-ping procps
  ok "Packages installed"
  printf '\nNext step:\n  sudo %q doctor\n' "$(readlink -f "$0")"
}

bridge_capability_test() {
  local t="mls1-cap-test"
  ip link del "$t" >/dev/null 2>&1 || true
  if ip link add "$t" type bridge >/dev/null 2>&1; then
    ip link del "$t" >/dev/null 2>&1 || true
    return 0
  fi
  return 1
}

doctor() {
  need_root
  banner
  local failed=0
  printf '%bENVIRONMENT%b\n' "$C_BOLD" "$C_RESET"
  printf 'Distribution   %s\n' "$(. /etc/os-release 2>/dev/null; printf '%s %s' "${NAME:-Unknown}" "${VERSION_ID:-}")"
  printf 'Kernel         %s\n' "$(uname -r)"
  printf 'WSL2 detected  %s\n' "$(is_wsl2 && echo yes || echo no)"
  if ! is_wsl2; then fail "WSL2 was not detected"; failed=1; fi
  printf 'Script path    %s\n' "$(readlink -f "$0")"

  if [[ "$(readlink -f "$0")" == /mnt/* ]]; then
    warn "Move the script into your Linux home directory before continuing"
    printf '  cp %q ~/mission-switching-lab-v2.sh\n' "$(readlink -f "$0")"
    printf '  cd ~\n  chmod +x mission-switching-lab-v2.sh\n'
    failed=1
  fi

  printf '\n%bTOOLS%b\n' "$C_BOLD" "$C_RESET"
  local cmd
  for cmd in "${required[@]}"; do
    if have "$cmd"; then ok "$cmd"; else fail "$cmd is missing"; failed=1; fi
  done

  if bridge_capability_test; then
    ok "Linux bridge creation"
  else
    fail "Cannot create a Linux bridge"
    printf 'Confirm that you are inside WSL2 Ubuntu and that the command uses sudo.\n'
    failed=1
  fi

  if (( failed )); then
    printf '\nFix the checks above, then run doctor again.\n'
    return 1
  fi
  ok "Environment ready"
}

exists() { ip link show "$1" >/dev/null 2>&1; }

set_bridge() {
  local br="$1" mac="$2" priority="$3"
  ip link add "$br" type bridge
  ip link set dev "$br" address "$mac"
  ip link set dev "$br" type bridge \
    vlan_filtering 1 \
    stp_state 1 \
    priority "$priority" \
    forward_delay 4 \
    hello_time 1 \
    max_age 6 \
    ageing_time 300
  ip link set "$br" up
}

remove_vlan1() {
  bridge vlan del dev "$1" vid 1 >/dev/null 2>&1 || true
}

add_trunk() {
  local a="$1" b="$2" br_a="$3" br_b="$4"
  ip link add "$a" type veth peer name "$b"
  ip link set "$a" master "$br_a"
  ip link set "$b" master "$br_b"
  ip link set "$a" up
  ip link set "$b" up
  remove_vlan1 "$a"
  remove_vlan1 "$b"
  bridge vlan add dev "$a" vid 110
  bridge vlan add dev "$a" vid 120
  bridge vlan add dev "$b" vid 110
  bridge vlan add dev "$b" vid 120
}

add_host() {
  local ns="$1" ns_if="$2" sw_if="$3" br="$4" vlan_id="$5" addr="$6" mac="$7"
  ip netns add "$ns"
  ip link add "$sw_if" type veth peer name "$ns_if"
  ip link set "$ns_if" netns "$ns"
  ip link set "$sw_if" master "$br"
  ip link set "$sw_if" up
  remove_vlan1 "$sw_if"
  bridge vlan add dev "$sw_if" vid "$vlan_id" pvid untagged

  ip -n "$ns" link set lo up
  ip -n "$ns" link set "$ns_if" name eth0
  ip -n "$ns" link set eth0 address "$mac"
  ip -n "$ns" link set eth0 up
  ip -n "$ns" address add "$addr" dev eth0
  ip netns exec "$ns" sysctl -qw net.ipv6.conf.all.disable_ipv6=1
}

start_app() {
  ensure_state
  ip netns exec "$NS_APP" python3 -m http.server 8080 --bind 10.120.0.50 \
    >"$LOG_DIR/app.log" 2>&1 &
  echo $! > "$STATE_DIR/app.pid"
  sleep 0.5
}

kill_pidfile() {
  local f="$1" pid
  [[ -f "$f" ]] || return 0
  pid="$(cat "$f" 2>/dev/null || true)"
  if [[ -n "$pid" ]] && kill -0 "$pid" 2>/dev/null; then
    kill "$pid" 2>/dev/null || true
    wait "$pid" 2>/dev/null || true
  fi
  rm -f "$f"
}

topology() {
  cat <<'EOF'

                              mls1-core
                         mls1-ce-c     mls1-cw-c
                            /           \
                           /             \
                    mls1-ce-e               mls1-cw-w
                     mls1-east ----------- mls1-west
                         mls1-ew-e     mls1-ew-w
                        /                 |      \
                 Host A               Host B    App
                 VLAN 110             VLAN 110  VLAN 120
                 10.110.0.11          .12       10.120.0.50

All three inter-switch links carry VLAN 110 and VLAN 120.
Spanning Tree keeps one triangle port in blocking state.
The timers are shortened for the classroom, not for production use.
EOF
}

build() {
  need_root
  banner
  doctor >/dev/null || die "Doctor checks failed. Run: sudo $0 doctor"
  destroy --quiet
  ensure_state

  log "Creating three Linux switches"
  set_bridge "$BR_CORE" "02:00:00:00:00:01" 4096
  set_bridge "$BR_EAST" "02:00:00:00:00:02" 8192
  set_bridge "$BR_WEST" "02:00:00:00:00:03" 12288

  log "Connecting the triangle with tagged trunks"
  add_trunk mls1-ce-c mls1-ce-e "$BR_CORE" "$BR_EAST"
  add_trunk mls1-cw-c mls1-cw-w "$BR_CORE" "$BR_WEST"
  add_trunk mls1-ew-e mls1-ew-w "$BR_EAST" "$BR_WEST"

  log "Adding isolated endpoints"
  add_host "$NS_A" mls1-a-eth mls1-a-port "$BR_EAST" 110 "10.110.0.11/24" "02:00:00:00:11:0a"
  add_host "$NS_B" mls1-b-eth mls1-b-port "$BR_WEST" 110 "10.110.0.12/24" "02:00:00:00:11:0b"
  add_host "$NS_APP" mls1-app-eth mls1-app-port "$BR_WEST" 120 "10.120.0.50/24" "02:00:00:00:12:50"

  start_app
  sleep 8
  rm -f "$SCENARIO_FILE"
  topology
  ok "Lab built"
  printf '\nNext step:\n  sudo %q verify\n' "$(readlink -f "$0")"
}

blocked_count() {
  bridge link show | grep -E 'master mls1-(core|east|west)' | grep -Ec 'state (blocking|disabled)' || true
}

verify() {
  need_root
  local br
  for br in "${BRIDGES[@]}"; do exists "$br" || die "Lab is incomplete. Missing bridge: $br"; done
  [[ ! -f "$SCENARIO_FILE" ]] || { fail "A scenario is still active. Run fix before baseline verification."; return 1; }
  banner
  local failed=0 blocked

  blocked="$(blocked_count)"
  if (( blocked >= 1 )); then ok "Spanning Tree has a protected path"; else fail "No blocking port found"; failed=1; fi

  if ip netns exec "$NS_A" ping -c 2 -W 2 10.110.0.12 >/dev/null; then
    ok "Host A reaches Host B inside VLAN 110"
  else
    fail "Host A cannot reach Host B"
    failed=1
  fi

  if ip netns exec "$NS_A" ping -c 1 -W 1 10.120.0.50 >/dev/null 2>&1; then
    fail "Host A unexpectedly reached VLAN 120 without a router"
    failed=1
  else
    ok "VLAN 110 cannot reach VLAN 120 without routing"
  fi

  if (( failed )); then
    printf '\nRun these commands and compare them with the guide:\n'
    printf '  sudo %q stp\n  sudo %q vlan\n' "$(readlink -f "$0")" "$(readlink -f "$0")"
    return 1
  fi
  ok "Baseline verified"
}

physical() {
  need_root
  printf '\n%bPHYSICAL%b\n' "$C_BOLD" "$C_RESET"
  ip -br link show | grep -E '^mls1-' || true
  printf '\nCore-West counters:\n'
  ip -s link show mls1-cw-c | sed -n '1,9p'
}

stp() {
  need_root
  printf '\n%bSPANNING TREE%b\n' "$C_BOLD" "$C_RESET"
  printf '\nCompact port states:\n'
  bridge link show | grep -E 'master mls1-(core|east|west)' || true
  printf '\nRoot and timer evidence:\n'
  brctl showstp "$BR_CORE" | sed -n '1,18p'
  brctl showstp "$BR_EAST" | sed -n '1,18p'
  brctl showstp "$BR_WEST" | sed -n '1,18p'
}

vlan() {
  need_root
  printf '\n%bVLAN AND TRUNKS%b\n' "$C_BOLD" "$C_RESET"
  bridge vlan show
}

fdb() {
  need_root
  printf '\n%bFORWARDING DATABASE%b\n' "$C_BOLD" "$C_RESET"
  local br
  for br in "${BRIDGES[@]}"; do
    printf '\n--- %s ---\n' "$br"
    bridge fdb show br "$br" | grep -vE 'permanent|self' || echo '(no dynamic entries yet)'
  done
}

checks() {
  need_root
  exists "$BR_CORE" || die "Lab is not built."
  physical
  stp
  fdb
  vlan
}

record() {
  ensure_state
  printf '%s\n' "$1" > "$SCENARIO_FILE"
}

scenario_failover() {
  need_root
  exists mls1-cw-c || die "Lab is not built."
  record failover
  log "Starting a continuous A to B ping"
  ip netns exec "$NS_A" ping -D -i 0.5 -w 35 10.110.0.12 >"$LOG_DIR/failover-ping.log" 2>&1 &
  echo $! > "$STATE_DIR/ping.pid"
  sleep 2
  log "Cutting the active Core-West trunk"
  ip link set mls1-cw-c down
  printf '\nWatch the protected path move to forwarding:\n'
  printf '  watch -n 1 "bridge link show | grep -E '\''mls1-'\''"\n\n'
  printf 'Ping log:\n  sudo tail -f %s\n\n' "$LOG_DIR/failover-ping.log"
  printf 'Repair:\n  sudo %q fix\n' "$(readlink -f "$0")"
}

scenario_vlan_outage() {
  need_root
  exists mls1-cw-w || die "Lab is not built."
  record vlan-outage
  log "Removing VLAN 110 from the active West trunk"
  bridge vlan del dev mls1-cw-w vid 110
  printf '\nExpected result: physical and STP stay healthy, but A to B fails.\n\n'
  bridge vlan show dev mls1-cw-w
  ip netns exec "$NS_A" ping -c 2 -W 1 10.110.0.12 || true
  printf '\nRepair:\n  sudo %q fix\n' "$(readlink -f "$0")"
}

move_b_east() {
  bridge vlan del dev mls1-b-port vid 110 >/dev/null 2>&1 || true
  ip link set mls1-b-port nomaster
  ip link set mls1-b-port master "$BR_EAST"
  remove_vlan1 mls1-b-port
  bridge vlan add dev mls1-b-port vid 110 pvid untagged
  ip link set mls1-b-port up
}

move_b_west() {
  bridge vlan del dev mls1-b-port vid 110 >/dev/null 2>&1 || true
  ip link set mls1-b-port nomaster
  ip link set mls1-b-port master "$BR_WEST"
  remove_vlan1 mls1-b-port
  bridge vlan add dev mls1-b-port vid 110 pvid untagged
  ip link set mls1-b-port up
}

scenario_mac_move() {
  need_root
  exists mls1-b-port || die "Lab is not built."
  record mac-move
  log "Moving Host B from West to East without changing MAC or IP"
  move_b_east
  ip netns exec "$NS_B" ping -c 2 -W 1 10.110.0.11 >/dev/null 2>&1 || true
  fdb
  printf '\nRepair:\n  sudo %q fix\n' "$(readlink -f "$0")"
}

scenario_routing() {
  need_root
  exists "$BR_CORE" || die "Lab is not built."
  record routing
  printf '\n%bROUTING HANDOFF%b\n' "$C_BOLD" "$C_RESET"
  printf 'Source       10.110.0.11/24, VLAN 110\n'
  printf 'Destination  10.120.0.50/24, VLAN 120\n\n'
  printf 'Host A route table:\n'
  ip netns exec "$NS_A" ip route
  printf '\nHTTP attempt:\n'
  ip netns exec "$NS_A" curl -v --connect-timeout 2 http://10.120.0.50:8080 2>&1 || true
  printf '\nPhysical: PASS\nSTP:      PASS\nFDB:      PASS\nVLAN:     PASS\n'
  printf 'Router:    NOT PRESENT\nGateway:   NOT CONFIGURED\n\n'
  printf '%bACT II: FROM YOUR LAN TO AN APPLICATION%b\n' "$C_PURPLE" "$C_RESET"
}

fix() {
  need_root
  local scenario
  scenario="$(cat "$SCENARIO_FILE" 2>/dev/null || true)"
  if [[ -z "$scenario" ]]; then
    warn "No recorded scenario. Applying safe baseline repairs."
  fi
  case "$scenario" in
    failover)
      log "Restoring Core-West"
      ip link set mls1-cw-c up
      kill_pidfile "$STATE_DIR/ping.pid"
      ;;
    vlan-outage)
      log "Restoring VLAN 110 on mls1-cw-w"
      bridge vlan add dev mls1-cw-w vid 110
      ;;
    mac-move)
      log "Moving Host B back to West"
      move_b_west
      ;;
    routing)
      warn "Routing is intentionally missing. No Layer 2 fix applies."
      ;;
    '') ;;
    *) warn "Unknown scenario: $scenario" ;;
  esac
  if exists "$BR_CORE"; then
    ip link set mls1-cw-c up >/dev/null 2>&1 || true
    ip link set mls1-cw-w up >/dev/null 2>&1 || true
    bridge vlan add dev mls1-cw-w vid 110 >/dev/null 2>&1 || true
    move_b_west >/dev/null 2>&1 || true
  fi
  kill_pidfile "$STATE_DIR/ping.pid"
  rm -f "$SCENARIO_FILE"
  sleep 8
  ok "Baseline restored"
}

destroy() {
  need_root
  local quiet="${1:-}"
  [[ "$quiet" == "--quiet" ]] || log "Removing the lab"
  kill_pidfile "$STATE_DIR/app.pid"
  kill_pidfile "$STATE_DIR/ping.pid"

  local ns v br
  for ns in "${NAMESPACES[@]}"; do ip netns del "$ns" >/dev/null 2>&1 || true; done
  for v in "${TRUNK_ROOTS[@]}" mls1-a-port mls1-b-port mls1-app-port; do ip link del "$v" >/dev/null 2>&1 || true; done
  for br in "${BRIDGES[@]}"; do ip link del "$br" >/dev/null 2>&1 || true; done
  rm -rf "$STATE_DIR"
  [[ "$quiet" == "--quiet" ]] || ok "All lab objects removed"
}

reset() {
  need_root
  destroy --quiet
  build
}

main() {
  case "${1:-help}" in
    install) install_packages ;;
    doctor) doctor ;;
    build) build ;;
    verify) verify ;;
    topology) topology ;;
    physical) physical ;;
    stp) stp ;;
    vlan) vlan ;;
    fdb) fdb ;;
    checks) checks ;;
    scenario)
      case "${2:-}" in
        failover) scenario_failover ;;
        vlan-outage) scenario_vlan_outage ;;
        mac-move) scenario_mac_move ;;
        routing) scenario_routing ;;
        *) die "Use: failover, vlan-outage, mac-move or routing" ;;
      esac
      ;;
    fix) fix ;;
    reset) reset ;;
    destroy) destroy ;;
    help|-h|--help) usage ;;
    *) usage; exit 1 ;;
  esac
}

main "$@"