#!/bin/bash
# Installs the contest firewall without Ansible, on the booted contest image
# (Ubuntu 26.04, as in the ICPC image). Run it as root from this directory.
#
#   sudo ./install_firewall.sh            # install, the firewall starts at the next boot
#   sudo ./install_firewall.sh --start    # install and start it now
set -euo pipefail
cd "$(dirname "$0")"

if [ "$(id -u)" -ne 0 ]; then
    echo "Run this script as root (sudo)." >&2
    exit 1
fi
START=no
[ "${1:-}" = "--start" ] && START=yes
TIMEZONE=${TIMEZONE:-America/Moncton}

entries() {
    sed -e 's/#.*//' -e 's/[[:space:]]//g' "$1" | grep -v '^$'
}

echo "Installing nftables, dig, chrony and the nginx stream module"
apt-get update
apt-get install -y nftables dnsutils chrony nginx libnginx-mod-stream dnsmasq-base
id contest-proxy >/dev/null 2>&1 || useradd --system --no-create-home --shell /usr/sbin/nologin contest-proxy

echo "Copying the lists, scripts and systemd units"
install -d -m 0755 /etc/contest-firewall
install -m 0644 allowed_domains.txt dns_servers.txt ntp_servers.txt files/contest-proxy.conf /etc/contest-firewall/
install -m 0500 files/firewall.sh files/refresh_rules.sh files/enable_firewall.sh files/disable_firewall.sh /
install -m 0755 test_firewall.sh /
install -m 0644 files/firewall.service files/refresh_rules.service files/refresh_rules.timer files/contest-proxy.service files/contest-dns.service /etc/systemd/system/

echo "Configuring DNS (systemd-resolved)"
dns_list=$(entries dns_servers.txt | tr '\n' ' ' | sed 's/ $//')
if grep -qE '^#?DNS=' /etc/systemd/resolved.conf; then
    sed -i -E "s|^#?DNS=.*|DNS=${dns_list}|" /etc/systemd/resolved.conf
else
    echo "DNS=${dns_list}" >> /etc/systemd/resolved.conf
fi
if grep -qE '^#?DNSStubListener=' /etc/systemd/resolved.conf; then
    sed -i -E 's|^#?DNSStubListener=.*|DNSStubListener=no|' /etc/systemd/resolved.conf
else
    echo "DNSStubListener=no" >> /etc/systemd/resolved.conf
fi

echo "Configuring NTP (chrony) and the timezone"
sed -i -E '/^# (BEGIN|END) CONTEST FIREWALL NTP SERVERS$/,/^# END CONTEST FIREWALL NTP SERVERS$/d' /etc/chrony/chrony.conf
sed -i -E '/^(pool|server)[[:space:]]/d' /etc/chrony/chrony.conf
{
    echo "# BEGIN CONTEST FIREWALL NTP SERVERS"
    for ntp in $(entries ntp_servers.txt); do
        echo "pool ${ntp} iburst"
    done
    echo "# END CONTEST FIREWALL NTP SERVERS"
    cat /etc/chrony/chrony.conf
} > /etc/chrony/chrony.conf.new
mv /etc/chrony/chrony.conf.new /etc/chrony/chrony.conf
# Ubuntu 26.04 adds NTS time sources that the firewall blocks, and chrony
# does not use the plain NTP servers while those sources are configured
if [ -f /etc/chrony/sources.d/ubuntu-ntp-pools.sources ]; then
    mv /etc/chrony/sources.d/ubuntu-ntp-pools.sources /etc/chrony/sources.d/ubuntu-ntp-pools.sources.disabled
fi
# The contest machines may also run Windows, which keeps the hardware clock
# in local time: chrony must not write its time into that clock
sed -i -E '/^rtcsync/d' /etc/chrony/chrony.conf
grep -q '^makestep' /etc/chrony/chrony.conf || echo 'makestep 1 3' >> /etc/chrony/chrony.conf
# Fallback: time from the contest site if NTP does not work
install -m 0500 files/time/time_from_contest.sh /
install -m 0644 files/time/time_from_contest.service files/time/time_from_contest.timer /etc/systemd/system/
ln -sf "/usr/share/zoneinfo/${TIMEZONE}" /etc/localtime

echo "Blocking open.kattis.com"
grep -q 'open\.kattis\.com' /etc/hosts || printf '127.0.0.1\topen.kattis.com\n' >> /etc/hosts

# Firefox must send the real site name, which encrypted client hello (ECH) hides
POLICIES=/usr/lib/firefox/distribution/policies.json
if [ -f "$POLICIES" ]; then
    echo "Turning off encrypted client hello in Firefox"
    python3 - "$POLICIES" <<'PY'
import json, sys
path = sys.argv[1]
data = json.load(open(path))
prefs = data.setdefault("policies", {}).setdefault("Preferences", {})
for name in ("network.dns.echconfig.enabled", "network.dns.http3_echconfig.enabled"):
    prefs[name] = {"Value": False, "Status": "locked"}
json.dump(data, open(path, "w"), indent=2)
PY
fi

echo "Restarting services"
systemctl restart systemd-resolved
systemctl enable chrony
systemctl restart chrony
timedatectl set-ntp true || true
systemctl daemon-reload
systemctl enable firewall.service refresh_rules.service refresh_rules.timer time_from_contest.timer
systemctl start refresh_rules.timer

if [ "$START" = yes ]; then
    systemctl start firewall.service
    /refresh_rules.sh
    echo "Firewall installed and active."
else
    echo "Firewall installed. It will start at the next boot (or run: sudo /enable_firewall.sh)."
fi