#!/bin/bash
# uninstall_clean.sh
# Full SafeDNS cleanup from the machine.
# Intended to be run as root via MDM script.

set -uo pipefail   # NO -e flag — rm failures do not abort the script

LOG_FILE="/var/log/safedns-uninstall.log"
STATUS_FILE="/Library/Group Containers/group.com.safedns.agent.macos.enterprise/MDMStatus.txt"

log_file() {
    echo "$(date '+%Y-%m-%d %H:%M:%S') [$1] $2" >> "$LOG_FILE"
}
log_info() { echo "== $1"; log_file "INFO" "$1"; }
log_ok()   { echo "✔  $1"; log_file  "OK"  "$1"; }
log_warn() { echo "⚠  $1"; log_file "WARN" "$1"; }
log_err()  { echo "✘  $1"; log_file "ERR" "$1"; }

if [[ $EUID -ne 0 ]]; then
    log_err "Run as root: sudo bash $0"
    exit 1
fi

SYSEX_BUNDLE="com.safedns.agent.SafeDNSMacProxy.dnsProxy"
DAEMON_LABEL="com.safedns.daemon"
DAEMON_PLIST="/Library/LaunchDaemons/${DAEMON_LABEL}.plist"
HOST_AGENT_LABEL="com.safedns.agent.SafeDNSMacProxy.host"
HOST_AGENT_PLIST="/Library/LaunchAgents/${HOST_AGENT_LABEL}.plist"
APP_PATH="/Applications/SafeDNSMacProxy.app"
SYSEX_DIR="/Library/SystemExtensions"

echo ""
log_info "SafeDNS — full cleanup"
echo "────────────────────────────────────────"

###
# 0. Deactivate DNS Proxy (user-space NE configuration)
###
log_info "0. Deactivating DNS Proxy (Network Extension)..."

CONSOLE_USER=$(stat -f "%Su" /dev/console 2>/dev/null || echo "")
if [[ -z "${CONSOLE_USER}" || "${CONSOLE_USER}" == "root" ]]; then
    log_warn "Failed to detect console user — skipping DNS Proxy deactivation"
else
    CONSOLE_UID=$(id -u "${CONSOLE_USER}" 2>/dev/null || echo "")
    if [[ -n "${CONSOLE_UID}" ]]; then
        DEACTIVATE_SCRIPT="/tmp/safedns_deactivate.sh"
        cat > "${DEACTIVATE_SCRIPT}" << 'INNER_EOF'
#!/bin/bash
NE_USER_PLIST="${HOME}/Library/Preferences/com.apple.networkextension.plist"
BUNDLE="com.safedns.agent.SafeDNSMacProxy.dnsProxy"

if [[ -f "${NE_USER_PLIST}" ]]; then
    if /usr/libexec/PlistBuddy -c "Print" "${NE_USER_PLIST}" 2>/dev/null | grep -q "${BUNDLE}"; then
        echo "[SafeDNS] Found NE configuration, disabling..."
        COUNT=$(/usr/libexec/PlistBuddy -c "Print :DNSProxyConfigurations" "${NE_USER_PLIST}" 2>/dev/null | grep -c "Dict" || echo 0)
        for i in $(seq 0 $((COUNT - 1))); do
            IDENTIFIER=$(/usr/libexec/PlistBuddy -c "Print :DNSProxyConfigurations:${i}:Identifier" "${NE_USER_PLIST}" 2>/dev/null || echo "")
            if [[ "${IDENTIFIER}" == *"${BUNDLE}"* ]]; then
                /usr/libexec/PlistBuddy -c "Set :DNSProxyConfigurations:${i}:Enabled false" "${NE_USER_PLIST}" 2>/dev/null \
                    && echo "[SafeDNS] DNS Proxy disabled (index ${i})" \
                    || echo "[SafeDNS] Failed to set Enabled=false"
            fi
        done
        killall -HUP configd 2>/dev/null || true
        echo "[SafeDNS] configd reloaded"
    else
        echo "[SafeDNS] No NE configuration found — already clean"
    fi
else
    echo "[SafeDNS] NE plist not found — already clean"
fi
INNER_EOF
        chmod +x "${DEACTIVATE_SCRIPT}"
        launchctl asuser "${CONSOLE_UID}" sudo -u "${CONSOLE_USER}" bash "${DEACTIVATE_SCRIPT}" 2>/dev/null || true
        sleep 2
        log_ok "DNS Proxy deactivation completed (user: ${CONSOLE_USER})"
        rm -f "${DEACTIVATE_SCRIPT}" 2>/dev/null || true
    else
        log_warn "Failed to get console user's UID — skipping DNS Proxy deactivation"
    fi
fi

# Flush DNS cache immediately to restore connectivity
dscacheutil -flushcache 2>/dev/null || true
killall -HUP mDNSResponder 2>/dev/null || true
log_ok "DNS cache flushed — connectivity restored"

###
# 1. LaunchDaemon (system)
###
log_info "1. Stopping LaunchDaemon..."
if launchctl list 2>/dev/null | grep -q "${DAEMON_LABEL}"; then
    launchctl bootout system "${DAEMON_PLIST}" 2>/dev/null \
        || launchctl unload -w "${DAEMON_PLIST}" 2>/dev/null \
        || true
    log_ok "LaunchDaemon unloaded"
else
    log_warn "LaunchDaemon was not loaded"
fi

if [[ -f "${DAEMON_PLIST}" ]]; then
    rm -f "${DAEMON_PLIST}" 2>/dev/null || true
    log_ok "Removed ${DAEMON_PLIST}"
else
    log_warn "${DAEMON_PLIST} not found"
fi

###
# 1b. LaunchAgent (host auto-start)
###
log_info "1b. Stopping LaunchAgent (host)..."

# Enumerate all GUI domains and boot out the agent
launchctl print system 1>/dev/null 2>&1 || true

while IFS= read -r LINE; do
    DOMAIN=$(echo "$LINE" | awk '{print $1}')
    [[ "$DOMAIN" != gui/* ]] && continue

    if launchctl print "$DOMAIN/${HOST_AGENT_LABEL}" >/dev/null 2>&1; then
        launchctl bootout "$DOMAIN" "$HOST_AGENT_PLIST" 2>/dev/null \
            || launchctl bootout "$DOMAIN/${HOST_AGENT_LABEL}" 2>/dev/null \
            || true
        log_ok "LaunchAgent unloaded from ${DOMAIN}"
    fi
done < <(launchctl print system 2>/dev/null | grep '^gui/' || true)

# Additionally: boot out from the active console user's GUI session
if [[ -n "${CONSOLE_USER}" && "${CONSOLE_USER}" != "root" ]]; then
    CONSOLE_UID=$(id -u "${CONSOLE_USER}" 2>/dev/null || echo "")
    if [[ -n "${CONSOLE_UID}" ]]; then
        launchctl bootout "gui/${CONSOLE_UID}" "$HOST_AGENT_PLIST" 2>/dev/null \
            || launchctl bootout "gui/${CONSOLE_UID}/${HOST_AGENT_LABEL}" 2>/dev/null \
            || true
        log_ok "LaunchAgent unloaded from gui/${CONSOLE_UID}"
    fi
fi

if [[ -f "${HOST_AGENT_PLIST}" ]]; then
    rm -f "${HOST_AGENT_PLIST}" 2>/dev/null || true
    log_ok "Removed ${HOST_AGENT_PLIST}"
else
    log_warn "${HOST_AGENT_PLIST} not found"
fi

###
# 2. Kill system extension process and remove binary (block nesessionmanager restart)
###
log_info "2. Killing system extension process..."
PIDS=$(pgrep -f "${SYSEX_BUNDLE}" 2>/dev/null || true)
if [[ -n "${PIDS}" ]]; then
    echo "${PIDS}" | xargs kill -9 2>/dev/null || true
    log_ok "System extension processes killed: ${PIDS}"

    # IMPORTANT: remove the binary immediately before nesessionmanager can restart it
    find "${SYSEX_DIR}" -maxdepth 5 \
        -name "${SYSEX_BUNDLE}" -type f 2>/dev/null \
        | xargs rm -f 2>/dev/null || true
    log_ok "System extension binary removed — restart blocked"
else
    log_warn "System extension process not running"
fi

###
# 3. Physical removal of the system extension directory
###
log_info "3. Removing system extension from ${SYSEX_DIR}..."
FOUND=0
if [[ -d "${SYSEX_DIR}" ]]; then
    while IFS= read -r -d '' BUNDLE_PATH; do
        UUID_DIR=$(echo "${BUNDLE_PATH}" | awk -F'/' '{print "/"$2"/"$3"/"$4"/"$5"/"$6}')
        rm -rf "${UUID_DIR}" 2>/dev/null || true
        log_ok "Removed: ${UUID_DIR}"
        FOUND=1
    done < <(find "${SYSEX_DIR}" -maxdepth 4 \
        -name "${SYSEX_BUNDLE}.systemextension" -print0 2>/dev/null)
fi
[[ $FOUND -eq 0 ]] && log_warn "System extension bundle not found in ${SYSEX_DIR}"

SYSEX_DB="/Library/SystemExtensions/db.plist"
if [[ -f "${SYSEX_DB}" ]]; then
    /usr/libexec/PlistBuddy -c "Print" "${SYSEX_DB}" 2>/dev/null | grep -q "${SYSEX_BUNDLE}" \
        && log_warn "Entry in db.plist found — it will be cleared after reboot" \
        || log_ok "No entries for SafeDNS in db.plist"
fi

###
# 4. Host App removal
###
log_info "4. Removing host application..."
if [[ -d "${APP_PATH}" ]]; then
    rm -rf "${APP_PATH}" 2>/dev/null || true
    log_ok "Removed ${APP_PATH}"
else
    log_warn "${APP_PATH} not found"
fi

###
# 5. Application Support
###
log_info "5. Removing Application Support..."
for DIR in \
    "/Library/Application Support/SafeDNS" \
    "/Library/Application Support/SafeDNSMacProxy"; do
    if [[ -d "${DIR}" ]]; then
        rm -rf "${DIR}" 2>/dev/null || true
        log_ok "Removed ${DIR}"
    else
        log_warn "Not found: ${DIR}"
    fi
done

for USER_HOME in /Users/*/; do
    [[ -d "${USER_HOME}" ]] || continue
    for DIR in \
        "${USER_HOME}Library/Application Support/SafeDNSMacProxy" \
        "${USER_HOME}Library/Application Support/com.safedns.agent.SafeDNSMacProxy" \
        "${USER_HOME}Library/Group Containers/group.com.safedns.SafeDNSmac"; do
        if [[ -d "${DIR}" ]]; then
            rm -rf "${DIR}" 2>/dev/null || true
            log_ok "Removed ${DIR}"
        fi
    done
done

# Enterprise group container (MDM)
if [[ -d "/Library/Group Containers/group.com.safedns.agent.macos.enterprise" ]]; then
    rm -rf "/Library/Group Containers/group.com.safedns.agent.macos.enterprise" 2>/dev/null || true
    log_ok "Removed enterprise group container"
fi

###
# 6. Preferences / UserDefaults
###
log_info "6. Removing Preferences..."
for PLIST in \
    "/Library/Preferences/com.safedns.agent.SafeDNSMacProxy.plist" \
    "/Library/Preferences/com.safedns.agent.SafeDNSMacProxy.dnsProxy.plist"; do
    if [[ -f "${PLIST}" ]]; then
        rm -f "${PLIST}" 2>/dev/null || true
        log_ok "Removed ${PLIST}"
    else
        log_warn "Not found: ${PLIST}"
    fi
done

for USER_HOME in /Users/*/; do
    [[ -d "${USER_HOME}" ]] || continue
    for PLIST in \
        "${USER_HOME}Library/Preferences/com.safedns.agent.SafeDNSMacProxy.plist" \
        "${USER_HOME}Library/Preferences/com.safedns.agent.SafeDNSMacProxy.dnsProxy.plist"; do
        if [[ -f "${PLIST}" ]]; then
            rm -f "${PLIST}" 2>/dev/null || true
            log_ok "Removed ${PLIST}"
        fi
    done
done

NE_PREF="/Library/Preferences/com.apple.networkextension.plist"
if [[ -f "${NE_PREF}" ]]; then
    log_warn "NE preferences: ${NE_PREF} still exists — entry will be cleared after reboot"
fi

###
# 7. Logs
###
log_info "7. Removing logs..."
for PATH_TO_DEL in \
    "/var/log/com.safedns.agent.SafeDNSMacProxy.postinstall.log" \
    "/var/log/safedns.log" \
    "/var/log/safedns_pf.log" \
    "/var/log/safedns_daemon.log" \
    "/var/log/safedns-installer.log" \
    "/var/log/safedns-daemon.log" \
    "/var/log/safedns-daemon.err" \
    "/var/log/safedns-daemon-install.log" \
    "/var/log/safedns-proxy.log" \
    "/Library/Logs/SafeDNS" \
    "/tmp/deactivate_dns_proxy.swift" \
    "/tmp/check_dns_proxy.swift" \
    "/tmp/safedns_deactivate.sh"; do
    if [[ -e "${PATH_TO_DEL}" ]]; then
        rm -rf "${PATH_TO_DEL}" 2>/dev/null || true
        log_ok "Removed ${PATH_TO_DEL}"
    fi
done

log erase --subsystem "com.safedns.agent.SafeDNSMacProxy.dnsProxy" 2>/dev/null \
    && log_ok "Unified log entries erased" \
    || log_warn "No unified log entries found"

###
# 8. Keychain
###
log_info "8. Cleaning Keychain..."
security delete-generic-password -s "SafeDNS" \
    /Library/Keychains/System.keychain 2>/dev/null \
    && log_ok "System Keychain: item removed" \
    || log_warn "System Keychain: no items found"

for USER_HOME in /Users/*/; do
    USERNAME=$(basename "${USER_HOME}")
    [[ "${USERNAME}" == "Shared" ]] && continue
    sudo -u "${USERNAME}" security delete-generic-password \
        -s "SafeDNS" 2>/dev/null \
        && log_ok "User Keychain (${USERNAME}): item removed" \
        || true
done

###
# 9. DNS final reset
###
log_info "9. Final DNS reset..."
dscacheutil -flushcache 2>/dev/null && log_ok "dscacheutil: cache flushed"
killall -HUP mDNSResponder 2>/dev/null && log_ok "mDNSResponder: restarted"

while IFS= read -r SERVICE; do
    [[ -z "${SERVICE}" ]] && continue
    [[ "${SERVICE}" == *"*"* ]] && continue
    networksetup -setdnsservers "${SERVICE}" "Empty" 2>/dev/null \
        && log_ok "DNS reset: ${SERVICE}" \
        || true
done < <(networksetup -listallnetworkservices 2>/dev/null | tail -n +2)

###
# 10. Verification
###
echo ""
echo "════════════════════ RESULT ════════════════════"

echo -n "  LaunchDaemon plist:     "
[[ -f "${DAEMON_PLIST}" ]] && log_err "REMAINS!" || log_ok "removed"

echo -n "  LaunchAgent plist:      "
[[ -f "${HOST_AGENT_PLIST}" ]] && log_err "REMAINS!" || log_ok "removed"

echo -n "  Host application:       "
[[ -d "${APP_PATH}" ]] && log_err "REMAINS!" || log_ok "removed"

echo -n "  System extension bundle:"
REMAINING=$(find "${SYSEX_DIR}" -maxdepth 4 \
    -name "${SYSEX_BUNDLE}.systemextension" 2>/dev/null | head -1)
[[ -n "${REMAINING}" ]] && log_err "REMAINS: ${REMAINING}" || log_ok "removed"

echo -n "  System extension proc:  "
PROC=$(pgrep -f "${SYSEX_BUNDLE}" 2>/dev/null || true)
[[ -n "${PROC}" ]] && log_err "RUNNING (PID ${PROC})" || log_ok "not running"

echo -n "  DNS cache:              "
log_ok "flushed"

echo ""
log_ok "══════════════════════════════════════════════════"
log_ok "Cleanup completed. Please reboot the machine."