mnf

#!/usr/bin/env bash

# ---------------------------
# CONFIG VARIABLES YOU SHOULD CHANGE
RULES_DIR="$HOME/Software/mullvad-tailscale/" # Path to the dir in which mullvad.rules file is located.
EXCLUDE_COUNTRY_CODES=(us ca jp au hk gb) # Country codes to avoid. Set to '' to allow all.
# INCLUDE_COUNTRY_CODES=(fr gb) # Country codes to connect to. Uncomment to use. Overridden by -c option. Overrides EXCLUDE_COUNTRY_CODES.
# DNS=(1.1.1.1 8.8.8.8) # Custom DNS servers for Mullvad. Uncomment to use. Overridden by -d option.
# ----------------------------
# CONFIG VARIABLES YOU CAN BUT SHOULD NOT MANUALLY CHANGE
# It is better to use the flags -f/--file and -t/--table to change this dynamically
RULES_FILE="mullvad.rules"
NFT_TABLENAME="mullvad-ts"
# Do not edit past this point.
# ----------------------------

# Global variables
script_name="$(basename "$0")"
args="$*"
rules_checked=false

# Functions
## Checks and errors
die() {
  printf "%s Exiting.\n" "$*" 1>&2
  exit 1
}

### Prefix external command output with a string
function prefix_output {
  if [[ $1 != sudo ]]; then
    local prefix=$1
  else
    local prefix=$2
  fi
  prefix=${prefix^^}
  "${@}" 2>&1 | while read -r line; do
    printf "[%s] %s\n" "${prefix}" "${line}"
  done
  # shellcheck disable=SC2086
  return ${PIPESTATUS[0]}
}

bash_version_check() {
  # This script requires Bash >= 4.4 since it uses readarray.
  local error="This script requires Bash >=4.4."
  if [[ ${BASH_VERSINFO[0]} = 4 ]]; then
    if [[ ${BASH_VERSINFO[1]} -lt 4 ]]; then
      die "$error"
    fi
  elif [[ ${BASH_VERSINFO[0]} -lt 4 ]]; then
    die "$error"
  fi
}

### Catch interruption signal.
trap_sigint() {
  # TODO Do a proper cleanup.
  die "Interruption signal caught. Your network might be in a broken state. Run '$script_name $args' again to fix."
}

### Check that the nft rules file is valid.
check_nft_rules() {
  if [[ $rules_checked = false ]]; then
    RULES_DIR="${RULES_DIR%/}/"
    if [[ ! -f $RULES_DIR/$RULES_FILE ]]; then
      die "[ERROR] The file '$RULES_DIR/$RULES_FILE' does not exist."
    fi

    printf "[INFO] Checking nftables rules validity...\n"
    if prefix_output sudo nft -cf "$RULES_DIR/$RULES_FILE"; then
      printf "[INFO] nftables rules are valid.\n"
    else
      die "[ERROR] nftables rules are not valid. Fix them before running the script."
    fi
    rules_checked=true
  fi
}

## Usage functions
usage() {
  printf "usage: %s <action> [OPTIONS]\n" "$script_name"
  printf "\n"
  printf "Action: up\n"
  upusage
  printf "\n"
  printf "Action: down\n"
  downusage
  printf "\n"
  printf "Action: conf\n"
  confusage
  exit 0
}

upusage() {
  cat <<EOF
Apply nftables configuration and connect to Mullvad and Tailscale/Zerotier
$script_name up [OPTIONS]
  -h, --help:         Show this help message
  -r, --ram:          No-disk/RAM only Mullvad relays (default: all servers)
  -z, --zerotier:     Use Zerotier instead of Tailscale
  -d, --dns:          Set custom Mullvad DNS Server (i.e. -d 1.1.1.1 or -d 8.8.8.8,1.1.1.1)
  -c, --country:      Specify a country code to connect to (i.e. -c gb or -c fr,pt,es)
  -f, --file:         Specify a particular NFT rules file. Indicate path to file (absolute or relative) (default: $RULES_FILE)
EOF
}

downusage() {
  cat <<EOF
Bring down Mullvad and remove nftables configuration
$script_name down [OPTIONS]
  -h, --help:         Show this help message
  -a, --all:          Stop Mullvad and Tailscale/Zerotier (default: only stop Mullvad)
  -z, --zerotier:     Use Zerotier instead of Tailscale
  -t, --table:        Indicate the nft tablename to bring down (default: $NFT_TABLENAME)
EOF
}

confusage() {
  cat <<EOF
Apply nftables configuration so Mullvad and Tailscale/Zerotier can work together and do nothing more
$script_name conf [OPTIONS]
  -u, --unconf:       Remove the nftables configuration
  -h, --help:         Show this help message
EOF
}

remove_nft_rules() {
  printf "[INFO] Removing nftables config...\n"
  if prefix_output sudo nft delete table inet "$NFT_TABLENAME"; then
    printf "[INFO] nftables config removed.\n"
  else
    die "[ERROR] Can't remove table '$NFT_TABLENAME'. Has it already been removed?"
  fi
}

## Action functions
up() {
  # Bring down current connections
  if [[ $zerotier = false ]]; then
    bash "$0" down -a || printf "[INFO] Ignoring error...\n"
  else
    bash "$0" down -a -z || printf "[INFO] Ignoring error...\n"
  fi

  # Update relay list and store it in an array
  prefix_output mullvad relay update
  readarray -t MULLVAD_RELAYS < <(mullvad relay list | awk '/wireguard|wg/ { print $1 }')

  # Include countries
  ## User-specified option overrides INCLUDE_COUNTRY_CODES.
  if [[ ! $country_code = false ]]; then
    # Turn into regex.
    country_code="$(tr ',' '|' <<<"$country_code")"
    readarray -t MULLVAD_RELAYS < <(printf '%s\n' "${MULLVAD_RELAYS[@]}" | grep -E "^$country_code")
    size=${#MULLVAD_RELAYS[@]}
    if [[ $size -lt 1 ]]; then
      die "[ERROR] Country code '$country_code' is not valid!"
    fi
  ## Use INCLUDE_COUNTRY_CODES unless command flag used.
  elif [[ -n $INCLUDE_COUNTRY_CODES && $country_code = false ]]; then
    # Check that the user-specified variable is valid.
    for c in "${INCLUDE_COUNTRY_CODES[@]}"; do
      if ! [[ $c =~ ^[a-z]{2}$ ]]; then
        die "The country code '$c' you have specified in INCLUDE_COUNTRY_CODES is not valid."
      fi
    done

    # Turn into regex.
    country_code="$(tr ',' '|' <<<"$INCLUDE_COUNTRY_CODES")"
    readarray -t MULLVAD_RELAYS < <(printf '%s\n' "${MULLVAD_RELAYS[@]}" | grep -E "^$country_code")

  # Exclude countries
  else
    # If EXCLUDE_COUNTRY_CODES is non-null
    if [[ -n ${EXCLUDE_COUNTRY_CODES[*]} ]]; then
      # Check that the user-specified variable is valid
      for c in "${EXCLUDE_COUNTRY_CODES[@]}"; do
        if ! [[ $c =~ ^[a-z]{2}$ ]]; then
          die "The country code '$c' you have specified in EXCLUDE_COUNTRY_CODES is not valid."
        fi
      done

      # Turn into regex
      exclude=$(tr " " "|" <<<"${EXCLUDE_COUNTRY_CODES[@]}")
    fi

    readarray -t MULLVAD_RELAYS < <(printf '%s\n' "${MULLVAD_RELAYS[@]}" | grep -Ev "^$exclude")
    size=${#MULLVAD_RELAYS[@]}
    if [[ $size -lt 1 ]]; then
      die "[ERROR] Too many countries have been excluded, none left."
    fi
  fi

  # Keep only RAM servers
  if [[ $ram = true ]]; then
    printf "Getting RAM-only servers..."
    readarray -t MULLVAD_RELAYS < <(printf '%s\n' "${MULLVAD_RELAYS[@]}" | grep -E 'wg')
  fi

  size=${#MULLVAD_RELAYS[@]}
  if [[ $size -lt 1 ]]; then
    die "[ERROR] No RAM-only servers are available with the options you specified."
  fi

  # Pick a random country code.
  size=${#MULLVAD_RELAYS[@]}
  index=$((RANDOM % size))
  relay=${MULLVAD_RELAYS[$index]}

  # DNS
  printf "[INFO] Setting DNS server(s) for Mullvad...\n"
  if [[ $dns = false && ${#DNS[*]} -eq 0 ]]; then
    prefix_output mullvad dns set default
  else
    if [[ $dns = false && ${#DNS[*]} -gt 0 ]]; then
      dns="${DNS[*]}"
    else
      dns="$(tr ',' ' ' <<<"$dns")"
    fi
    printf "[INFO] Setting up Mullvad DNS to %s...\n" "$dns"
    # shellcheck disable=SC2086
    prefix_output mullvad dns set custom $dns
  fi

  # Zerotier
  if [[ "$zerotier" = false ]]; then
    printf "[INFO] Restarting tailscale connection...\n"
    prefix_output sudo tailscale down
    prefix_output sudo tailscale up
  else
    printf "[INFO] Restarting zerotier-one connection...\n"
    prefix_output sudo systemctl start zerotier-one.service
    prefix_output sudo systemctl restart zerotier-one.service
  fi

  if [[ "$zerotier" = false ]]; then
    printf "[INFO] Applying nft rules...\n"
    prefix_output sudo nft -f "$RULES_DIR""$RULES_FILE"
  fi

  if [[ "$zerotier" = true ]]; then
    printf "[INFO] Applying nft rules...\n"
    # This repetition is necessary. If not used it does not work wiht ZT.
    # I need to investigate, but this is a workaround for the time being.
    prefix_output sudo nft -f "$RULES_DIR""$RULES_FILE"
    prefix_output sudo nft add table inet "$NFT_TABLENAME"
    sleep .5
    remove_nft_rules
    sleep .5
    prefix_output sudo nft -f "$RULES_DIR""$RULES_FILE"
    prefix_output sudo nft add table inet "$NFT_TABLENAME"
  fi

  # Connect to Mullvad
  printf "[INFO] Connecting to server '%s'...\n" "$relay"
  prefix_output mullvad relay set hostname "$relay"
  prefix_output mullvad connect -w

  if [[ "$zerotier" = true ]]; then
    printf "[INFO] Applying nft rules...\n"
    # This repetition is necessary. If not used it does not work wiht ZT.
    # I need to investigate, but this is a workaround for the time being.
    prefix_output sudo nft -f "$RULES_DIR""$RULES_FILE"
    prefix_output sudo nft add table inet "$NFT_TABLENAME"
    sleep .5
    prefix_output remove_nft_rules
    sleep .5
    prefix_output sudo nft -f "$RULES_DIR""$RULES_FILE"
    prefix_output sudo nft add table inet "$NFT_TABLENAME"

    sleep 6 &
    PID=$!
    printf 'Progress: \n'
    while [[ -d /proc/$PID ]]; do
      for s in █ █ █; do
        printf "%s" "$s"
        sleep .1
      done
    done
  fi

  if [[ $zerotier = true ]]; then
    printf "\n\nDone! You may need to wait a few seconds for the connection to be ready."
  fi
  exit 0
}

down() {
  if [[ $all = true ]]; then
    if [[ $zerotier = true ]]; then
      printf "Stopping zerotier-one...\n"
      prefix_output sudo systemctl stop zerotier-one
      printf "[INFO] Zerotier-one disconnected.\n"
    else
      printf "[INFO] Disconnecting tailscale...\n"
      prefix_output sudo tailscale down
      printf "[INFO] Tailscale disconnected.\n"
    fi
  fi

  printf "[INFO] Disconnecting Mullvad...\n"
  prefix_output mullvad disconnect && printf "[INFO] Mullvad disconnected.\n" || printf "[WARN] Mullvad didn't disconnect properly.\n"

  remove_nft_rules
}

conf() {
  if [[ $unconf = true ]]; then
    remove_nft_rules
  fi
  conf

  printf "[INFO] Applying nft rules...\n"
  if prefix_output sudo nft -f "$RULES_DIR/$RULES_FILE"; then
    printf "[INFO] nft rules applied\n"
    exit 0
  else
    die "[ERROR] There was an error while applying nft rules."
  fi
}

## Main function
main() {
  bash_version_check

  # Get the action to perform
  local actions=("up" "down" "conf")
  action="$1"
  if [[ ! ${actions[*]} =~ ${action} || -z $action ]]; then
    usage
    exit 1
  fi

  # Get options and flags
  ARGS=$(getopt -a --options d:f:c:t:auhzr --long "dns:,file:,table:,country:,all,unconf,zerotier,help,ram" -- "$@")
  eval set -- "$ARGS"

  help=false
  all=false
  unconf=false
  zerotier=false
  dns=false
  ram=false
  country_code=false

  while true; do
    case "$1" in
      -d | --dns)
        dns="$2"
        shift 2
        ;;
      -f | --file)
        RULES_FILE="$2"
        RULES_DIR=""
        shift 2
        ;;
      -c | --country)
        country_code="$2"
        shift 2
        ;;
      -t | --table)
        NFT_TABLENAME="$2"
        shift 2
        ;;
      -r | --ram)
        ram=true
        shift
        ;;
      -a | --all)
        all=true
        shift
        ;;
      -u | --unconf)
        unconf=true
        shift
        ;;
      -z | --zerotier)
        zerotier=true
        shift
        ;;
      -h | --help)
        help=true
        shift
        ;;
      --) break ;;
      *)
        die "Unknown option '$1'."
        ;;
    esac
  done

  if [[ $help = false && $unconf = false && $action != down ]]; then
    check_nft_rules
  fi

  case "$action" in
    "up")
      if [[ $help = true ]]; then
        upusage
        exit 0
      fi

      up
      ;;

    "down")
      if [[ $help = true ]]; then
        downusage
        exit 0
      fi

      down
      ;;

    "conf")
      if [[ $help = true ]]; then
        confusage
        exit 0
      fi

      conf
      ;;
  esac
}

# Body of script
trap trap_sigint SIGINT
main "$@"

# TODO Improve trap_sigint function.
# TODO prefix subcommand output