#!/bin/sh
# StudyLock's browser-rules helper. Runs as root through pkexec (polkit action
# com.abrar.studylock.browser-rules, allowed for the person at the computer without a password).
#
# It only ever writes or removes StudyLock's own browser policy files:
#   Chrome, Chromium, Edge, Brave, Vivaldi: <managed policy dir>/studylock.json
#   Firefox: /etc/firefox/policies/policies.json, only if there is none or StudyLock wrote it
# and every site entry is checked against a strict character set before anything is written.
#
#   studylock-policy check
#   studylock-policy apply [--chromium FILTER]... [--firefox PATTERN]...
#   studylock-policy clear
set -euf
umask 022
PATH=/usr/sbin:/usr/bin:/sbin:/bin

NAME=studylock.json
FIREFOX_DIR=/etc/firefox/policies
FIREFOX_MARK=$FIREFOX_DIR/.studylock
NL='
'

# managed-policy directory : what shows that browser is installed
CHROMIUM="/etc/opt/chrome/policies/managed:/opt/google/chrome
/etc/chromium/policies/managed:/usr/lib/chromium
/etc/chromium/policies/managed:/usr/lib64/chromium-browser
/etc/chromium-browser/policies/managed:/usr/lib/chromium-browser
/var/snap/chromium/current/policies/managed:/var/snap/chromium/current
/etc/opt/edge/policies/managed:/opt/microsoft/msedge
/etc/brave/policies/managed:/opt/brave.com/brave
/etc/opt/vivaldi/policies/managed:/opt/vivaldi"

die() { echo "studylock-policy: $*" >&2; exit 2; }

valid() {
    case "$1" in *"$NL"*|"") return 1 ;; esac
    [ "${#1}" -le 300 ] || return 1
    printf '%s' "$1" | grep -Eq '^[A-Za-z0-9._~:/?#@!$&()*+,;=%-]+$'
}

json_list() {   # "a" "b" -> "a", "b"
    first=1
    for v in "$@"; do
        if [ $first = 1 ]; then first=0; else printf ', '; fi
        printf '"%s"' "$v"
    done
}

firefox_installed() {
    [ -e /usr/lib/firefox ] || [ -e /usr/lib64/firefox ] || [ -e /usr/lib/firefox-esr ] || [ -e /snap/firefox ] || [ -e /opt/firefox ]
}

write_file() {   # path content: atomically, root-owned, world-readable
    dir=$(dirname "$1")
    mkdir -p "$dir"
    tmp=$(mktemp "$dir/.studylock.XXXXXX")
    printf '%s\n' "$2" > "$tmp"
    chmod 644 "$tmp"
    mv -f "$tmp" "$1"
}

clear_all() {
    printf '%s\n' "$CHROMIUM" | while IFS=: read -r dir marker; do
        rm -f "$dir/$NAME"
    done
    if [ -e "$FIREFOX_MARK" ]; then
        rm -f "$FIREFOX_DIR/policies.json" "$FIREFOX_MARK"
    fi
}

[ "$(id -u)" = 0 ] || die "must run as root (through pkexec)"
cmd=${1:-}
[ $# -gt 0 ] && shift

case "$cmd" in
check)
    echo ok
    ;;
clear)
    clear_all
    echo cleared
    ;;
apply)
    chromium=""; firefox=""; nc=0; nf=0
    while [ $# -gt 0 ]; do
        case "$1" in
        --chromium) [ $# -ge 2 ] || die "missing value"; valid "$2" || die "rejected entry"; nc=$((nc + 1)); chromium="$chromium$2$NL"; shift 2 ;;
        --firefox)  [ $# -ge 2 ] || die "missing value"; valid "$2" || die "rejected entry"; nf=$((nf + 1)); firefox="$firefox$2$NL"; shift 2 ;;
        *) die "unknown argument" ;;
        esac
    done
    [ $nc -le 400 ] && [ $nf -le 400 ] || die "too many entries"

    oldifs=$IFS; IFS=$NL
    # shellcheck disable=SC2086
    allow=$(json_list $chromium)
    # shellcheck disable=SC2086
    exceptions=$(json_list $firefox)
    IFS=$oldifs

    chrome_json=$(printf '{\n  "URLBlocklist": ["*"],\n  "URLAllowlist": [%s]\n}' "$allow")
    written=0
    printf '%s\n' "$CHROMIUM" | {
        while IFS=: read -r dir marker; do
            if [ -e "$marker" ]; then write_file "$dir/$NAME" "$chrome_json"; fi
        done
    }
    if firefox_installed; then
        if [ ! -e "$FIREFOX_DIR/policies.json" ] || [ -e "$FIREFOX_MARK" ]; then
            ff_json=$(printf '{\n  "policies": {\n    "WebsiteFilter": { "Block": ["<all_urls>"], "Exceptions": [%s] },\n    "Preferences": { "browser.sessionstore.resume_session_once": { "Value": true, "Status": "user" } }\n  }\n}' "$exceptions")
            write_file "$FIREFOX_DIR/policies.json" "$ff_json"
            : > "$FIREFOX_MARK"
        else
            echo "firefox: /etc/firefox/policies/policies.json belongs to someone else; left alone" >&2
            exit 3
        fi
    fi
    echo applied
    ;;
*)
    die "usage: studylock-policy check | apply [--chromium F]... [--firefox P]... | clear"
    ;;
esac
