#!/bin/bash
set -e

source ${SCRIPT_DIR}/logging
source ${SCRIPT_DIR}/kasm_api

wait_for_container() {
  info "Waiting for container to be ready."

  container_name=
  container_pid=
  container_ip=
  container_ports=

  # get the container details based on the network namespaces
  until [ -e "${container_ns}" ]; do
    sleep 1
  done

  # There is a delay from when the container_ns is ready and the container actually shows up in docker commands in
  # Docker 28.x . So we have to wait until it shows up.
  retries=0
  while ((retries < 20)); do
    for container_id in $(docker ps -q); do
      if docker inspect $container_id | grep -q "$container_ns"; then
        container_name=$(docker inspect --format '{{.Name}}' ${container_id} | sed 's/^\///')
        container_ports=($(docker inspect -f '{{range $p, $conf := .NetworkSettings.Ports}}{{if $conf}}{{(index (split $p "/") 0)}} {{end}}{{end}}' $container_name))
        container_pid=$(docker inspect --format '{{.State.Pid}}' ${container_id})
        container_ip=$(docker inspect --format '{{.NetworkSettings.Networks.kasm_sidecar_network.IPAddress}}' ${container_id})
        break
      fi
    done
    if [[ -n "$container_name" && -n "$container_pid" && -n "$container_ip" ]]; then
      break
    fi
    sleep 1
    ((retries++))
  done

  if [ -z "$container_name" ] ; then
    info "Could not find container name. Exiting."
    exit 1
  fi

  if [ -z "$container_pid" ] ; then
    info "Could not find container pid. Exiting."
    exit 1
  fi

  if [ -z "$container_ip" ]; then
    info "Could not find container ip. Exiting."
    exit 1
  fi

  until [ -e "/proc/$container_pid/ns/mnt" ]; do
    sleep 1
  done

  info "Container \"$container_name\" is ready after $retries seconds. Name: $container_name, PID: $container_pid, IP: $container_ip."
}

find_iptables_rules() {
  local type=$1
  local chain=$2
  local rules=()

  if [ -z "$chain" ]; then
    rules=($(iptables -t "$type" -L --line-numbers -n -v | awk '/br_kasm_sidecar/ {print $1}' | sort -rn 2>/dev/null))
  else
    rules=($(iptables -t "$type" -L "$chain" --line-numbers -n -v | awk '/br_kasm_sidecar/ {print $1}' | sort -rn 2>/dev/null))
  fi

  echo "${rules[@]}"
}

create_bridge_interface() {
   should_reconnect_proxy=false

  # ensure that any existing bridge interface has the correct gateway IP
  if ip link show $bridge_name > /dev/null 2>&1; then
    current_bridge_ip=$(ip addr show br_kasm_sidecar | grep -o -E "inet [0-9.]+" | sed 's/inet //')
    bridge_gateway_ip_no_cidr=$(echo $bridge_gateway_ip | cut -d/ -f1)

    if [ "$current_bridge_ip" == "$bridge_gateway_ip_no_cidr" ]; then
      return
    fi

    info "Detected existing bridge interface \"$bridge_name\" with incorrect gateway IP \"$current_bridge_ip\". Removing."
    ip link del $bridge_name

    if [[ "$container_name" != "kasm_proxy" ]]; then
      should_reconnect_proxy=true
    fi
  fi

  info "Creating bridge interface \"$bridge_name\" with gateway IP \"$bridge_gateway_ip\"."

  # create the bridge interface
  bridge_broadcast=$(echo $bridge_gateway_ip | cut -d. -f1-2).255.255
  ip link add name $bridge_name type bridge
  ip addr add $bridge_gateway_ip dev $bridge_name broadcast $bridge_broadcast
  ip link set $bridge_name up

  # clear existing iptables rules
  info "Clearing existing iptables rules for bridge interface \"$bridge_name\"."

  RULES=$(find_iptables_rules "filter" "FORWARD")
  for RULE in ${RULES[@]}; do
    iptables -t "filter" -D "FORWARD" ${RULE}
  done

  RULES=$(find_iptables_rules "filter" "DOCKER")
  for RULE in ${RULES[@]}; do
    iptables -t "filter" -D "DOCKER" ${RULE}
  done

  RULES=$(find_iptables_rules "filter" "DOCKER-ISOLATION-STAGE-1")
  for RULE in ${RULES}; do
    iptables -t "filter" -D "DOCKER-ISOLATION-STAGE-1" ${RULE}
  done

  RULES=$(find_iptables_rules "filter" "DOCKER-ISOLATION-STAGE-2")
  for RULE in ${RULES}; do
    iptables -t "filter" -D "DOCKER-ISOLATION-STAGE-2" ${RULE}
  done

  RULES=$(find_iptables_rules "nat" "POSTROUTING")
  for RULE in ${RULES}; do
    iptables -t "nat" -D "POSTROUTING" ${RULE}
  done

  # set up iptables rules
  info "Setting up iptables rules for bridge interface \"$bridge_name\" using \"$container_ns_id\" namespace."
  rule_idx=$(iptables -L FORWARD -v -n --line-numbers | grep docker0 | awk 'END{if (NF) print $1}')
  iptables -I FORWARD $((rule_idx + 1)) -o $bridge_name -m conntrack --ctstate RELATED,ESTABLISHED -j ACCEPT
  iptables -I FORWARD $((rule_idx + 2)) -o $bridge_name -j DOCKER
  iptables -I FORWARD $((rule_idx + 3)) -i $bridge_name ! -o $bridge_name -j ACCEPT
  iptables -I FORWARD $((rule_idx + 4)) -i $bridge_name -o $bridge_name -j ACCEPT
  iptables -I DOCKER-ISOLATION-STAGE-1 2 -i $bridge_name ! -o $bridge_name -j DOCKER-ISOLATION-STAGE-2
  iptables -I DOCKER-ISOLATION-STAGE-2 2 -o $bridge_name -j DROP
  iptables -A POSTROUTING -t nat ! -o $bridge_name -s $bridge_gateway_ip -j MASQUERADE
  iptables -I DOCKER -i $bridge_name -j RETURN

  info "Bridge \"$bridge_name\" interface created."

  if [ "$should_reconnect_proxy" = true ]; then
    info "Disconnecting proxy container to bridge."
    docker network disconnect kasm_sidecar_network kasm_proxy

    info "Reconnecting proxy container to bridge."
    docker network connect kasm_sidecar_network kasm_proxy
  fi
}

configure_proxy_container() {
  info "Configuring proxy container."

  # create a veir pair
  veth_id=$(cat /dev/urandom | tr -dc 'a-f0-9' | fold -w 6 | head -n 1)
  veth_name=k-p-e-$veth_id
  vpeer_name=k-p-p-$veth_id
  ip link add $veth_name type veth peer name $vpeer_name
  echo $veth_name > /var/run/kasm-sidecar/$container_ns_id/ifaces

  # attach one end to the bridge
  ip link set $veth_name up
  ip link set $veth_name master $bridge_name

  # attach the other end to the container
  ip link set $vpeer_name netns ${container_ns}
  nsenter --net=$container_ns ip link set $vpeer_name up
  nsenter --net=$container_ns ip addr add $container_ip/16 dev $vpeer_name broadcast $bridge_broadcast

  # ensure the plugin uses the correct API hostname
  if [ ! -f "/var/run/kasm-sidecar/api_hostname_external" ]; then
    default_network_container_ip=$(docker inspect -f '{{if .NetworkSettings.Networks.kasm_default_network}}{{.NetworkSettings.Networks.kasm_default_network.IPAddress}}{{end}}' $container_name)
    proxy_container_ip=${container_ip:-$default_network_container_ip}
    proxy_container_port=${container_ports[0]:-443}
    echo ${proxy_container_ip}:${proxy_container_port} > /var/run/kasm-sidecar/api_hostname
    info "Proxy container configured using internal IP (${proxy_container_ip}:${proxy_container_port})."
  else
    proxy_container_hostname=$(cat /var/run/kasm-sidecar/api_hostname)
    info "Proxy container configured using external hostname. (${proxy_container_hostname})."
  fi
}

configure_session_container() {
  info "Configuring session container."

  info "Removing existing network interfaces."
  nsenter --net=$container_ns sh -c 'for iface in $(ip link show | grep -oP "^\d+: \K[^:]+(?=:)" | grep -v lo); do iface_clean=${iface%@*}; ip link del $iface_clean || true; done'

  info "Setting container status to starting"
  set_container_status "starting" "0" "Configuring container network..."

  if [ -z "$script_file" ] || [ -z "$script_data" ]; then
    info "Missing script file or data. Exiting."
    exit 1
  fi

  # clear the egress signal file
  docker exec --user root $container_name rm -rf /dockerstartup/.egress_status

  # create a veth pair
  info "Creating veth pair for container."
  veth_id=$(cat /dev/urandom | tr -dc 'a-f0-9' | fold -w 6 | head -n 1)
  veth_name=k-e-$veth_id
  vpeer_name=k-p-$veth_id

  ip link add $veth_name type veth peer name $vpeer_name
  echo $veth_name > /var/run/kasm-sidecar/$container_ns_id/ifaces

  # attach one end to the bridge
  info "Attaching veth pair to bridge."
  ip link set $veth_name up
  ip link set $veth_name master $bridge_name

  # attach the other end to the container
  info "Attaching veth pair to container."
  ln -s $container_ns /var/run/netns/$container_ns_id
  ip link set $vpeer_name netns $container_ns_id

  # ensure kasm_proxy is reachable from the container
  proxy_ip=$(docker inspect kasm_proxy | jq -r ".[0].NetworkSettings.Networks.kasm_sidecar_network.IPAddress")
  if [ -z "$proxy_ip" ]; then
    info "Could not find proxy container IP. Exiting."
    exit 1
  fi

  # as the session container is set up to use the network plugin,
  # it must communicate with the proxy container on the same kasm_sidecar_network network.
  kasm_api_host=$(docker exec $container_name env | grep KASM_API_HOST | cut -d= -f2)
  kasm_api_host_ip=$kasm_api_host
  kasm_api_port=$(docker exec $container_name env | grep KASM_API_PORT | cut -d= -f2)

  if ! [[ "$kasm_api_host" =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
    kasm_api_host_ip=$(getent hosts $kasm_api_host | awk '{ print $1 }')
    kasm_api_host_ip=${kasm_api_host_ip:-$proxy_ip}
    info "Adding entry for \"${kasm_api_host}\" in hosts file (\"$kasm_api_host_ip\")."
    nsenter --mount=/proc/$container_pid/ns/mnt sh -c "printf '%s\t%s\n' \"$kasm_api_host_ip\" \"$kasm_api_host\" >> /etc/hosts"
  fi

  if [[ "$script_file" != "egress_custom_setup" ]]; then
    # set the container DNS servers
    nsenter --mount=/proc/$container_pid/ns/mnt sh -c "echo '' > /etc/resolv.conf"
    dns_servers=$(echo $script_data | base64 -d | jq -r '.dns[]')
    info "Setting initial DNS servers for container: $dns_servers"

    for dns_server in $dns_servers; do
      nsenter --mount=/proc/$container_pid/ns/mnt sh -c "echo \"nameserver $dns_server\" >> /etc/resolv.conf"
    done

    # set the container iptables rules
    info "Setting up iptables rules for container."
    nsenter --net=$container_ns iptables -F

    info "Allowing direct access to $kasm_api_host:$kasm_api_port via $vpeer_name interface."
    nsenter --net=$container_ns iptables -A INPUT -s $kasm_api_host_ip -i $vpeer_name -j ACCEPT
    nsenter --net=$container_ns iptables -A OUTPUT -d $kasm_api_host_ip -o $vpeer_name -p tcp --dport $kasm_api_port -j ACCEPT

    nsenter --net=$container_ns iptables -A INPUT -i lo -j ACCEPT
    nsenter --net=$container_ns iptables -A INPUT -i tun0 -j ACCEPT
    nsenter --net=$container_ns iptables -A INPUT -i wg -j ACCEPT
    nsenter --net=$container_ns iptables -A INPUT -i tailscale0 -j ACCEPT
    nsenter --net=$container_ns iptables -A INPUT -s $bridge_gateway_ip -j ACCEPT

    nsenter --net=$container_ns iptables -A OUTPUT -o lo -j ACCEPT
    nsenter --net=$container_ns iptables -A OUTPUT -o tun0 -j ACCEPT
    nsenter --net=$container_ns iptables -A OUTPUT -o wg -j ACCEPT
    nsenter --net=$container_ns iptables -A OUTPUT -o tailscale0 -j ACCEPT
    nsenter --net=$container_ns iptables -A OUTPUT -d $bridge_gateway_ip -j ACCEPT

    # allow traffic to the VPN server
    vpn_server_ip=$(echo $script_data | base64 -d | jq -r '.server')
    if [ -n "$vpn_server_ip" ]; then
      nsenter --net=$container_ns iptables -A INPUT -s $vpn_server_ip -j ACCEPT
      nsenter --net=$container_ns iptables -A OUTPUT -d $vpn_server_ip -j ACCEPT
    fi

    for dns_server in $dns_servers; do
      nsenter --net=$container_ns iptables -A INPUT -s $dns_server -j ACCEPT
      nsenter --net=$container_ns iptables -A OUTPUT -d $dns_server -j ACCEPT
    done

    nsenter --net=$container_ns iptables -A INPUT -j DROP
    nsenter --net=$container_ns iptables -A OUTPUT -j DROP
  fi

  # configure the interface
  info "Configuring network interface for container."
  bridge_subnet_mask=$(echo $bridge_gateway_ip | cut -d/ -f2)
  nsenter --net=$container_ns ip link set $vpeer_name up
  nsenter --net=$container_ns ip addr add $container_ip/$bridge_subnet_mask dev $vpeer_name broadcast $bridge_broadcast
  nsenter --net=$container_ns ip route add default via $bridge_ip dev $vpeer_name

  if [ "$kasm_api_host_ip" != "$proxy_ip" ]; then
    info "Adding route for $kasm_api_host_ip via $bridge_ip."
    nsenter --net=$container_ns ip route add $kasm_api_host_ip via $bridge_ip dev $vpeer_name
  fi

  # run the notification script
  info "Running notification script in container."
  notification_pipe="/var/run/kasm-sidecar/$container_ns_id/notifications"
  mkfifo $notification_pipe
  ${SCRIPT_DIR}/notification_bridge "${kasm_id}" "${container_name}" "${container_ns_id}" "${notification_pipe}" &
  echo $! >> "/var/run/kasm-sidecar/$container_ns_id/pids"

  # run the script
  info "Running script \"$script_file\" in container."
  ${SCRIPT_DIR}/${script_file} "${kasm_id}" "${container_name}" "${container_ns_id}" "${container_pid}" "${script_data}" "${notification_pipe}" &
  echo $! >> "/var/run/kasm-sidecar/$container_ns_id/pids"
}

container_ns=$1
bridge_name=$2
bridge_gateway_ip=$3
kasm_id=$4
script_file=$5
script_data=$6

if [ -z "$container_ns" ] || [ -z "$bridge_name" ] || [ -z "$bridge_gateway_ip" ]; then
  info "Missing parameters. Got $@. Exiting."
  exit 1
fi

bridge_ip=$(echo $bridge_gateway_ip | cut -d. -f1-3).1
bridge_broadcast=$(echo $bridge_gateway_ip | cut -d. -f1-2).255.255
container_ns_id=$(echo $container_ns | grep -o -E "[a-z0-9]+$")

if [ -z "$kasm_id" ]; then
  kasm_id="kasm_proxy"
fi

info "Setting up network namespace for \"$kasm_id\" (/var/run/kasm-sidecar/$container_ns_id)" 
mkdir -p /var/run/netns
mkdir -p $LOG_DIR/$container_ns_id
mkdir -p /var/run/kasm-sidecar/$container_ns_id
echo "$$" >> "/var/run/kasm-sidecar/$container_ns_id/pids"

wait_for_container
create_bridge_interface

if [[ "$container_name" == "kasm_proxy" ]]; then
  configure_proxy_container
elif [[ -n "$container_name" ]]; then
  configure_session_container
fi