#!/bin/bash
#
# Purpose : Show the network ports open on this machine: who can reach each one
#           through the firewall, which service opened it and why, and let you
#           stop or start those services.
#
# A service you stop here is also kept from starting at boot, and update leaves
# it off (see enable_services in utils.sh) until you start it here again.

# Wrapped in { ...; exit; } so bash reads the whole script up front; update
# replaces this file while it runs.
{

COREDESKTOP="$HOME/.local/share/coredesktop"
source "$COREDESKTOP/utils.sh"

# Why a service or app opens ports, by systemd unit or process name. Services
# with ports here are also listed while they're stopped, so they can be started.
# Format: name|ports|why
KNOWN=(
  "sshd.service|22/tcp|SSH: log in to this machine from other machines"
  "systemd-resolved.service||Looks up domain names (DNS); 5355 is LLMNR, which answers name lookups from the LAN"
  "avahi-daemon.service|5353/udp|mDNS: lets machines on your LAN reach this one as $(uname -n).local"
  "cups.service|631/tcp|Printing (CUPS); only this machine can reach it"
  "cups.socket|631/tcp|Printing (CUPS); only this machine can reach it"
  "tailscaled.service|41641/udp|Tailscale: your Core's private network"
  "docker.service||Docker"
  "sonarr.service|8989/tcp|Sonarr: finds and downloads TV shows; its web page is on port 8989"
  "hermes-gateway.service||The sandbox's proxy for Hermes Agent's gateway (srt); only this machine can reach it"
  "gvfs-daemon.service||Files: finds Windows computers and shares on your network (WS-Discovery)"
  "kdeconnectd||KDE Connect (GSConnect): links your phone to this computer"
  "localsend||LocalSend: sends files to and from nearby devices"
  "com.vscodium.codium||VSCodium: its extensions' local servers; only this machine can reach them"
)

# Stopping these would break the desktop, Files or the network, so audit
# doesn't offer to
PROTECTED="gvfs-daemon.service systemd-resolved.service NetworkManager.service gdm.service dbus.service dbus-broker.service systemd-networkd.service"

usage() {
  cat << EOF
Usage: audit [--list]

  audit          Show the open ports and choose a service to stop or start
  audit --list   Only show the open ports

Run it with sudo rights: without them, the services run by root can't be named.
EOF
}

known_field() { # <name> <field: 2=ports 3=why>
  local entry
  for entry in "${KNOWN[@]}"; do
    [ "${entry%%|*}" = "$1" ] && { cut -d'|' -f"$2" <<< "$entry"; return 0; }
  done
  return 1
}

# The running known service that uses <proto> <port>, for a socket whose
# process only root can see
known_port_owner() {
  local entry unit
  for entry in "${KNOWN[@]}"; do
    unit="${entry%%|*}"
    [[ ",$(cut -d'|' -f2 <<< "$entry")," == *",$2/$1,"* ]] || continue
    systemctl is-active --quiet "$unit" 2> /dev/null && { echo "system $unit"; return 0; }
  done
  return 1
}

# Firewall rules that let connections in, as "proto port iface" ("any" for any
# protocol, port or interface), from ufw's rule files
load_firewall() {
  FIREWALL_ON=0
  FIREWALL_RULES=()
  grep -qs '^ENABLED=yes' /etc/ufw/ufw.conf && systemctl is-active --quiet ufw && FIREWALL_ON=1
  local action proto dport iface
  while read -r _ _ _ action proto dport _ _ _ iface _; do
    [[ "$action" == allow || "$action" == limit ]] || continue
    iface="${iface#in}"
    FIREWALL_RULES+=("$proto $dport ${iface#_}")
  done < <(grep -hs '^### tuple ###' /etc/ufw/user.rules /etc/ufw/user6.rules)
}

# Whether the firewall lets <proto> <port> in on <iface> ("" for every interface)
firewall_allows() {
  local proto="$1" port="$2" want="$3" rule rproto rport riface part
  [ "$FIREWALL_ON" = 0 ] && return 0
  for rule in "${FIREWALL_RULES[@]}"; do
    read -r rproto rport riface <<< "$rule"
    [[ "$rproto" == any || "$rproto" == "$proto" ]] || continue
    [[ -z "$riface" || "$riface" == "$want" ]] || continue
    [ "$rport" = any ] && return 0
    for part in ${rport//,/ }; do
      if [[ "$part" == *:* ]]; then
        (( port >= ${part%:*} && port <= ${part#*:} )) && return 0
      elif [ "$part" = "$port" ]; then
        return 0
      fi
    done
  done
  return 1
}

# Which network an interface is on: local, lan, tailscale or containers
iface_kind() {
  case "$1" in
    lo) echo local ;;
    tailscale*) echo tailscale ;;
    docker*|br-*|veth*|virbr*) echo containers ;;
    *) echo lan ;;
  esac
}

# The systemd unit a process runs in, as "system <unit>", "user <unit>",
# "app <name>" or "none <name>"
process_owner() {
  local pid="$1" comm="$2" cgroup unit
  cgroup="$(cut -d: -f3- "/proc/$pid/cgroup" 2> /dev/null | head -n 1)"
  unit="${cgroup##*/}"
  case "$cgroup" in
    */app-flatpak-*.scope)
      unit="${unit#app-flatpak-}"
      echo "app ${unit%-*.scope}" ;;
    *.scope) echo "app $comm" ;;
    /system.slice/*.service) echo "system $unit" ;;
    /user.slice/*/user@*.service/*.service) echo "user $unit" ;;
    *) echo "none $comm" ;;
  esac
}

# A port systemd itself listens on, for a service it starts on demand
socket_unit() { # <port>
  systemctl list-sockets --all --no-legend 2> /dev/null \
    | awk -v port="$1" '$1 ~ ":" port "$" { print $2; exit }'
}

# Collect the listening sockets into ROWS, one per service and port, with the
# networks each listens on
load_ports() {
  local -A iface_of seen
  local netid local_addr process addr port iface kind pid comm owner key
  ROW_KEYS=()
  declare -gA ROW_PROTO ROW_PORT ROW_OWNER ROW_KINDS
  ROW_PROTO=() ROW_PORT=() ROW_OWNER=() ROW_KINDS=()

  while read -r _ iface _ addr _; do
    iface_of["${addr%/*}"]="$iface"
  done < <(ip -o addr show 2> /dev/null)

  while read -r netid _ _ _ local_addr _ process; do
    port="${local_addr##*:}"
    addr="${local_addr%:*}"
    addr="${addr#[}" addr="${addr%]}"
    iface=""
    [[ "$addr" == *%* ]] && iface="${addr#*%}" addr="${addr%%%*}"
    case "$addr" in
      0.0.0.0|::|\*) kind=all ;;
      127.*|::1) kind=local ;;
      *) kind="$(iface_kind "${iface:-${iface_of[$addr]:-lan}}")" ;;
    esac

    pid="$(grep -o 'pid=[0-9]*' <<< "$process" | head -n 1 | cut -d= -f2)"
    comm="$(sed -nE 's/^users:\(\("([^"]*)".*/\1/p' <<< "$process")"
    if [ -z "$pid" ]; then
      owner="$(known_port_owner "$netid" "$port")" || owner="none ?"
    elif [ "$comm" = systemd ] && unit="$(socket_unit "$port")" && [ -n "$unit" ]; then
      owner="system $unit"
    elif [ "$comm" = docker-proxy ]; then
      owner="none docker-proxy"
    else
      owner="$(process_owner "$pid" "$comm")"
    fi

    # Random high ports bound to one address are a service's replies and
    # discovery, not something to connect to: one row for all of them
    [[ "$kind" != all && "$port" -ge 32768 ]] && port=random
    key="$netid $port $owner"
    if [ -z "${seen[$key]}" ]; then
      seen[$key]=1
      ROW_KEYS+=("$key")
      ROW_PROTO[$key]="$netid" ROW_PORT[$key]="$port" ROW_OWNER[$key]="$owner"
    fi
    [[ " ${ROW_KINDS[$key]} " == *" $kind "* ]] || ROW_KINDS[$key]+=" $kind"
  done < <($SS -tulpnH 2> /dev/null)
}

# Who can connect to a row's port, given where it listens and the firewall
reach_label() {
  local proto="$1" port="$2" kinds="$3" lan=0 ts=0 ct=0 blocked=0 parts=()
  [ "$port" = random ] && port=0
  for kind in $kinds; do
    case "$kind" in
      all) lan=1 ts=1 ;;
      lan) lan=1 ;;
      tailscale) ts=1 ;;
      containers) ct=1 ;;
    esac
  done
  if [ "$lan" = 1 ]; then
    if firewall_allows "$proto" "$port" ""; then parts+=("\e[31mLAN\e[0m"); else blocked=1; fi
  fi
  if [ "$ts" = 1 ] && firewall_allows "$proto" "$port" tailscale0; then
    parts+=("\e[33mTailscale\e[0m")
  fi
  if [ "$ct" = 1 ] && firewall_allows "$proto" "$port" ""; then
    parts+=("\e[33mcontainers\e[0m")
  fi
  [ ${#parts[@]} -eq 0 ] && parts=("\e[32mthis machine\e[0m")
  local label="${parts[0]}" part
  for part in "${parts[@]:1}"; do label+=", $part"; done
  [ "$blocked" = 1 ] && label+=" (firewall blocks LAN)"
  echo -n "$label"
}

owner_label() {
  local scope="${1%% *}" name="${1#* }"
  case "$scope" in
    system) echo "${name%.service}" ;;
    user) echo "${name%.service} (yours)" ;;
    app) echo "$name (app)" ;;
    *) [ "$name" = docker-proxy ] && echo "Docker container" || echo "$name" ;;
  esac
}

why() {
  local scope="${1%% *}" name="${1#* }" text
  text="$(known_field "$name" 3)" && { echo "$text"; return; }
  case "$scope" in
    system) systemctl show -p Description --value "$name" 2> /dev/null ;;
    user) systemctl --user show -p Description --value "$name" 2> /dev/null ;;
    none)
      case "$name" in
        "?") echo "Run audit with sudo rights to see which service this is" ;;
        docker-proxy) echo "A port a Docker container publishes; see: docker ps" ;;
      esac ;;
  esac
}

systemctl_for() { # <scope> <arguments>
  if [ "$1" = user ]; then systemctl --user "${@:2}"; else sudo systemctl "${@:2}"; fi
}

is_active() { # <scope> <unit>
  if [ "$1" = user ]; then systemctl --user is-active --quiet "$2"; else systemctl is-active --quiet "$2"; fi
}

# Services that can be stopped or started here, as "scope unit": those with
# open ports, then the known ones that are installed
load_services() {
  local key owner entry unit
  SERVICES_LIST=()
  for key in "${ROW_KEYS[@]}"; do
    owner="${ROW_OWNER[$key]}"
    [[ "$owner" == system\ * || "$owner" == user\ * ]] || continue
    [[ " $PROTECTED " == *" ${owner#* } "* ]] && continue
    [[ " ${SERVICES_LIST[*]} " == *" $owner "* ]] || SERVICES_LIST+=("$owner")
  done
  for entry in "${KNOWN[@]}"; do
    unit="${entry%%|*}"
    [[ "$unit" == *.service || "$unit" == *.socket ]] || continue
    [ -n "$(known_field "$unit" 2)" ] || continue
    [[ " ${SERVICES_LIST[*]} " == *" system $unit "* ]] && continue
    # Listed even while running without its port open yet, as right after
    # starting, so it doesn't seem to vanish
    systemctl cat "$unit" &> /dev/null || continue
    # cups.socket and cups.service are one service to the user
    [[ "$unit" == *.socket && " ${SERVICES_LIST[*]} " == *" system ${unit%.socket}.service "* ]] && continue
    [[ "$unit" == *.service && " ${SERVICES_LIST[*]} " == *" system ${unit%.service}.socket "* ]] && continue
    SERVICES_LIST+=("system $unit")
  done
}

show_ports() {
  local key reach
  echo
  if [ "$FIREWALL_ON" = 1 ]; then
    echo "Firewall (ufw) on: it blocks connections from other machines except to the ports it allows."
  else
    echo -e "\e[31mFirewall (ufw) off: every port listening beyond this machine can be reached.\e[0m"
  fi
  [ -n "$NO_SUDO" ] && echo "Without sudo rights, the services run by root show as \"?\"."
  echo
  printf "\e[1m%-7s %-5s %-34s %-26s %s\e[0m\n" PORT PROTO "OPEN TO" SERVICE WHY
  local sorted visible
  # By port, with the random ports last
  mapfile -t sorted < <(for key in "${ROW_KEYS[@]}"; do
      printf '%s\t%s\n' "${ROW_PORT[$key]/random/99999}" "$key"; done | sort -n -k1,1 | cut -f2-)
  for key in "${sorted[@]}"; do
    reach="$(reach_label "${ROW_PROTO[$key]}" "${ROW_PORT[$key]}" "${ROW_KINDS[$key]}")"
    # printf counts the color codes as width, so pad by the visible length
    visible="$(sed 's/\\e\[[0-9;]*m//g' <<< "$reach")"
    printf "%-7s %-5s %b%*s %-26s %s\n" "${ROW_PORT[$key]}" "${ROW_PROTO[$key]}" \
      "$reach" $((34 - ${#visible})) "" "$(owner_label "${ROW_OWNER[$key]}")" "$(why "${ROW_OWNER[$key]}")"
  done
}

show_services() {
  local i scope unit state
  echo
  echo "Services you can stop or start:"
  echo
  for i in "${!SERVICES_LIST[@]}"; do
    scope="${SERVICES_LIST[$i]%% *}" unit="${SERVICES_LIST[$i]#* }"
    if is_active "$scope" "$unit"; then state="\e[32mrunning\e[0m"; else state="\e[33mstopped\e[0m"; fi
    printf "  %2d) %-30b %b\n" $((i + 1)) "$(owner_label "${SERVICES_LIST[$i]}")" "$state"
  done
}

OFF_FILE="$HOME/.config/coredesktop/services-off"

toggle() { # <scope> <unit>
  local scope="$1" unit="$2" units=("$2") pair
  if [[ " $PROTECTED " == *" $unit "* ]]; then
    echo "$unit keeps the desktop or the network working, so audit won't stop it."
    return 1
  fi
  if [ "$unit" = sshd.service ] && [ -n "$SSH_CONNECTION" ] && is_active "$scope" "$unit"; then
    echo "You're logged in over SSH: stopping sshd would lock you out."
    return 1
  fi
  # A service started on demand has a socket unit too; both go together
  case "$unit" in
    *.socket) pair="${unit%.socket}.service" ;;
    *.service) pair="${unit%.service}.socket" ;;
  esac
  systemctl_for "$scope" cat "$pair" &> /dev/null && units+=("$pair")

  mkdir -p "$(dirname "$OFF_FILE")"
  if is_active "$scope" "$unit"; then
    confirm "Stop ${units[*]} and keep it from starting at boot?" || return 0
    systemctl_for "$scope" disable --now "${units[@]}" || return 1
    for unit in "${units[@]}"; do
      grep -qsx "$scope $unit" "$OFF_FILE" || echo "$scope $unit" >> "$OFF_FILE"
    done
  else
    confirm "Start ${units[*]} and start it at boot?" || return 0
    systemctl_for "$scope" enable --now "${units[@]}" || return 1
    for unit in "${units[@]}"; do
      sed -i "\|^$scope $unit\$|d" "$OFF_FILE" 2> /dev/null
    done
    wait_for_ports "$2"
  fi
}

# A service takes a moment to open its ports once started: wait up to 10
# seconds for its known ones, so the table shows them
wait_for_ports() { # <unit>
  local ports port proto
  ports="$(known_field "$1" 2)" || return 0
  for port in ${ports//,/ }; do
    proto="${port#*/}" port="${port%/*}"
    echo "Waiting for $1 to open port $port..."
    for _ in $(seq 40); do
      [ -n "$(ss -lnH "--$proto" "sport = :$port" 2> /dev/null)" ] && break
      sleep 0.25
    done
  done
}

confirm() {
  local answer
  read -rp "$1 [y/N] " answer
  [[ "$answer" == [yY]* ]]
}

case "$1" in
  ""|-l|--list) ;;
  -h|--help) usage; exit 0 ;;
  *) usage; exit 2 ;;
esac

# Root's sockets name their process only to root
SS="sudo ss"
NO_SUDO=""
if ! sudo -v; then
  SS="ss"
  NO_SUDO=1
fi

while true; do
  load_firewall
  load_ports
  show_ports
  [ -n "$1" ] && exit 0

  load_services
  show_services
  echo
  echo "   q) Quit"
  echo
  while true; do
    read -rp "Choose a service to stop or start: " choice || exit 0
    case "$choice" in
      q|Q|"") exit 0 ;;
      *[!0-9]*) ;;
      *)
        if (( choice >= 1 && choice <= ${#SERVICES_LIST[@]} )); then
          entry="${SERVICES_LIST[$((choice - 1))]}"
          toggle "${entry%% *}" "${entry#* }"
          break
        fi ;;
    esac
    echo "Choose a number from the list, or q."
  done
done
}
