#!/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 a command as another user
# Note: sudo and su syntax is slightly different and, more importantly, passing
# su arguments with spaces needs escaping quotes and/or possibly spaces; YMMV.
# Ex.    orjail ls "/tmp/a\\ b" # systems with su but without sudo
run () {
  if [ "$SUDOBIN" ]; then
    $SUDOBIN -u "$USERNAME" --preserve-env=XDG_RUNTIME_DIR,WAYLAND_DISPLAY "$@"
  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
# umount is not needed as per `man 7 mount_namespaces`:
# "mount is implicitly unmounted because a mount namespace is removed"
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

    print G " * Remove temporary /etc dir ($ETC_DIR)..."
    eno rm -fr "$ETC_DIR"
  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 "                       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 "    -p, --private      Private home"
  print N "                       Mount a sandboxed home directory and set current directory."
  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> <private_home> <userhome> <name> <command> <arguments...>
if [ "$1" = "--inside" ]; then
  REALPATH=$(realpath "$0")
  REALPATH=$(dirname "$REALPATH")

  shift
  USERNAME="$1"
  shift
  ETC_DIR="$1"
  shift
  VERBOSE="$1"
  shift
  PRIVATE_HOME="$1"
  shift
  USERHOME="$1"
  shift
  NAME="$1"
  shift

  if [ -n "$ETC_DIR" ]; then
    print G " * Mirroring /etc"
    mount --bind "$ETC_DIR" /etc ||
      die "Failed to mount /etc from $ETC_DIR"
  fi

  if [ "$PRIVATE_HOME" = "y" ]; then
    DIR="$USERHOME/.orjail/$NAME"
    print G " * Mount a private $USERHOME 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 "$DIR" "$USERHOME" ||
      die "Failed to mount $USERHOME"
    cd "$USERHOME"
  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")

      if ! [ -x "$TORBIN" ]; then
        die "Tor executable '$2' is invalid."
      fi
      shift
    ;;

    # 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=${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

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

for cmd in ip iptables bc mkdir chown grep getent $FIREJAILBIN ${SUDOBIN:-su}; do
  if ! [ -x "$(command -v "$cmd")" ]; then
    die "Cannot locate $cmd executable, please install all needed dependencies:\r\nsudo apt-get install bc iptables grep sudo iproute2 libc-bin"
  fi
done

USERHOME=$(getent passwd "$USERNAME" | cut -d: -f 6)
if [ "$USERHOME" = "" ]; then
  die "User $USERNAME: invalid name or no home directory."
fi

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=$USERHOME/.orjail/$NAME")
  fi

  FIREJAILBIN=firejail
fi

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

# check if network namespace already exists
if ! ip netns list | eno grep \\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 torXXXXXX)
  TORCONFIGFILE=$(realpath "$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

  # this chown causes issues in archlinux where tmpfs seems to be protected...
  # so call this only after the config file is ready. see PR#71
  chown "$USERNAME" "$TORCONFIGFILE"

  # 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}"
    if [ -f "/tmp/orjail-${NAME}/lock" ]; then
      rm "/tmp/orjail-${NAME}/lock"
    fi
  else
    mkdir "/tmp/orjail-$NAME"
  fi
  chown -R "$USERNAME" "/tmp/orjail-${NAME}"
  chmod -R 700 "/tmp/orjail-${NAME}"

  # 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
    ETC_DIR="$USERHOME/.orjail/etc/$NAME"
    mkdir -p "$ETC_DIR"
    cp -a /etc/. "$ETC_DIR"

    print G " * Create a temporary /etc/resolv.conf and /etc/nsswitch.conf"
    # remove it in case it's a symlink
    rm "$ETC_DIR/resolv.conf" "$ETC_DIR/nsswitch.conf"
    echo "nameserver $IPHOST" > "$ETC_DIR/resolv.conf"
    chmod a+r "$ETC_DIR/resolv.conf"

    cat > "$ETC_DIR/nsswitch.conf" <<EOF
hosts: dns files
passwd:files
shadow:files
group:files
hosts:dns files
bootparams:files
ethers:files
netmasks:files
networks:files
protocols:files
rpc:files
services:files
automount:files
aliases:files
EOF

  fi
else
  print Y "$NAME network namespace already exists!"
  NAMESPACE_EXIST=y
  KEEP=y

  if ! ip netns pids "$NAME" | eno grep . ; then
    die "No process found for namespace $NAME. Please specify another/no namespace."
  fi
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)
    # This is like function run() - read comments there
    if [ "$SUDOBIN" ]; then
      nsenter -p -n  -i -m  -t "$pid" "$SUDOBIN" -u "$USERNAME" "$@"
    else
      nsenter -p -n  -i -m  -t "$pid" su "$USERNAME" -c "$*"
    fi
  else
    ip netns exec "$NAME" \
      unshare --ipc --fork --pid --mount --mount-proc \
      "$0" --inside "$USERNAME" "$ETC_DIR" "$VERBOSE" "$PRIVATE_HOME" "$USERHOME" "$NAME" "$@"
  fi
fi


# All done!
