fleet-ssh

#!/bin/bash
# ═══════════════════════════════════════════════════════════════
# fleet-ssh — Mac Fleet Control Tool
#
# Usage:
#   fleet-ssh list                    Show all machines with status
#   fleet-ssh <#> "<command>"         Run command on machine by number
#   fleet-ssh <name> "<command>"      Run command on machine by name
#   fleet-ssh all "<command>"         Run command on ALL online machines
#   fleet-ssh ping                    Quick ping all machines
#   fleet-ssh add <name> <user> <ip>  Manually add a machine
#   fleet-ssh remove <name>           Remove a machine
#   fleet-ssh shell <#|name>          Interactive SSH session
#
# Machine registry: ~/.fleet-machines.json
# ═══════════════════════════════════════════════════════════════

FLEET_FILE="${FLEET_FILE:-$HOME/.fleet-machines.json}"
SSH_OPTS="-o ConnectTimeout=5 -o StrictHostKeyChecking=no -o LogLevel=ERROR"
# Prefix commands with PATH setup so homebrew tools are found on remote
REMOTE_PREFIX="export PATH=/opt/homebrew/bin:/usr/local/bin:\$PATH;"

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
DIM='\033[2m'
BOLD='\033[1m'
NC='\033[0m'

# ── Ensure fleet file exists ──
init_fleet() {
  if [ ! -f "$FLEET_FILE" ]; then
    echo '{"machines":[]}' > "$FLEET_FILE"
  fi
}

# ── Get machine count ──
machine_count() {
  python3 -c "
import json
with open('$FLEET_FILE') as f: d=json.load(f)
print(len(d.get('machines',[])))
"
}

# ── List machines ──
list_machines() {
  init_fleet
  local count=$(machine_count)
  if [ "$count" = "0" ]; then
    echo -e "${YELLOW}No machines registered.${NC}"
    echo -e "Add one: ${CYAN}fleet-ssh add <name> <user> <ip>${NC}"
    echo -e "Or run ${CYAN}worker-setup.sh${NC} on a remote Mac."
    return
  fi
  
  echo ""
  echo -e "${BOLD}  #   Name                    User          IP                Status${NC}"
  echo -e "  ─── ─────────────────────── ───────────── ───────────────── ──────"
  
  python3 -c "
import json, subprocess, sys

with open('$FLEET_FILE') as f:
    d = json.load(f)

for i, m in enumerate(d.get('machines', []), 1):
    name = m.get('name', '?')
    user = m.get('user', '?')
    ip = m.get('ip', '?')
    
    # Quick SSH check
    try:
        r = subprocess.run(
            ['ssh', '-o', 'ConnectTimeout=3', '-o', 'StrictHostKeyChecking=no', '-o', 'LogLevel=ERROR',
             f'{user}@{ip}', 'echo ok'],
            capture_output=True, text=True, timeout=5
        )
        status = '\033[0;32m● online\033[0m' if r.returncode == 0 else '\033[0;31m● offline\033[0m'
    except:
        status = '\033[0;31m● timeout\033[0m'
    
    print(f'  {i:<3} {name:<23} {user:<13} {ip:<17} {status}')
    print(f'      \033[2mVNC: open vnc://{user}@{ip}\033[0m')
" 2>/dev/null
  echo ""

  # ── Visual topology ──
  local master_name=$(hostname | sed 's/\.local$//')
  local master_ip
  master_ip=$(/Applications/Tailscale.app/Contents/MacOS/Tailscale ip -4 2>/dev/null || tailscale ip -4 2>/dev/null || echo "100.x.x.x")

  echo -e "  ${CYAN}┌─────────────────────┐${NC}"
  echo -e "  ${CYAN}│${NC} ${BOLD}🖥  MASTER${NC}           ${CYAN}│${NC}"
  printf "  ${CYAN}│${NC} %-20s${CYAN}│${NC}\n" "$master_name"
  printf "  ${CYAN}│${NC} ${DIM}%-20s${NC}${CYAN}│${NC}\n" "$master_ip"
  echo -e "  ${CYAN}└──────────┬──────────┘${NC}"
  echo -e "             ${CYAN}│${NC} ${DIM}Tailscale E2EE${NC}"

  # Read machines and draw
  python3 -c "
import json
with open('$FLEET_FILE') as f:
    machines = json.load(f).get('machines', [])

n = len(machines)
if n == 0:
    exit()

C = '\033[0;36m'
G = '\033[0;32m'
W = '\033[1;37m'
B = '\033[1m'
D = '\033[2m'
NC = '\033[0m'

# Collect status results (reuse from earlier check via simple display)
for i, m in enumerate(machines):
    name = m.get('name', '?')[:18]
    ip = m.get('ip', '?')
    connector = '├' if i < n - 1 else '└'
    print(f'             {C}{connector}── {G}🖥  {NC}{B}{name}{NC} {D}({ip}){NC}')

print()
" 2>/dev/null

  echo -e "  ${DIM}github.com/celestwong0920/mac-fleet-control${NC}"
  echo ""
}

# ── Get machine by number or name ──
resolve_machine() {
  local target="$1"
  python3 -c "
import json, sys

with open('$FLEET_FILE') as f:
    d = json.load(f)

machines = d.get('machines', [])
target = '$target'

# Try by number
try:
    idx = int(target) - 1
    if 0 <= idx < len(machines):
        m = machines[idx]
        print(f\"{m['user']}@{m['ip']}|{m['name']}\")
        sys.exit(0)
except ValueError:
    pass

# Try by name (partial match)
for m in machines:
    if target.lower() in m.get('name', '').lower():
        print(f\"{m['user']}@{m['ip']}|{m['name']}\")
        sys.exit(0)

# Try by IP
for m in machines:
    if target == m.get('ip', ''):
        print(f\"{m['user']}@{m['ip']}|{m['name']}\")
        sys.exit(0)

sys.exit(1)
" 2>/dev/null
}

# ── Add machine ──
add_machine() {
  local name="$1" user="$2" ip="$3"
  [ -z "$name" ] || [ -z "$user" ] || [ -z "$ip" ] && {
    echo "Usage: fleet-ssh add <name> <user> <ip>"
    exit 1
  }
  init_fleet
  python3 -c "
import json, datetime

with open('$FLEET_FILE') as f:
    d = json.load(f)

machines = d.get('machines', [])

# Remove existing with same name
machines = [m for m in machines if m.get('name') != '$name']

machines.append({
    'name': '$name',
    'user': '$user',
    'ip': '$ip',
    'added': datetime.datetime.now().isoformat()
})

d['machines'] = machines

with open('$FLEET_FILE', 'w') as f:
    json.dump(d, f, indent=2)

print(f'Added: $name ($user@$ip)')
"
}

# ── Remove machine ──
remove_machine() {
  local name="$1"
  [ -z "$name" ] && { echo "Usage: fleet-ssh remove <name>"; exit 1; }
  python3 -c "
import json

with open('$FLEET_FILE') as f:
    d = json.load(f)

before = len(d.get('machines', []))
d['machines'] = [m for m in d.get('machines', []) if m.get('name') != '$name']
after = len(d['machines'])

with open('$FLEET_FILE', 'w') as f:
    json.dump(d, f, indent=2)

if before > after:
    print(f'Removed: $name')
else:
    print(f'Not found: $name')
"
}

# ── Run on all ──
run_all() {
  local cmd="$1"
  init_fleet
  echo -e "${BOLD}Running on all machines:${NC} $cmd"
  echo ""
  
  python3 -c "
import json
with open('$FLEET_FILE') as f:
    d = json.load(f)
for m in d.get('machines', []):
    print(f\"{m['user']}@{m['ip']}|{m['name']}\")
" 2>/dev/null | while IFS='|' read -r target name; do
    echo -e "${YELLOW}── $name ($target) ──${NC}"
    if ssh $SSH_OPTS "$target" "$REMOTE_PREFIX $cmd" 2>&1; then
      echo -e "${GREEN}[✓]${NC} $name"
    else
      echo -e "${RED}[✗]${NC} $name"
    fi
    echo ""
  done
}

# ── Show master info (paste-able worker setup command) ──
show_master() {
  local ts_cli=""
  for p in tailscale /opt/homebrew/bin/tailscale /usr/local/bin/tailscale "/Applications/Tailscale.app/Contents/MacOS/Tailscale"; do
    if command -v "$p" &>/dev/null || [ -x "$p" ]; then
      "$p" version &>/dev/null && { ts_cli="$p"; break; }
    fi
  done

  if [ -z "$ts_cli" ]; then
    echo -e "${RED}Tailscale not found on this machine.${NC}"
    exit 1
  fi

  local ts_ip
  ts_ip=$("$ts_cli" ip -4 2>/dev/null | head -1)
  local me
  me=$(whoami)
  local hn
  hn=$("$ts_cli" status --json 2>/dev/null | python3 -c "import sys,json; print(json.load(sys.stdin).get('Self',{}).get('HostName',''))" 2>/dev/null || hostname -s)

  if [ -z "$ts_ip" ]; then
    echo -e "${RED}Tailscale not connected.${NC} Open the app and log in."
    exit 1
  fi

  echo ""
  echo -e "${BOLD}This master:${NC}"
  echo -e "  Hostname:     ${CYAN}$hn${NC}"
  echo -e "  Tailscale IP: ${CYAN}$ts_ip${NC}"
  echo -e "  Username:     ${CYAN}$me${NC}"
  echo ""
  echo -e "${BOLD}Run this on each new worker Mac (copy-paste):${NC}"
  echo ""
  echo -e "${CYAN}git clone https://github.com/willau95/mac-fleet-control.git ~/mac-fleet-control \\"
  echo -e "  && cd ~/mac-fleet-control \\"
  echo -e "  && bash worker-setup.sh --master $me@$ts_ip${NC}"
  echo ""
  echo -e "${DIM}(Will prompt once for the master password — afterwards passwordless.)${NC}"
  echo ""
}

# ── Ping all ──
ping_all() {
  init_fleet
  echo -e "${BOLD}Pinging fleet...${NC}"
  echo ""
  
  python3 -c "
import json
with open('$FLEET_FILE') as f:
    d = json.load(f)
for m in d.get('machines', []):
    print(f\"{m['user']}@{m['ip']}|{m['name']}\")
" 2>/dev/null | while IFS='|' read -r target name; do
    START=$(python3 -c "import time; print(int(time.time()*1000))")
    if ssh $SSH_OPTS "$target" "echo ok" &>/dev/null; then
      END=$(python3 -c "import time; print(int(time.time()*1000))")
      MS=$((END - START))
      echo -e "  ${GREEN}●${NC} $name — ${MS}ms"
    else
      echo -e "  ${RED}●${NC} $name — unreachable"
    fi
  done
  echo ""
}

# ── Main ──
case "${1:-}" in
  ""|help|-h|--help)
    echo ""
    echo -e "${BOLD}fleet-ssh — Mac Fleet Control${NC}"
    echo ""
    echo "  fleet-ssh list                    Show all machines"
    echo "  fleet-ssh master                  Print paste-able worker setup command"
    echo "  fleet-ssh <#> \"<command>\"         Run by number"
    echo "  fleet-ssh <name> \"<command>\"      Run by name"
    echo "  fleet-ssh all \"<command>\"         Run on ALL machines"
    echo "  fleet-ssh ping                    Ping all"
    echo "  fleet-ssh shell <#|name>          Interactive SSH"
    echo "  fleet-ssh add <name> <user> <ip>  Add machine"
    echo "  fleet-ssh remove <name>           Remove machine"
    echo ""
    ;;
  list) list_machines ;;
  master|info|whoami) show_master ;;
  ping) ping_all ;;
  add) add_machine "$2" "$3" "$4" ;;
  remove) remove_machine "$2" ;;
  all) run_all "$2" ;;
  shell)
    RESOLVED=$(resolve_machine "$2")
    if [ $? -ne 0 ]; then
      echo -e "${RED}Machine not found: $2${NC}"
      echo "Run 'fleet-ssh list' to see available machines."
      exit 1
    fi
    TARGET=$(echo "$RESOLVED" | cut -d'|' -f1)
    NAME=$(echo "$RESOLVED" | cut -d'|' -f2)
    echo -e "${CYAN}Connecting to $NAME ($TARGET)...${NC}"
    ssh $SSH_OPTS "$TARGET"
    ;;
  *)
    # Resolve target + run command
    if [ -z "$2" ]; then
      echo "Usage: fleet-ssh <target> \"<command>\""
      echo "Run 'fleet-ssh list' to see available machines."
      exit 1
    fi
    RESOLVED=$(resolve_machine "$1")
    if [ $? -ne 0 ]; then
      echo -e "${RED}Machine not found: $1${NC}"
      echo "Run 'fleet-ssh list' to see available machines."
      exit 1
    fi
    TARGET=$(echo "$RESOLVED" | cut -d'|' -f1)
    NAME=$(echo "$RESOLVED" | cut -d'|' -f2)
    echo -e "${DIM}[$NAME]${NC}"
    ssh $SSH_OPTS "$TARGET" "$REMOTE_PREFIX $2"
    ;;
esac