#!/bin/bash

# Copyright (C) 2018 - underscore@autistici.org
# This work is free. You can redistribute it and/or modify it under the
# terms of the Do What The Fuck You Want To Public License, Version 2,
# as published by Sam Hocevar. See the COPYING file for more details.

REALPATH=$(realpath "$0")
REALPATH=$(dirname "$REALPATH")
DEFAULTNAME=orjail
NAME=$DEFAULTNAME
USERNAME=${SUDO_USER:-$(whoami)}
USEFIREJAIL=n
FIREJAILARGS=
VERBOSE=
TORBIN=
KEEP=
HIDDENSERVICE=
IPHOST=
IPHOSTMASK=10.200.x.2
IPNETNS=
IPNETNSMASK=10.200.x.1
NETMASK=30
TRANSPORT=9040
DNSPORT=5354
SUDOBIN=$(command -v sudo)
NOSETUPERROR=n
HOSTTORRC=n
NAMESPACE_EXIST=n
PRIVATE_HOME=n
unset FIREJAILARGS

# Functions
# ~~~~~~~~~

print_real() {
  if [ "$VERBOSE" != y ]; then
    return
  fi

  if [ -t 1 ]; then
    NCOLORS=$(tput colors)

    if test -n "$NCOLORS" && test "$NCOLORS" -ge 8; then
      NORMAL="$(tput sgr0)"
      RED="$(tput setaf 1)"
      GREEN="$(tput setaf 2)"
      YELLOW="$(tput setaf 3)"
    fi
  fi

  if [[ $2 = 'G' ]]; then
    # shellcheck disable=SC2086
    echo $1 -e "${GREEN}$3${NORMAL}"
  elif [[ $2 = 'Y' ]]; then
    # shellcheck disable=SC2086
    echo $1 -e "${YELLOW}$3${NORMAL}"
  elif [[ $2 = 'N' ]]; then
    # shellcheck disable=SC2086
    echo $1 -e "$3"
  else
    # shellcheck disable=SC2086
    echo $1 -e "${RED}$3${NORMAL}"
  fi
}

print() {
  print_real '' "$1" "$2"
}

printn() {
  print_real "-n" "$1" "$2"
}

printv() {
  OLDVERBOSE=$VERBOSE
  VERBOSE=y
  print_real '' "$1" "$2"
  VERBOSE=$OLDVERBOSE
}

printvn() {
  OLDVERBOSE=$VERBOSE
  VERBOSE=y
  print_real "-n" "$1" "$2"
  VERBOSE=$OLDVERBOSE
}

run () {
  if [ "$SUDOBIN" ]; then
    $SUDOBIN -u "$USERNAME" "$@"
  else
    su "$USERNAME" -c "$*"
  fi
}
# exec no output
eno() {
  if [ "$VERBOSE" != y ]; then
    "$@" &>/dev/null
  else
    "$@"
  fi
}

# kill tor, remove network namespace, cleanup added iptables rules
cleanup() {
  if [ "$NOSETUPERROR" != y ]; then
    printv R "[Error] in command: $BASH_COMMAND"
    if [ "$VERBOSE" != y ]; then
      printv R "[Error] Enable verbose mode to debug (using -v)"
    fi
  fi
  set +e

  TORPID=$(cat /tmp/orjail-"$NAME"/pid 2> /dev/null)
  if [ "$KEEP" = y ]; then
    print G " * Keep Tor process $TORPID running"
    print G " * Keep $NAME namespace active"
  else
    print G " * Remove Tor temporary configuration"
    [ -f "$TORCONFIGFILE" ] && rm "$TORCONFIGFILE"
    print G " * Killing Tor process $TORPID"
    [ "$TORPID" ] && eno kill -9 "$TORPID" && eno wait $!
    print G " * Killed $TORPID"

    print G " * Remove Tor DataDirectory: /tmp/orjail-$NAME"
    eno rm -fr "/tmp/orjail-$NAME"

    print G " * Remove in-$NAME network interface"
    eno ip link del "in-$NAME"

    print G " * Delete network namespace $NAME"
    eno ip netns delete "$NAME"

    print G " * Cleaning up iptables rules..."
    iptables -S | grep \\b"in-$NAME"\\b | while read -r line; do
      # shellcheck disable=SC2086
      eno iptables ${line//-A/-D}
    done
    iptables -t nat -S | grep \\b"in-$NAME"\\b | while read -r line; do
      # shellcheck disable=SC2086
      eno iptables -t nat ${line//-A/-D}
    done


  fi

  if [ ! -z "$RESOLVEFILE" ]; then
  print G " * Remove the temporary resolve file $RESOLVEFILE..."
  eno rm "$RESOLVEFILE"
  fi
}

error() {
  printv R "$1"
}

die() {
  error "$1"
  exit 1
}

help_and_exit() {
  VERBOSE=y
  printn N "Usage: "
  printn Y "$DEFAULTNAME"
  print G " <options> [command <arguments>...]"
  print N "Options:"
  print N "    -h, --help         It shows this menu."
  print N "    -u, --user <user>  Execute the command with this user permission. By default '$USERNAME'."
  print N "    -n, --name <name>  Set a custom namespace name. By default '$DEFAULTNAME'."
  print N "    -v, --verbose      Verbose mode."
  print N "    -k, --keep         Don't delete namespace and don't kill tor after the execution."
  print N "    -f, --firejail     Use firejail as a security container ($SUDOBIN orjail -f pidgin)."
  print N "        --firejail-args \"<args>\""
  print N "                       Set arguments to pass to firejail surrounded by quotes. (\"--hostname=host --env=PS1=[orjail]\")"
  print N "    -H, --hidden <port>"
  print N "    -p, --private      Private home"
  print N "                       Enable Tor as an hidden service forwarding request from/to specified port."
  print N "    -d, --hiddendir <dir>"
  print N "                       Specify where to search for hidden service 'hostname' and 'private_key'."
  print N "    -s, --shell        Execute a shell (using your current one)"
  print N "        --host-torrc   Include your torrc host."
  print N "    -t, --tor-exec     Select a Tor executable to use. The path can be full, relative or be in \$PATH"
  print N "    -r, --routing <ip_host> <ip_ns> <netmask>"
  print N "                       Set custom IPs. By default $IPHOSTMASK/$IPNETNSMASK/$NETMASK."
  print N "        --trans-port <port>"
  print N "                       Set tor TransPort. By default $TRANSPORT"
  print N "        --dns-port <port>"
  print N "                       Set custom DnsPort. By default $DNSPORT"
  print N "        --port-range <port>-<port>"
  print N "                       Generate random TransPort and DnsPort in the defined range."
  exit "$1"
}

# Inside part
# ~~~~~~~~~~~

# This script calls itself. yeah \o/ This part is executed only inside the
# namespace. The arguments are:
# --inside <username> <resolvefile> <verbose> <command> <arguments...>
if [ "$1" = "--inside" ]; then
  REALPATH=$(realpath "$0")
  REALPATH=$(dirname "$REALPATH")

  shift
  USERNAME="$1"
  shift
  RESOLVEFILE="$1"
  shift
  VERBOSE="$1"
  shift
  PRIVATE_HOME="$1"
  shift
  NAME="$1"
  shift


  print G " * Replacing resolv.conf..."
  mount --bind -o users "$RESOLVEFILE" /etc/resolv.conf || \
    error "Failed to mount /etc/resolv.conf."

  if [ "$PRIVATE_HOME" = "y" ]; then
    DIR="/home/$USERNAME/.orjail/$NAME"
    print G " * Mount a private /home/$USERNAME from $DIR"
    if ! [ -d "$DIR" ]; then
      mkdir -p "$DIR"
      echo "PS1=\"[orjail@$NAME] %n@%m:%~%#  \"" > "$DIR"/.zshrc
      echo "PS1=\"[orjail@$NAME] \[\e]0;\u@\h: \w\a\]$  \"" > "$DIR"/.bashrc
    fi
    mount --bind -o users "/home/$USERNAME/.orjail/$NAME" "/home/$USERNAME" || \
      error "Failed to mount /home/$USERNAME"
    cd "/home/$USERNAME"
   fi

  print G " * Executing..."

  run "$@"
  exit
fi

# The tool
# ~~~~~~~~

# Arguments check
while [[ $# -gt 0 ]]; do
  key="$1"
  case $key in
    # Replacing the name
    -n|--name)
      NAME="$2"
      shift

      if [ "$NAME" = "" ]; then
        die "$key requires an argument."
      fi
    ;;

    # Username
    -u|--username)
      USERNAME="$2"
      shift

      if [ "$USERNAME" = "" ]; then
        die "$key requires an argument."
      fi
    ;;

    -v|--verbose)
      VERBOSE=y
      ;;

    -p|--private)
       PRIVATE_HOME=y
      ;;

    -k|--keep)
      KEEP=y
      ;;

    -H|--hidden)
      HIDDENSERVICE=y
      HSERVICEPORT="$2"
      shift

      if [ "$HSERVICEPORT" = "" ]; then
        die "$key requires an argument."
      fi
      ;;

    -d|--hiddendir)
      HIDDENSERVICEDIR="$2"
      shift

      if [ "$HIDDENSERVICEDIR" = "" ]; then
        die "$key requires an argument."
      fi
      ;;

    -r|--routing)
      IPHOST="$2"
      shift
      IPNETNS="$2"
      shift
      NETMASK="$2"
      shift

      if [ "$IPHOST" = "" ] ||
         [ "$IPNETNS" = "" ] ||
         [ "$NETMASK" = "" ]; then
        die "$key requires 3 arguments."
      fi
      ;;

    -f|--firejail)
      USEFIREJAIL=y
      ;;

    --firejail-args)
      IFS=' ' read -r -a FIREJAILARGS <<< "$2"
      shift

      if [ "${#FIREJAILARGS[@]}" -eq 0 ]; then
        die "$key requires an argument."
      fi
      USEFIREJAIL=y
      ;;

    -s|--shell)
      set -- "$@" "$SHELL"
      ;;

    # TransPort
    --trans-port)
      TRANSPORT="$2"
      shift
      [ "$TRANSPORT" ] || die "$key requires an argument."
    ;;

    # EnableHostTorrc
    --host-torrc)
      HOSTTORRC=y
    ;;

    # Select a Tor executable
    -t|--tor-exec)
      TORBIN=$(command -v "$2")
      shift

    if ! [ -f "$TORBIN" ]; then
      die "This Tor executable does not exist."
    fi
    ;;

    # DnsPort
    --dns-port)
      DNSPORT="$2"
      shift
      [ "$DNSPORT" ] || die "$key requires an argument."
    ;;

     # PortRange
    --port-range)
      range="$2"
      shift

      [ "$range" ] || die "$key requires an argument."

      if ! [[ "$range" =~ [0-9]{2,5}-[0-9]{2,5} ]]; then
        die "port range should be like 1000-9000."
      fi

      # disabled bacause read -a and mapfile are not available in zsh
      # shellcheck disable=SC2207
      rnd_range=( $(shuf -i "$range" -n 2) )
      TRANSPORT="${rnd_range[0]}"
      DNSPORT="${rnd_range[1]}"

      printv G "random generated TransPort: $TRANSPORT"
      printv G "random generated DnsPort: $DNSPORT"
    ;;

    # Help menu
    -h|--help)
      help_and_exit 0
    ;;

    # End my options
    --)
      shift
      break
    ;;

    # Illegal options
    -*)
      die "$key unknown option."
    ;;

    # The rest
    *)
    break
    ;;
  esac
  shift
done

if [[ $EUID -ne 0 ]]; then
   die "$DEFAULTNAME must be run as root."
fi

TORBIN=$(command -v tor)
if ! [ -x "$TORBIN" ]; then
  die "Can't locate tor executable.";
fi

# No arguments, no party
if [ "$1" = "" ]; then
  help_and_exit 1
fi

# exit on error, and call cleanup on exit
set -e
trap cleanup EXIT

# Check linux kernel
if [ "$(uname)" != "Linux" ]; then
  die "No Linux no party"
fi

err=0
for cmd in ip iptables bc mkdir chown grep $FIREJAILBIN ${SUDOBIN:-su}; do
  if ! [ -x "$(command -v "$cmd")" ]; then
    printv R "Cannot locate $cmd executable"
    err=1
  fi
done
! [ 0 = $err ] && exit $err

FIREJAILBIN=
if [ $USEFIREJAIL = y ]; then
  firejail_version=$(firejail --version | grep -io "9.[0-9]\\{2\\}")
  if [[ $(echo "$firejail_version>9.44" | bc) -eq 0 ]]; then
	  die "orjail requires at least firejail 0.9.44.10 to run."
  fi

  if [ "$PRIVATE_HOME" = "y" ]; then
    FIREJAILARGS+=("--private=/home/$USERNAME/.orjail/$NAME")
  fi

  FIREJAILBIN=firejail
fi
# check if network namespace already exists
if ! ip netns list | eno grep -e \\b"$NAME"\\b; then

  # generate a random available address from specified subnet mask
  set +e
  USEDADDR=$(ip addr show type veth| grep -Po 'inet \d+.\d+.\K(\d+)')
  for available_subnet in $(shuf -i 1-255 -n 254); do
    if ! echo "$USEDADDR" | eno grep "^$available_subnet$"; then
      break
    fi
  done
  set -e

  IPHOST=${IPHOSTMASK//x/$available_subnet}
  IPNETNS=${IPNETNSMASK//x/$available_subnet}
  print G " * $IPHOST <---> $IPNETNS"
  print G " * Creating a $NAME namespace..."

  # add network namespace
  ip netns add "$NAME"

  # Create veth link.
  print G " * Creating a veth link..."
  ip link add "in-$NAME" type veth peer name "out-$NAME"

  # Add out to NS.
  print G " * Sharing the veth interface..."
  ip link set "out-$NAME" netns "$NAME"

  ## setup ip address of host interface
  print G " * Setting up IP address of host interface ($IPHOST/$NETMASK)"
  ip addr add "$IPHOST/$NETMASK" dev "in-$NAME"
  ip link set "in-$NAME" up

  # setup ip address of peer
  print G " * Setting up IP address of peer interface ($IPNETNS/$NETMASK)"
  ip netns exec "$NAME" ip addr add "$IPNETNS/$NETMASK" dev "out-$NAME"
  ip netns exec "$NAME" ip link set "out-$NAME" up

  # default route
  print G " * Default routing up..."
  ip netns exec "$NAME" ip route add default via "$IPHOST"

  # bring loopback interface up inside sandbox
  print G " * Bringing orjail loopback up..."
  ip netns exec "$NAME" ip link set lo up

  # resolve with tor
  print G " * Resolving via Tor"
  iptables -t nat -A  PREROUTING -i "in-$NAME" -p udp -d "$IPHOST" --dport 53 -j DNAT \
        --to-destination "$IPHOST":"$DNSPORT"

  # traffic througth tor
  print G " * Traffic via Tor..."
  iptables -t nat -A  PREROUTING -i "in-$NAME" -p tcp --syn -j DNAT \
           --to-destination "$IPHOST":"$TRANSPORT"
  iptables -A OUTPUT -m state -o "in-$NAME" --state ESTABLISHED,RELATED -j ACCEPT

  # REJECT all traffic coming from orjail
  # this is needed to avoid reaching other interfaces
  iptables -I INPUT -i "in-$NAME" -p udp --destination "$IPHOST" --dport "$DNSPORT" -j ACCEPT
  iptables -I INPUT -i "in-$NAME" -p tcp --destination "$IPHOST" --dport "$TRANSPORT" -j ACCEPT
  if [[ $HIDDENSERVICE = y ]]; then
    iptables -I INPUT -i "in-$NAME" -p tcp --source "$IPNETNS" --sport "$HSERVICEPORT" -j ACCEPT
  fi
  # while we inserted the rules above, the DROP rule must be appended instead
  iptables -A INPUT -i "in-$NAME" -j DROP

  # disable forwarding (no packets from here should be forwarded!)
  iptables -I FORWARD -i "in-$NAME" -j DROP
  iptables -I FORWARD -o "in-$NAME" -j DROP

  sysctl -w -q "net.ipv4.conf.in-$NAME".forwarding=0

  # everything coming/redirected from orjail does not have to reach any other interface
  iptables -t nat -I POSTROUTING 1 \! -o "in-$NAME" -s "$IPHOST/$NETMASK" -j RETURN
  iptables -t nat -I PREROUTING 1 \! -i "in-$NAME" -d "$IPHOST/$NETMASK" -j RETURN

  # prevent port redirection to be made in orjail
  iptables -t nat -A PREROUTING -i "in-$NAME" -j RETURN

  # prevent external traffic to reach orjail
  iptables -A INPUT ! -i "in-$NAME" -s "$IPHOST/$NETMASK" -j DROP
  iptables -A INPUT ! -i "in-$NAME" -d "$IPHOST/$NETMASK" -j DROP

  # execute tor
  print G " * Creating the Tor configuration file..."

  # automatically detect tor version and use appropriate syntax
  TORVERSION="$($TORBIN --version|grep -Eo ' ([0-9.]+)'|xargs)"
  print G " * Tor version is $TORVERSION"

  TORCONFIGFILE=$(mktemp /tmp/torXXXXXX)
  chown "$USERNAME" "$TORCONFIGFILE"
  if [ "$HOSTTORRC" = "y" ] && [ -f /etc/tor/torrc ]; then
    echo '%include /etc/tor/torrc' >> "$TORCONFIGFILE"
  fi

  cat >> "$TORCONFIGFILE" <<EOF
  DataDirectory /tmp/orjail-${NAME}
  AutomapHostsSuffixes .onion,.exit
  AutomapHostsOnResolve 1
  PidFile      /tmp/orjail-${NAME}/pid
  User         ${USERNAME}
  VirtualAddrNetworkIPv4 ${IPNETNS}/16
  TransPort ${IPHOST}:${TRANSPORT}
  DNSPort ${IPHOST}:${DNSPORT}
  SOCKSPort 0
  RunAsDaemon 1
EOF

  if [[ "$HIDDENSERVICE" = y ]]; then
    HIDDENSERVICEDIR=${HIDDENSERVICEDIR:-/tmp/orjail-$NAME}
    HIDDENSERVICEDIR=$(realpath "$HIDDENSERVICEDIR")
    print G " * Hidden Service Dir $HIDDENSERVICEDIR"
    { echo "HiddenServiceDir $HIDDENSERVICEDIR" >> "$TORCONFIGFILE";
      echo "HiddenServiceVersion 3" >> "$TORCONFIGFILE";
      echo "HiddenServicePort $HSERVICEPORT $IPNETNS:$HSERVICEPORT"; } >> "$TORCONFIGFILE"
  fi

  # reuse tor host's cache
  if [ -d "/var/lib/tor" ]; then
  print G " * Copying host's tor cache"
  cp -d -R /var/lib/tor "/tmp/orjail-${NAME}"
  chown -R "$USERNAME" "/tmp/orjail-${NAME}/"
  chmod -R 700 "/tmp/orjail-${NAME}/"
    if [ -f "/tmp/orjail-${NAME}/lock" ]; then
      rm "/tmp/orjail-${NAME}/lock"
    fi
  else
  mkdir "/tmp/orjail-$NAME"
  fi

  # executing tor
  print G " * Executing Tor..."
  if [ "$VERBOSE" != y ]; then
     "$TORBIN" -f "$TORCONFIGFILE" &>/dev/null
  else
     "$TORBIN" -f "$TORCONFIGFILE"
  fi

  if [[ "$HIDDENSERVICE" = y ]]; then
    print G " * Your hidden service domain:"
    cat "$HIDDENSERVICEDIR"/hostname
  fi

  if ! [ $USEFIREJAIL = y ]; then
    RESOLVEFILE=$(mktemp /tmp/resolveXXXXXX)
    print G " * Creating a temporary /etc/resolv.conf ($RESOLVEFILE)..."
    echo "nameserver $IPHOST" > "$RESOLVEFILE"
    chmod a+r "$RESOLVEFILE"
  fi
else
  print Y "$NAME network namespace already exists!"
  NAMESPACE_EXIST=y
  KEEP=y
fi

NOSETUPERROR=y
# use firejail as security container
if [ $USEFIREJAIL = y ]; then
  # shellcheck disable=SC2068
  if [ $NAMESPACE_EXIST = y ]; then
    run "$FIREJAILBIN" ${FIREJAILARGS[@]} --join="$NAME" "$@"
  else
    run "$FIREJAILBIN" \
      --quiet --dns="$IPHOST" --name="$NAME" --netns="$NAME" \
      --hostname=host --noroot --private-tmp --private-dev \
      ${FIREJAILARGS[@]} "$@"
  fi
else #or without
  if [ $NAMESPACE_EXIST = y ]; then
    pid=$(ip netns pids "$NAME" | tail -1)
    nsenter -p -n  -i -m  -t "$pid" "$SUDOBIN" -u "$USERNAME" "$@"
  else
    ip netns exec "$NAME" \
      unshare --ipc --fork --pid --mount --mount-proc \
      "$0" --inside "$USERNAME" "$RESOLVEFILE" "$VERBOSE" "$PRIVATE_HOME" "$NAME" "$@"
  fi
fi


# All done!
