#!/usr/bin/env bash
# zeek-rita.sh - import Zeek logs into RITA every 15 minutes
# Usage: sudo ./zeek-rita.sh enable|disable|status

set -euo pipefail

SERVICE_UNIT=/etc/systemd/system/rita-import.service
TIMER_UNIT=/etc/systemd/system/rita-import.timer

# Re-exec with sudo if not root
if [[ ${EUID} -ne 0 ]]; then
    exec sudo -E "$0" "$@"
fi

usage() {
    echo "Usage: $0 {enable|disable|status}"
    exit 2
}

enable() {
    if [[ ! -x /usr/local/bin/rita || ! -x /opt/zeek/bin/zeekctl ]]; then
        echo "[!] rita (/usr/local/bin/rita) and zeek (/opt/zeek/bin/zeekctl) must be installed first" >&2
        echo "[!] run ../rita.sh add and ../zeek.sh add first" >&2
        exit 1
    fi

    echo "[*] Creating ${SERVICE_UNIT}"
    cat > "${SERVICE_UNIT}" <<'EOF'
[Unit]
Description=Import Zeek logs into RITA
After=zeek.service

[Service]
Type=oneshot
# rita import exits 1 on the idempotent no-op ("all files were previously
# imported"); tolerate exactly that exit code so the timer stays green.
ExecStart=/bin/sh -c '/usr/bin/flock -n /run/rita-import.lock /usr/local/bin/rita import --database shaula --logs=/tmp/zeek_logs /opt/zeek/logs || [ $? -eq 1 ]'
EOF

    echo "[*] Creating ${TIMER_UNIT}"
    cat > "${TIMER_UNIT}" <<'EOF'
[Unit]
Description=Run rita-import every 15 minutes

[Timer]
OnCalendar=*:0/15
Persistent=true

[Install]
WantedBy=timers.target
EOF

    systemctl daemon-reload
    systemctl enable --now rita-import.timer
    echo "[+] rita-import.timer enabled (every 15 min)"
}

disable() {
    echo "[*] Stopping and disabling rita-import timer/service"
    systemctl disable --now rita-import.timer 2>/dev/null || true
    systemctl stop rita-import.service 2>/dev/null || true
    systemctl disable rita-import.service 2>/dev/null || true

    echo "[*] Removing units"
    rm -f "${SERVICE_UNIT}" "${TIMER_UNIT}"
    systemctl daemon-reload

    echo "[+] rita-import disabled"
}

status() {
    echo "[*] rita-import.timer:"
    systemctl is-active rita-import.timer 2>/dev/null && systemctl is-enabled rita-import.timer 2>/dev/null || echo "inactive/not enabled"
    echo "[*] Last import service run:"
    systemctl status rita-import.service --no-pager -n 5 2>/dev/null || true
}

case "${1:-}" in
    enable)  enable ;;
    disable) disable ;;
    status)  status ;;
    *)       usage ;;
esac