#!/usr/bin/env bash
# tests/unit/06_entrypoint.sh -- the container-side enforcement script.
source "$TESTS/lib.sh"

PROJ="$WORK/proj"; mkdir -p "$PROJ"; cd "$PROJ" || exit 1
export AISB_CONFIG_DIR="$WORK/cfg-empty"; mkdir -p "$AISB_CONFIG_DIR"

# Extract the entrypoint exactly as it is baked into the image.
EP="$WORK/entrypoint.sh"
"$AISB" --print-entrypoint > "$EP" 2>"$WORK/ep.err"; RC=$?
assert_eq 0 "$RC" "the embedded entrypoint can be printed"
bash -n "$EP" && _ok "entrypoint is syntactically valid bash" || _bad "entrypoint is syntactically valid bash"
grep -q 'set -Eeuo pipefail' "$EP" && _ok "entrypoint is strict" || _bad "entrypoint is strict"

t_name "the image is described by the script itself"
sb --print-containerfile
assert_contains "FROM docker.io/library/ubuntu:24.04" "ubuntu 24.04 base"
assert_contains "nftables" "image ships nftables"
assert_contains "util-linux" "image ships setpriv (util-linux)"
assert_contains "sha256sum -c" "node tarball is checksum verified"
assert_contains "useradd" "image creates the sandbox user"
assert_contains 'ENTRYPOINT ["/usr/local/lib/aisb/entrypoint.sh"]' "entrypoint is baked in"
assert_contains "WORKDIR /project" "project is the default workdir"

t_name "generated ruleset shape"
# Run the rule generator outside a container (nft is not needed to render it).
POL="$WORK/mnt"; mkdir -p "$POL"
cat > "$POL/policy.env" <<EOF
AISB_NET=allow
AISB_UID=1000
AISB_GID=1000
AISB_UNENFORCED=0
AISB_KEEP_ROOT=0
AISB_ALLOW_ICMP=1
AISB_ALLOW_DNS=0
AISB_PROXY_IPS=100.64.0.1
AISB_PROXY_PORTS=18081
EOF
printf '10.0.0.5 5432\n10.1.0.0/16 80,443\n' > "$POL/direct4.list"
printf '::1 8080\n' > "$POL/direct6.list"

RULES="$AISB_STATE_DIR/rules.f"
# Render the ruleset by extracting the generator and feeding it the policy above.
# Render the ruleset the way the entrypoint would. publish_rules() is a function of its own
# in there, so it comes along: extracting only gen_rules() would ask a half-program to run.
render_rules() { # render_rules <AISB_PUBLISH_PORTS value> -> the ruleset on stdout
  local g="$WORK/render.$$.sh"
  { awk '/^publish_rules\(\) \{/{f=1} f{print} f&&/^\}$/{exit}' "$EP"
    awk '/^gen_rules\(\) \{/{f=1} f{print} f&&/^\}$/{exit}' "$EP"
    printf 'POLICY_DIR=%q\nAISB_PROXY_IPS="100.64.0.1"\nAISB_PROXY_PORTS="18081"\nAISB_ALLOW_ICMP=1\nAISB_ALLOW_DNS=0\nAISB_PUBLISH_PORTS=%q\n' "$POL" "$1"
    printf 'gen_rules\n'; } > "$g"
  bash "$g" 2>"$WORK/render.err"
}

GEN="$WORK/gen.sh"
awk '/^publish_rules\(\) \{/{f=1} f{print} f&&/^\}$/{exit}' "$EP" > "$GEN"
awk '/^gen_rules\(\) \{/{f=1} f{print} f&&/^\}$/{exit}' "$EP" >> "$GEN"
{ printf 'POLICY_DIR=%q\nAISB_PROXY_IPS="100.64.0.1"\nAISB_PROXY_PORTS="18081"\nAISB_ALLOW_ICMP=1\nAISB_ALLOW_DNS=0\nAISB_PUBLISH_PORTS=""\n' "$POL"
  printf 'gen_rules\n'; } >> "$GEN"
bash "$GEN" > "$RULES" 2>"$WORK/gen.err" && _ok "rules render" || _bad "rules render" "$(tail -3 "$WORK/gen.err")"

grep -q 'table inet aisb' "$RULES"            && _ok "inet table (covers v4 and v6)"    || _bad "inet table"
grep -q 'policy drop' "$RULES"                && _ok "default policy is drop"           || _bad "default policy drop"
grep -q 'aisb-default-drop' "$RULES"          && _ok "explicit default-drop rule exists" || _bad "default-drop rule"
grep -q 'oifname "lo" accept' "$RULES"        && _ok "loopback stays usable"             || _bad "loopback accept"
grep -q 'ip daddr 10.0.0.5 tcp dport { 5432 }' "$RULES" && _ok "per-entry host:port rule" || _bad "host:port rule: $(grep -c . "$RULES") lines"
grep -q 'ip daddr 10.1.0.0/16 tcp dport { 80,443 }' "$RULES" && _ok "CIDR rule with default ports" || _bad "CIDR rule"
grep -q 'ip6 daddr ::1 tcp dport { 8080 }' "$RULES"          && _ok "ipv6 rule"                    || _bad "ipv6 rule"
grep -q 'ip daddr 100.64.0.1 tcp dport 18081' "$RULES"       && _ok "gate endpoint is allowed"      || _bad "gate rule"
grep -q 'aisb-est' "$RULES"                   && _ok "established traffic accepted"      || _bad "established rule"
grep -c 'accept comment "aisb-direct"' "$RULES" | grep -qx 3 && _ok "exactly 3 direct rules" || _bad "direct rule count: $(grep -c 'aisb-direct' "$RULES")"

t_name "the published ports appear in the input chain, and nowhere else"
# Nothing published means no ingress rule at all: the ruleset must not open a port that was
# never asked for, and an empty set must not render an empty port list (nft refuses those).
PUBRULES="$(render_rules "8000 18080 8000")"
# Is a line matching <regex> present inside the named chain? An ingress rule in the output
# chain would be an egress hole, which is the one thing this feature must never become.
in_chain() { # in_chain <chain> <regex>
  awk -v c="chain $1 {" -v re="$2" 'index($0, c){inside=1; next} /^  }/{inside=0}
       inside && $0 ~ re {found=1} END{exit !found}' <<< "$PUBRULES"
}
printf '%s' "$PUBRULES" | grep -qF 'comment "aisb-publish"' && _ok "a rule names the published ports" || _bad "aisb-publish rule" "$(printf '%s' "$PUBRULES" | sed -n '/chain input/,/}/p')"
printf '%s' "$PUBRULES" | grep -qF 'tcp dport { 8000,18080 }' && _ok "both ports, de-duplicated" || _bad "port set" "$(printf '%s' "$PUBRULES" | grep aisb-publish)"
in_chain input 'aisb-publish'  && _ok "the publish rule is in the input chain" || _bad "the publish rule is in the input chain"
in_chain output 'aisb-publish' && _bad "the publish rule is also an egress hole" || _ok "publishing opens no egress path"
in_chain input 'policy drop'   && _ok "input still defaults to drop with a port open" || _bad "input policy drop survived"
printf '%s' "$PUBRULES" | grep -qF 'aisb-default-drop' && _ok "the egress default-drop is untouched" || _bad "default-drop rule"
# The policy file is host-written and mounted read-only, but a ruleset is not somewhere to
# splice an unvalidated string: anything that is not a port is dropped on the floor.
BADRULES="$(render_rules '8000 nft delete table inet aisb accept')"
printf '%s' "$BADRULES" | grep -qF 'delete table' && _bad "a port value reached the ruleset as code" || _ok "a non-numeric port never reaches the ruleset"
printf '%s' "$BADRULES" | grep -qF 'tcp dport { 8000 }' && _ok "  ... and the valid port in it still lands" || _bad "valid port kept" "$(printf '%s' "$BADRULES" | grep aisb-publish)"
render_rules '' | grep -qF 'aisb-publish' && _bad "a rule appeared with nothing published" || _ok "no publish rule when nothing is published"

t_name "no DNS hole when the gate is in front"
grep -E 'th dport 53' "$RULES" >/dev/null && _bad "DNS is closed in gate mode" || _ok "DNS is closed in gate mode"

t_name "the seeding script survives being passed in single quotes"
# Everything between `bash -c '` and `' aisb-seed` is one single-quoted argument of the
# outer script. One quote inside it closes that argument early and the rest of the file
# silently becomes code in the *outer* shell -- valid bash, wrong program, so neither
# `bash -n` nor shellcheck notices. Hence this test.
INNER="$WORK/seeding-inner.sh"
awk "/setpriv .*bash -c '/{f=1;next} /^ *' aisb-seed/{f=0} f" "$EP" > "$INNER"
[[ -s "$INNER" ]] && _ok "the seeding script was extracted" || _bad "the seeding script was extracted"
if grep -q "'" "$INNER"; then
  _bad "it contains no single quote" "found: $(grep -n "'" "$INNER" | head -2 | tr '\n' ' ')"
else
  _ok "it contains no single quote"
fi
bash -n "$INNER" && _ok "and it is valid bash on its own" || _bad "and it is valid bash on its own"
grep -q '\$lvl' "$INNER" && _ok "the thinking level is enforced inside it" || _bad "thinking level enforced"

t_name "settings.json is merged per key, not managed whole"
# Run the merge exactly as the entrypoint does, with the program extracted from the image
# script: T = what the sandbox has, B = what we last wrote, N = what the host says now.
MJQ="$WORK/merge.jq"
awk '/^SETTINGS_MERGE_JQ=/{f=1;next} f && /^.$/{exit} f' "$EP" > "$MJQ"
[[ "$(wc -l < "$MJQ")" -gt 3 ]] && _ok "the merge program was extracted" || _bad "the merge program was extracted"
M_IN="$WORK/m"; mkdir -p "$M_IN"
cat > "$M_IN/new.json" <<'EOF'
{ "defaultProvider": "gateway", "defaultModel": "turbo", "theme": "dark/dark",
  "thinkingBudgets": { "xhigh": 65536 } }
EOF
cat > "$M_IN/wrote.json" <<'EOF'
{ "defaultProvider": "gateway", "defaultModel": "turbo", "theme": "dark/dark",
  "thinkingBudgets": { "xhigh": 65536 } }
EOF
cat > "$M_IN/sandbox.json" <<'EOF'
{ "defaultProvider": "gateway", "defaultModel": "turbo", "theme": "light/light",
  "thinkingBudgets": { "xhigh": 65536 }, "lastChangelogVersion": "1.4.0", "mine": 7 }
EOF
merge_out() { jq -c -n --slurpfile T "$1" --slurpfile B "$2" --slurpfile N "$3" -f "$MJQ"; }
M="$(merge_out "$M_IN/sandbox.json" "$M_IN/wrote.json" "$M_IN/new.json")"
[[ "$(printf '%s' "$M" | jq -r .theme)" == "light/light" ]] \
  && _ok "a key edited in the sandbox stays edited" || _bad "sandbox edit kept: $M"
[[ "$(printf '%s' "$M" | jq -r .lastChangelogVersion)" == "1.4.0" ]] \
  && _ok "a key only pi wrote survives" || _bad "sandbox-only key: $M"
[[ "$(printf '%s' "$M" | jq -r .mine)" == "7" ]] \
  && _ok "and so does one the sandbox invented" || _bad "invented key: $M"

# host moves two keys, drops a third, keeps one the sandbox changed
cat > "$M_IN/new2.json" <<'EOF'
{ "defaultProvider": "gateway", "defaultModel": "deepthink",
  "thinkingBudgets": { "xhigh": 131072 } }
EOF
M2="$(merge_out "$M_IN/sandbox.json" "$M_IN/wrote.json" "$M_IN/new2.json")"
[[ "$(printf '%s' "$M2" | jq -r .defaultModel)" == "deepthink" ]] \
  && _ok "an untouched key follows the host again" || _bad "host key arrives: $M2"
[[ "$(printf '%s' "$M2" | jq -r .thinkingBudgets.xhigh)" == "131072" ]] \
  && _ok "including a nested one (whole key, whole value)" || _bad "nested key: $M2"
[[ "$(printf '%s' "$M2" | jq -r 'has("theme")')" == "true" ]] \
  && _ok "a key the host dropped stays if the sandbox had changed it" || _bad "dropped-but-edited key: $M2"
# the same drop against a file the sandbox never touched: it goes away, as it did when
# the whole file was re-seeded -- otherwise settings only ever accumulate
cat > "$M_IN/clean.json" <<'EOF'
{ "defaultProvider": "gateway", "defaultModel": "turbo", "theme": "dark/dark",
  "thinkingBudgets": { "xhigh": 65536 } }
EOF
M3="$(merge_out "$M_IN/clean.json" "$M_IN/wrote.json" "$M_IN/new2.json")"
[[ "$(printf '%s' "$M3" | jq -r 'has("theme")')" == "false" ]] \
  && _ok "an untouched key the host dropped is dropped" || _bad "untouched dropped key: $M3"
# no base copy (the first run after this shipped, or .wrote deleted): the host wins on
# every key it has, which is what repairs a settings file pi had frozen.
M4="$(merge_out "$M_IN/sandbox.json" /dev/null "$M_IN/new.json")"
[[ "$(printf '%s' "$M4" | jq -r '.thinkingBudgets.xhigh')" == "65536" ]] \
  && _ok "with no base, host preferences come back (the frozen-file repair)" || _bad "no-base merge: $M4"
[[ "$(printf '%s' "$M4" | jq -r .lastChangelogVersion)" == "1.4.0" ]] \
  && _ok "and pi's own keys are kept" || _bad "no-base, pi key: $M4"
printf '{oops' > "$M_IN/broken.json"
jq -n --slurpfile T "$M_IN/broken.json" --slurpfile B /dev/null --slurpfile N "$M_IN/new.json" \
   -f "$MJQ" > /dev/null 2>&1 \
  && _bad "unparseable input fails loudly for the caller to handle" || _ok "unparseable input fails loudly for the caller to handle"

t_name "the sandbox keeps its own chatter off the stream the command owns"
# stdout inside the container is data. ai-sandbox-autophase reads its preflight answer from
# it and progress.mjs reads the child pi event stream from it, so every word the sandbox says
# about itself belongs on stderr. Five notices in this very block were once on stdout, and one
# of them fires on every pi-bridged run -- which made the driver refuse a working sandbox.
# Hence a test that runs the block and separates the two streams, not one that reads it.
if grep -nE '(echo|printf).*\[aisb\]' "$EP" | grep -v '>&2' | grep -q .; then
  _bad "every [aisb] line in the entrypoint is redirected" \
       "$(grep -nE '(echo|printf).*\[aisb\]' "$EP" | grep -v '>&2' | head -2 | tr '\n' ' ')"
else
  _ok "every [aisb] line in the entrypoint is redirected"
fi

# The extracted block, run as the entrypoint runs it (bash -c with the policy directory, the
# level and the merge program as $1..$3). Only this snippet sees the substituted HOME, and it
# starts no container, so podman's own HOME is untouched.
SP="$WORK/seed"; mkdir -p "$SP/pi"
printf '{ "providers": { "gateway": { "models": [ { "id": "turbo" } ] } } }\n' > "$SP/pi/models.json"
printf '{ "theme": "dark/dark" }\n'                                             > "$SP/pi/settings.json"
SEED_OUT="$WORK/seed.out"; SEED_ERR="$WORK/seed.err"
seed_run() { # <level>
  HOME="$SP/home" bash -c "$(cat "$INNER")" aisb-seed "$SP" "${1-xhigh}" "$(cat "$MJQ")" \
    > "$SEED_OUT" 2> "$SEED_ERR"
}
seed_says()  { grep -qF -- "$1" "$SEED_ERR"; }
seed_stdout() { [[ -s "$SEED_OUT" ]] && printf '%s' "$(tr '\n' ';' < "$SEED_OUT")"; }
seed_level() { jq -r '.defaultThinkingLevel // "absent"' "$SP/home/.pi/agent/settings.json" 2>/dev/null; }

seed_run xhigh; RC=$?
assert_eq 0 "$RC" "seeding runs outside a container"
[[ -z "$(seed_stdout)" ]] && _ok "and it writes nothing on stdout (cold seed)" \
                          || _bad "stdout belongs to the command (cold seed): $(seed_stdout)"
seed_says "seeded models.json"   && _ok "the notice went to stderr" || _bad "the notice went to stderr"
seed_says "seeded settings.json" && _ok "including the settings one" || _bad "including the settings one"
# the level is policy, so a home we create must not run one unit at the default of pi itself
assert_eq "xhigh" "$(seed_level)" "the first seed already carries the thinking policy"

seed_run xhigh; RC=$?
assert_eq 0 "$RC" "and on the next run"
[[ -z "$(seed_stdout)" ]] && _ok "nothing on stdout once the files exist (the steady state)" \
                          || _bad "steady-state stdout is not clean: $(seed_stdout)"
seed_says "refreshed models.json from the host config" \
  && _ok "the line that broke the driver is on stderr" || _bad "the line that broke the driver is on stderr"

sed -i 's/turbo/juggernaut/' "$SP/home/.pi/agent/models.json"
seed_run xhigh
[[ -z "$(seed_stdout)" ]] && _ok "an in-sandbox edit is reported without touching stdout" \
                          || _bad "edited-file stdout: $(seed_stdout)"
seed_says "left models.json alone" && _ok "and the explanation still arrives" || _bad "and the explanation still arrives"

printf '{oops' > "$SP/home/.pi/agent/settings.json"
seed_run xhigh
[[ -z "$(seed_stdout)" ]] && _ok "even a repair is announced off stdout" || _bad "broken-json stdout: $(seed_stdout)"
seed_says "not usable JSON" && _ok "with the warning on stderr" || _bad "with the warning on stderr"

seed_run low
assert_eq "low" "$(seed_level)" "a changed level is applied over the merged file"
seed_says "thinking -> low" && _ok "and announced there too" || _bad "and announced there too"
[[ -z "$(seed_stdout)" ]] && _ok "with stdout still empty" || _bad "level-change stdout: $(seed_stdout)"

t_name "privilege drop is the last thing that happens"
tail -5 "$EP" | grep -q 'setpriv' && _ok "exec setpriv at the end" || _bad "setpriv last: $(tail -3 "$EP" | tr '\n' ' ')"
grep -q -- '--bounding-set=-all' "$EP"  && _ok "bounding set cleared"   || _bad "bounding set cleared"
grep -q -- '--inh-caps=-all' "$EP"      && _ok "inheritable caps cleared" || _bad "inheritable caps cleared"
grep -q -- '--no-new-privs' "$EP"       && _ok "no new privileges"        || _bad "no-new-privs"
grep -q -- '--reuid=' "$EP"             && _ok "uid switched"             || _bad "reuid"
grep -q 'fail "policy file' "$EP"       && _ok "missing policy = refuse to start" || _bad "fail closed on missing policy"
grep -q 'refusing to start\|Refusing to run unenforced' "$EP" && _ok "missing NET_ADMIN = refuse to start" || _bad "fail closed without NET_ADMIN"