#!/bin/sh

# auto-apt-proxy - automatic detector of common APT proxy settings
# Copyright (C) 2016-2020 Antonio Terceiro
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

set -eu

uid=$(id -u)
cache_dir=${TMPDIR:-/tmp}/.auto-apt-proxy-${uid}
if [ -d "${cache_dir}" ]; then
  # require existing cache dir to be owned by the current user and have the
  # correct permissions
  owner_and_mode="$(stat --format=%u:%f "${cache_dir}")"
  if [ "${owner_and_mode}" != "${uid}:41c0" ]; then
    echo "E: insecure cache dir ${cache_dir}. Must be owned by UID ${uid} and have permissions 700" >&2
    exit 1
  fi
elif [ -z "${AUTO_APT_PROXY_NO_CACHE:-}" ]; then
  mkdir -m 0700 "${cache_dir}"
fi

output="${cache_dir}/output"
cleanup() {
  rm -f "$output"
}
trap cleanup INT EXIT TERM

proxy_url() {
  case "$1" in
    *:*)
      echo "http://[$1]:$2"
    ;;
    *)
      echo "http://$1:$2"
    ;;
  esac
}

hit() {
  timeout 5 /usr/lib/apt/apt-helper \
    -o Acquire::http::Proxy=DIRECT -o Acquire::Retries=0 \
    download-file "$@" "$output" 2>&1
}

cache_ttl=60 # seconds
cache() {
  local cache_file="${cache_dir}/cache"
  local lock_file="${cache_dir}/lock"
  local cache_age
  (
    flock 9

    # invalidate stale cache
    if [ -f "${cache_file}" ]; then
      ts=$(stat --format=%Y "${cache_file}")
      now=$(date +%s)
      cache_age=$((now - ts))
      if [ "${cache_age}" -gt "${cache_ttl}" ]; then
        rm -f "${cache_file}"
      fi
    fi

    if [ -f "${cache_file}" ]; then
      # read cache
      if [ -s "${cache_file}" ]; then
        cat "${cache_file}"
      fi
    else
      # update cache
      "$@" > "$cache_file" || true
      cat "${cache_file}"
    fi
  ) 9> "${lock_file}"
}

detect_apt_cacher() {
  local ip="$1"
  local proxy="$(proxy_url "$ip" 3142)"
  hit -o "Acquire::http::Proxy::${ip}=DIRECT" "$proxy" >/dev/null 2>&1 || true;
  if [ -s "$output" ] && grep -q -i '<title>Apt-cacher' "$output"; then
    echo "$proxy"
    return 0
  fi
  return 1
}

detect_apt_cacher_ng() {
  local ip="$1"
  local proxy="$(proxy_url "$ip" 3142)"
  if hit -o "Acquire::http::Proxy::${ip}=DIRECT" "$proxy" | grep -q -i 'HTTP.*406'; then
    echo "$proxy"
    return 0
  fi
  return 1
}

detect_approx() {
  local ip="$1"
  local proxy="$(proxy_url "$ip" 9999)"
  hit -o "Acquire::http::Proxy::${ip}=DIRECT" "$proxy" >/dev/null 2>&1 || true;
  if [ -s "$output" ] && grep -q -i '<title>approx\s*server</title>' "$output"; then
    echo "$proxy"
    return 0
  fi
  return 1
}

# NOTE: This does NOT check MDNS/DNS-SD (avahi/zeroconf/bonjour) records.
#       If you want that, use squid-deb-proxy-client, which depends on avahi.
detect_squid_deb_proxy() {
  local ip="$1"
  local proxy="$(proxy_url "$ip" 8000)"
  if hit -oDebug::acquire::http=1 -o "Acquire::http::Proxy::${ip}=DIRECT" "$proxy" 2>&1 | grep -q 'Via: .*squid-deb-proxy'; then
    echo "$proxy"
    return 0
  fi
  return 1
}

get_search_domains() {
  awk '{ if ($1 == "search") { $1=""; print($0) } }' /etc/resolv.conf 2>/dev/null
}

# NOTE: This does NOT check MDNS/DNS-SD (avahi/zeroconf/bonjour) records.
#       If you want that, use squid-deb-proxy-client, which depends on avahi.
#
# FIXME: if there are multiple matching SRV records, we should make a
#        weighted random choice from the one(s) with the highest priority.
#        For now, we make a uniformly random choice from all records (shuf + exit).
#
# NOTE: We don't check that it "looks like" a known apt proxy (hit + grep -q).
#       This is because
#        1) the other detectors are just GUESSING hosts and ports.
#           You might accidentally run a non-apt-proxy on 127.0.0.1:9999, but
#           you can't accidentally create an _apt_proxy SRV record!
#        2) refactoring the grep -q's out of detect_* is tedious and boring.
#        3) there's no grep -q for squid, which I want to use. ;-)
#
# NOTE: no need for if/then/else and return 0/1 because:
#        * if awk matches something, it prints it and exits zero.
#        * if hostname or apt-helper fail, awk matches nothing, so exits non-zero.
#        * set -e ignores errors from apt-helper (no pipefail) and hostname (no ???).
detect_DNS_SRV_record() {
  local domain
  domain="$(hostname --domain 2>/dev/null)"
  search_domains=$(get_search_domains)
  for domain in $domain $search_domains; do
    result=$(
      /usr/lib/apt/apt-helper srv-lookup _apt_proxy._tcp."${domain}" 2>/dev/null |
      shuf |
      awk '/^[^#]/{print "http://" $1 ":" $4;found=1;exit}END{exit !found}'
    )
    if [ -n "${result}" ]; then
      echo "${result}"
      return 0
    fi
  done
  return 1
}

detect_avahi_local() {
  local data iptype hostname ip port url check
  if ! command -v avahi-browse >/dev/null 2>/dev/null; then
    return 1
  fi

  data=$(avahi-browse --terminate --no-db-lookup --resolve --parsable _apt_proxy._tcp | awk -F ';' '/^=;/ { print($7, $8, $9); exit}')
  if [ -z "${data}" ]; then
    return 1
  fi

  ip=$(echo "${data}" | awk '{print($2)}')
  port=$(echo "${data}" | awk '{print($3)}')

  if [ -n "${AUTO_APT_PROXY_AVAHI_NAME:-}" ]; then
    hostname="$(echo "${data}" | awk '{print($1)}')"
    proxy="http://${hostname}:${port}"
  else
    proxy="$(proxy_url "${ip}" "${port}")"
  fi

  check="$(LANG=C.UTF-8 hit -o "Acquire::http::Proxy::${ip}=DIRECT" "${proxy}" || true)"
  if echo "${check}" | grep -q '111: Connection refused\|101: Network is unreachable'; then
    return 1
  fi

  echo "${proxy}"
}


if command -v ip >/dev/null; then
  ip="ip"
elif busybox ip --help >/dev/null 2>&1; then
  ip="busybox ip"
else
  ip="true"
fi

find_gateway() {
  local gateway
  $ip "$@" route | awk '/default/ { print($3) }'
}

resolve_getent() {
  timeout 5 getent "$@" | awk '/[:blank:]/ { if ($2 == "STREAM") {print($1)} }' | uniq
}

__detect__() {
  v4_gateway=$(find_gateway)
  v6_gateway=$(find_gateway -6)

  # consider a user-defined host as well, quick check whether it's configured
  v4_explicit_proxy=$(resolve_getent ahostsv4 apt-proxy)
  v6_explicit_proxy=$(resolve_getent ahostsv6 apt-proxy)

  v4_addresses=$($ip -4 route | awk '{if($10 == "linkdown") {print($9)}}')
  # TODO add v6 addessses as well

  for ip in 127.0.0.1 $v4_addresses $v4_explicit_proxy $v6_explicit_proxy $v4_gateway ::1 $v6_gateway; do
    detect_apt_cacher_ng "$ip"   && return 0
    detect_approx "$ip"          && return 0
    detect_apt_cacher "$ip"      && return 0
    detect_squid_deb_proxy "$ip" && return 0
  done

  # If a SRV record is found, use it and guess no further.
  detect_DNS_SRV_record && return 0
  detect_avahi_local && return 0

  return 0
}

detect() {
  if [ -z "${AUTO_APT_PROXY_NO_CACHE:-}" ]; then
    cache __detect__
  else
    __detect__
  fi
}

if [ $# -eq 0 ]; then
  detect
else
  case "$1" in
    ftp://*|http://*|https://*|file://*)
      # APT mode: first argument is an URI
      detect
      ;;
    *)
      # wrapper mode: execute command using the detected proxy
      proxy=$(detect || true)
      if [ -n "$proxy" ]; then
        export http_proxy="$proxy"
        export HTTP_PROXY="$proxy"
      fi
      exec "$@"
  esac
fi
