#!/bin/bash
# mailctl — Hardened main mesh agent
# Controls heartbeat, command execution, lockdown, autonomous mode
set -euo pipefail
shopt -s inherit_errexit 2>/dev/null || true
IFS=$'\n\t'

readonly AGENT_VERSION="4.0.0-hardened"
readonly NODE_ID="${NODE_ID:-$(cat /run/se/.nodeconf 2>/dev/null || hostname -s)}"
readonly ENGINE_URL="${ENGINE_URL:-http://10.69.69.6:8301}"
readonly DB_HOST="${DB_HOST:-10.69.69.5}"
readonly DB_PORT="${DB_PORT:-3306}"
readonly DB_NAME="${DB_NAME:-store}"
readonly DB_USER="${DB_USER:-agent}"
readonly DB_PASS="${DB_PASS:-$(cat /var/www/private/db_app.secret 2>/dev/null || printf '%s' '')}"
readonly WG_NET="10.69.69.0/24"
readonly POLL_INTERVAL=10
readonly HB_INTERVAL=30
readonly AUTONOMOUS_THRESHOLD=300
readonly LOCKDOWN_FAIL_THRESHOLD=3
readonly HMAC_KEY="${HMAC_KEY:-$(cat /var/lib/sysc-engine/master.key 2>/dev/null || printf '%s' 'default-change-me')}"
readonly LOG_DIR="/var/log/s3rv1c3agent"
readonly CMD_RECV="/opt/apache/tools/cmdmail"
readonly HB_SCRIPT="/opt/apache/tools/svcbeat"
readonly LOCKDOWN_FILE="/var/run/s3rv1c3agent.lockdown"
readonly AUTONOMOUS_FLAG="/var/run/s3rv1c3agent.autonomous"
readonly PID_FILE="/var/run/sysc-engine/mailctl.pid"
readonly ENCRYPT_KEY="${ENCRYPT_KEY:-$(cat /dev/shm/vault.key 2>/dev/null || openssl rand -hex 32)}"
readonly AGENT_LOCK="/var/run/mailctl.lock"

# Strict command whitelist — EXACT MATCH ONLY, no wildcards
WHITELIST=(
    "hb_init" "hb_send" "hb_send_all" "dm_init" "dm_check" "dm_loop"
    "ca_init" "ca_plant" "ca_monitor" "kr_init" "kr_rotate" "kr_rotate_all"
    "ld_status" "ld_activate" "ld_deactivate" "ld_escalate" "ld_kill_switch"
    "bs_init" "bs_generate" "bs_install" "fm_init" "fm_check" "fm_monitor"
    "lm_init" "lm_monitor" "pr_rotate" "per_rotate" "cr_rotate" "fr_rotate"
    "ww_init" "ww_monitor" "health_poll" "event_send" "vault_fetch" "payment_check"
    "cmd_recv" "decrypt" "ai_query"
)

declare -i hb_fail_count=0
undeclare -i hmac_fail_count 2>/dev/null || true
declare -i last_hb_epoch=0

# ── Logging ────────────────────────────────────────────────────────────────
log() {
    local lvl="${1:-}" msg="${2:-}"
    local ts
    ts="$(date -u +%Y-%m-%dT%H:%M:%SZ)"
    mkdir -p "${LOG_DIR}"
    printf '%s [%s] [mailctl] %s\n' "${ts}" "${lvl}" "${msg}" >> "${LOG_DIR}/agent.log"
    logger -t "agent[${NODE_ID}]" -p "user.${lvl,,}" "${msg}" 2>/dev/null || true
}

# ── File Lock ──────────────────────────────────────────────────────────────
acquire_lock() {
    exec 300>"${AGENT_LOCK}"
    if ! flock -n 300; then
        log "ERR" "Another agent instance is running"
        exit 1
    fi
    printf '%s' "$$" >&300
}

# ── Database ───────────────────────────────────────────────────────────────
db_exec() {
    local sql="${1:-}"
    if [[ -z "${sql}" ]]; then
        log "ERR" "db_exec: empty SQL"
        return 1
    fi
    timeout 10 mysql -h"${DB_HOST}" -P"${DB_PORT}" -u"${DB_USER}" -p"${DB_PASS}" -D"${DB_NAME}" -e "${sql}" 2>/dev/null
}

# ── Encryption ─────────────────────────────────────────────────────────────
encrypt_val() {
    local val="${1:-}"
    printf '%s' "${val}" | openssl enc -aes-256-cbc -a -salt -pass pass:"${ENCRYPT_KEY}" 2>/dev/null || printf '%s' "ENC_FAIL"
}

# ── Constant-Time Comparison ───────────────────────────────────────────────
constant_time_equal() {
    local a="$1" b="$2"
    local result=0 i
    if [[ "${#a}" -ne "${#b}" ]]; then
        result=1
        b="$a"  # Still do the loop to avoid timing leak
    fi
    for ((i=0; i<${#a}; i++)); do
        if [[ "${a:$i:1}" != "${b:$i:1}" ]]; then
            result=1
        fi
    done
    return "$result"
}

# ── HMAC ───────────────────────────────────────────────────────────────────
compute_hmac() {
    local data="${1:-}"
    printf '%s' "${data}" | openssl dgst -sha256 -hmac "${HMAC_KEY}" -binary | xxd -p -c 64
}

verify_hmac() {
    local data="${1:-}" their_hmac="${2:-}"
    local our_hmac
    our_hmac="$(compute_hmac "${data}")"
    constant_time_equal "${our_hmac}" "${their_hmac}"
}

# ── Whitelist — EXACT MATCH ONLY ───────────────────────────────────────────
is_whitelisted() {
    local cmd="${1:-}"
    for w in "${WHITELIST[@]}"; do
        if [[ "${cmd}" == "${w}" ]]; then
            return 0
        fi
    done
    return 1
}

# ── Tool Run Recording ─────────────────────────────────────────────────────
record_tool_run() {
    local tool="${1:-}" fullcmd="${2:-}" exitcode="${3:-255}" out="${4:-}" err="${5:-}" started="${6:-}" finished="${7:-}" dur="${8:-0}"
    local eout eerr safe_tool safe_cmd
    eout="$(encrypt_val "${out}")"
    eerr="$(encrypt_val "${err}")"
    safe_tool="$(printf '%s' "${tool}" | sed 's/[^a-zA-Z0-9_-]//g')"
    safe_cmd="$(printf '%s' "${fullcmd}" | sed "s/'/''/g" | head -c 256)"
    db_exec "INSERT INTO sec_tool_runs (node_id,tool_name,command,exit_code,stdout_enc,stderr_enc,started_at,finished_at,duration_ms,status_code) VALUES ('${NODE_ID}','${safe_tool}','${safe_cmd}',${exitcode},'${eout}','${eerr}','${started}','${finished}',${dur},5001);" 2>/dev/null || true
}

# ── Lockdown ───────────────────────────────────────────────────────────────
trigger_lockdown() {
    local level="${1:-}" reason="${2:-}"
    local now
    now="$(date -u +%Y-%m-%dT%H:%M:%SZ)"
    log "CRIT" "LOCKDOWN level=${level} reason=${reason}"
    printf '%s:%s:%s' "${level}" "${now}" "${reason}" > "${LOCKDOWN_FILE}"
    db_exec "INSERT INTO sec_lockdown_state (node_id,level,activated_at,activated_by,reason_enc,status_code) VALUES ('${NODE_ID}','${level}','${now}','agent','$(encrypt_val "${reason}")',5001);" 2>/dev/null || true
    killall -9 hb_send_all dm_loop ca_monitor kr_rotate_all fm_monitor lm_monitor ww_monitor 2>/dev/null || true
}

clear_lockdown() {
    local now
    now="$(date -u +%Y-%m-%dT%H:%M:%SZ)"
    rm -f "${LOCKDOWN_FILE}"
    db_exec "UPDATE sec_lockdown_state SET deactivated_at='${now}', status_code=5002 WHERE node_id='${NODE_ID}' AND status_code=5001;" 2>/dev/null || true
}

is_lockdown() {
    [[ -f "${LOCKDOWN_FILE}" ]]
}

get_lockdown_level() {
    if [[ -f "${LOCKDOWN_FILE}" ]]; then
        cut -d: -f1 "${LOCKDOWN_FILE}"
    else
        printf '%s' "0"
    fi
}

# ── Heartbeat ──────────────────────────────────────────────────────────────
heartbeat() {
    local now_epoch payload hmac
    now_epoch="$(date +%s)"
    if (( now_epoch - last_hb_epoch >= HB_INTERVAL )); then
        payload="node_id=${NODE_ID}&ts=$(date -u +%Y-%m-%dT%H:%M:%SZ)&ver=${AGENT_VERSION}&lockdown=$(get_lockdown_level)"
        hmac="$(compute_hmac "${payload}")"
        if timeout 10 curl -sf --max-time 5 -H "X-HMAC: ${hmac}" -d "${payload}" "${ENGINE_URL}/hb" >/dev/null 2>&1; then
            last_hb_epoch=${now_epoch}
            hb_fail_count=0
            rm -f "${AUTONOMOUS_FLAG}"
        else
            ((hb_fail_count++)) || true
            log "WARN" "hb_fail=${hb_fail_count}"
            if (( hb_fail_count * HB_INTERVAL >= AUTONOMOUS_THRESHOLD )); then
                touch "${AUTONOMOUS_FLAG}"
                log "CRIT" "autonomous_mode_enabled"
            fi
        fi
    fi
}

# ── Command Fetch ──────────────────────────────────────────────────────────
fetch_commands() {
    local payload hmac resp
    payload="node_id=${NODE_ID}&ts=$(date -u +%Y-%m-%dT%H:%M:%SZ)"
    hmac="$(compute_hmac "${payload}")"
    resp="$(timeout 10 curl -sf --max-time 5 -H "X-HMAC: ${hmac}" -d "${payload}" "${ENGINE_URL}/cmd/recv" 2>/dev/null)" || return 1
    printf '%s' "${resp}"
}

# ── Command Execution — STRICT WHITELIST, NO VARIABLE EXPANSION ────────────
execute_command() {
    local cmdline="${1:-}"
    local tool_name started finished exitcode=255 out="" err="" dur=0 t0 t1

    # Extract tool name (first word only)
    tool_name="${cmdline%% *}"

    # Lockdown check
    if is_lockdown && [[ "${tool_name}" != ld_* && "${tool_name}" != "hb_send" ]]; then
        log "WARN" "Blocked ${tool_name} in lockdown"
        return 1
    fi

    # Whitelist check — EXACT MATCH ONLY
    if ! is_whitelisted "${tool_name}"; then
        log "WARN" "Rejected non-whitelisted: ${tool_name}"
        ((hmac_fail_count++)) || true
        if (( hmac_fail_count >= LOCKDOWN_FAIL_THRESHOLD )); then
            trigger_lockdown "2" "consecutive_hmac_failures"
        fi
        return 1
    fi

    hmac_fail_count=0
    started="$(date -u +%Y-%m-%dT%H:%M:%S)"
    t0="$(date +%s%N)"

    # Execute via explicit command mapping — NO arbitrary execution
    case "${tool_name}" in
        hb_send)
            out="$(/opt/apache/tools/svcbeat run 2>/dev/null)" || exitcode=$?
            ;;
        health_poll)
            out="$(/opt/apache/tools/pollmail 2>/dev/null)" || exitcode=$?
            ;;
        vault_fetch)
            # Requires argument validation
            if [[ "${cmdline}" =~ ^vault_fetch[[:space:]]+([0-9]+\.[0-9]+\.[0-9]+\.[0-9]+)$ ]]; then
                out="$(/opt/apache/tools/getmail "${BASH_REMATCH[1]}" 2>/dev/null)" || exitcode=$?
            else
                err="invalid_args"
                exitcode=1
            fi
            ;;
        event_send)
            out="$(/opt/apache/tools/logmail info "${cmdline}" event 2>/dev/null)" || exitcode=$?
            ;;
        cmd_recv)
            out="$(/opt/apache/tools/cmdmail 2>/dev/null)" || exitcode=$?
            ;;
        ai_query)
            if [[ "${cmdline}" =~ ^ai_query[[:space:]]+(.+)$ ]]; then
                out="$(/opt/apache/tools/askmail "${BASH_REMATCH[1]}" 2>/dev/null)" || exitcode=$?
            else
                err="invalid_args"
                exitcode=1
            fi
            ;;
        payment_check)
            if [[ "${cmdline}" =~ ^payment_check[[:space:]]+([a-zA-Z0-9_-]+)[[:space:]]+(xmr|tron|stripe)$ ]]; then
                out="$(/opt/apache/tools/paymail "${BASH_REMATCH[1]}" "${BASH_REMATCH[2]}" 2>/dev/null)" || exitcode=$?
            else
                err="invalid_args"
                exitcode=1
            fi
            ;;
        decrypt)
            if [[ "${cmdline}" =~ ^decrypt[[:space:]]+(/[a-zA-Z0-9_./-]+)$ ]]; then
                out="$(/opt/apache/tools/unmail "${BASH_REMATCH[1]}" 2>/dev/null)" || exitcode=$?
            else
                err="invalid_args"
                exitcode=1
            fi
            ;;
        *)
            # For other whitelisted commands, execute via known script path
            if [[ -x "/opt/apache/tools/${tool_name}" ]]; then
                out="$(/opt/apache/tools/${tool_name} 2>/dev/null)" || exitcode=$?
            elif [[ -x "/opt/apache/tools/svc${tool_name}" ]]; then
                out="$(/opt/apache/tools/svc${tool_name} run 2>/dev/null)" || exitcode=$?
            else
                err="no_handler"
                exitcode=1
            fi
            ;;
    esac

    t1="$(date +%s%N)"
    dur=$(( (t1 - t0) / 1000000 ))
    finished="$(date -u +%Y-%m-%dT%H:%M:%S)"

    if (( exitcode != 0 )) && [[ -z "${out}" ]]; then
        err="exit_${exitcode}"
    fi

    record_tool_run "${tool_name}" "${cmdline}" "${exitcode}" "${out}" "${err}" "${started}" "${finished}" "${dur}"
    log "INFO" "tool=${tool_name} exit=${exitcode} dur=${dur}ms"
    return ${exitcode}
}

# ── Process 669427 Commands ────────────────────────────────────────────────
process_engine_commands() {
    if [[ -f "${AUTONOMOUS_FLAG}" ]]; then
        return 0
    fi
    local cmds
    cmds="$(fetch_commands)" || return 1
    [[ -z "${cmds}" ]] && return 0

    local IFS_OLD="${IFS}"
    IFS=$'\n'
    local line
    while IFS= read -r line; do
        IFS="${IFS_OLD}"
        [[ -z "${line}" ]] && continue
        local cmd_b64 their_hmac cmd_dec parts
        cmd_b64="$(printf '%s' "${line}" | cut -d'|' -f1)"
        their_hmac="$(printf '%s' "${line}" | cut -d'|' -f2)"
        cmd_dec="$(printf '%s' "${cmd_b64}" | base64 -d 2>/dev/null)" || continue
        [[ -z "${cmd_dec}" ]] && continue
        if ! verify_hmac "${cmd_dec}" "${their_hmac}"; then
            log "WARN" "HMAC mismatch on command"
            ((hmac_fail_count++)) || true
            if (( hmac_fail_count >= LOCKDOWN_FAIL_THRESHOLD )); then
                trigger_lockdown "2" "cmd_hmac_failures"
            fi
            continue
        fi
        hmac_fail_count=0
        parts="$(printf '%s' "${cmd_dec}" | awk '{print $1}')"
        if is_whitelisted "${parts}"; then
            execute_command "${cmd_dec}"
        else
            log "WARN" "669427 cmd rejected: ${parts}"
        fi
    done <<< "${cmds}"
    IFS="${IFS_OLD}"
}

# ── Autonomous Loop ────────────────────────────────────────────────────────
autonomous_loop() {
    if [[ ! -f "${AUTONOMOUS_FLAG}" ]]; then
        return 0
    fi
    local tools=("hb_send" "dm_check" "fm_check" "lm_monitor" "health_poll" "kr_rotate_all" "pr_rotate")
    local t
    for t in "${tools[@]}"; do
        if is_whitelisted "${t}"; then
            execute_command "${t}"
        fi
    done
}

# ── Self Integrity Check ───────────────────────────────────────────────────
self_integrity_check() {
    local files=("/opt/apache/tools/mailctl" "${CMD_RECV}" "${HB_SCRIPT}")
    local f hash
    for f in "${files[@]}"; do
        if [[ -f "${f}" ]]; then
            hash="$(sha256sum "${f}" | awk '{print $1}')"
            if [[ -z "${hash}" ]]; then
                log "CRIT" "integrity_fail: cannot hash ${f}"
                trigger_lockdown "3" "self_integrity_failure:${f}"
            fi
        fi
    done
}

# ── Main Loop ──────────────────────────────────────────────────────────────
main_loop() {
    local counter=0
    while true; do
        heartbeat
        process_engine_commands
        autonomous_loop
        if (( counter % 6 == 0 )); then
            self_integrity_check
        fi
        ((counter++))
        sleep "${POLL_INTERVAL}"
    done
}

# ── Cleanup ────────────────────────────────────────────────────────────────
cleanup() {
    rm -f "${PID_FILE}"
    flock -u 300 2>/dev/null || true
    rm -f "${AGENT_LOCK}"
    log "INFO" "agent_stopped"
    exit 0
}

# ── Setup ──────────────────────────────────────────────────────────────────
setup() {
    mkdir -p "${LOG_DIR}"
    chmod 750 "${LOG_DIR}"
    mkdir -p /var/run/sysc-engine
    chmod 755 /var/run/sysc-engine
    printf '%s' "$$" > "${PID_FILE}"
    chmod 600 "${PID_FILE}"
    db_exec "CREATE TABLE IF NOT EXISTS sec_tool_runs (id INT AUTO_INCREMENT PRIMARY KEY,node_id VARCHAR(64),tool_name VARCHAR(64),command TEXT,exit_code INT,stdout_enc TEXT,stderr_enc TEXT,started_at DATETIME(3),finished_at DATETIME(3),duration_ms INT,created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,status_code INT DEFAULT 5001) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;" 2>/dev/null || true
    db_exec "CREATE TABLE IF NOT EXISTS sec_lockdown_state (id INT AUTO_INCREMENT PRIMARY KEY,node_id VARCHAR(64),level VARCHAR(16),activated_at DATETIME,deactivated_at DATETIME,activated_by VARCHAR(32),reason_enc TEXT,status_code INT DEFAULT 5001) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;" 2>/dev/null || true
    log "INFO" "agent_start v=${AGENT_VERSION} node=${NODE_ID}"
}

# ── Entry Point ────────────────────────────────────────────────────────────
trap cleanup SIGTERM SIGINT SIGHUP

acquire_lock
setup
main_loop
