#!/usr/bin/env bash
#
# SoftSol Martyria - install and enroll the Wazuh endpoint agent on Ubuntu/Debian.
#
# Self-contained installer for internet-connected Linux workstations or servers.
# Adds the Wazuh apt repository, installs the agent, enrolls with Baruch, and
# applies the workstations group config.
#
# Canonical file (update here when changing the public installer):
#   midas/static/downloads/SoftSol-Martyria.sh
#
# Public download URL:
#   https://softsol.co.za/downloads/SoftSol-Martyria.sh
#
# Run as root, or approve the sudo prompt when the script re-launches itself.

set -euo pipefail

# =============================================================================
# Configuration - edit these values when updating the public installer
# =============================================================================
MANAGER_PUBLIC='martyria.softsol.co.za'
MANAGER_INTERNAL='172.17.12.34'
REGISTRATION_PORT=1515
EVENT_PORT=1514
GROUP='workstations'
VERSION='4.14.5'
REGISTRATION_PASSWORD='MartyriaWitness2026!'
AGENT_NAME_SUFFIX='martyria'
STAGE_DIR='/var/lib/martyria/wazuh-install'
PRODUCT_NAME='SoftSol Martyria'
OSSEC_CONF='/var/ossec/etc/ossec.conf'
CLIENT_KEYS='/var/ossec/etc/client.keys'
AGENT_AUTH='/var/ossec/bin/agent-auth'
OSSEC_LOG='/var/ossec/logs/ossec.log'

log() {
  printf '%s\n' "$*"
}

die() {
  log "ERROR: $*" >&2
  exit 1
}

require_root() {
  if [[ "${EUID:-$(id -u)}" -ne 0 ]]; then
    log '==> Root privileges required - re-running with sudo...'
    exec sudo -E bash "$0" "$@"
  fi
}

require_apt() {
  command -v apt-get >/dev/null 2>&1 || die 'This installer supports Ubuntu/Debian (apt) only.'
}

test_tcp_port_open() {
  local host="$1" port="$2" timeout_sec="${3:-2}"
  timeout "$timeout_sec" bash -c "echo >/dev/tcp/${host}/${port}" 2>/dev/null
}

resolve_martyria_endpoint() {
  if test_tcp_port_open "$MANAGER_INTERNAL" "$REGISTRATION_PORT"; then
    log "==> SoftSol LAN/VPN detected - using internal manager ${MANAGER_INTERNAL}"
    MANAGER="$MANAGER_INTERNAL"
    AGENT_GROUP='workstations-internal'
  else
    log "==> Using public manager ${MANAGER_PUBLIC}"
    MANAGER="$MANAGER_PUBLIC"
    AGENT_GROUP="$GROUP"
  fi
}

default_agent_name() {
  local name
  name="$(hostname -s 2>/dev/null || hostname)-${AGENT_NAME_SUFFIX}"
  name="${name//[^a-zA-Z0-9._-]/-}"
  name="${name#-}"
  name="${name%-}"
  [[ ${#name} -gt 128 ]] && name="${name:0:128}"
  printf '%s' "$name"
}

set_agent_remote_config() {
  local address="$1" port="${2:-1514}"
  [[ -f "$OSSEC_CONF" ]] || return 0
  python3 - "$address" "$port" "$OSSEC_CONF" <<'PY'
import re
import sys

address, port, path = sys.argv[1:4]

with open(path, 'r', encoding='utf-8', errors='replace') as f:
    content = f.read()

match = re.search(r'(<client>)(.*?)(</client>)', content, re.DOTALL)
if not match:
    sys.exit(0)

block = match.group(2)

def set_tag(block, tag, value):
    pattern = rf'(<{tag}>)[^<]*(</{tag}>)'
    if re.search(pattern, block):
        return re.sub(pattern, rf'\g<1>{value}\g<2>', block, count=1)
    return block

if re.search(r'<server>', block):
    if re.search(r'<address>', block):
        block = set_tag(block, 'address', address)
    else:
        block = re.sub(r'(<server>)', rf'\g<1>\n      <address>{address}</address>', block, count=1)
    if re.search(r'<port>', block):
        block = set_tag(block, 'port', str(port))
    else:
        block = re.sub(r'(</server>)', f'      <port>{port}</port>\n    \\g<1>', block, count=1)
    if re.search(r'<protocol>', block):
        block = set_tag(block, 'protocol', 'tcp')
    else:
        block = re.sub(r'(</server>)', '      <protocol>tcp</protocol>\n    \\g<1>', block, count=1)
else:
    block = (
        block.rstrip()
        + f'\n    <server>\n      <address>{address}</address>\n'
        + f'      <port>{port}</port>\n      <protocol>tcp</protocol>\n    </server>\n'
    )

if re.search(r'<crypto_method>', block):
    block = set_tag(block, 'crypto_method', 'aes')
else:
    block = block.rstrip() + '\n    <crypto_method>aes</crypto_method>\n'

content = content[:match.start(2)] + block + content[match.end(2):]

with open(path, 'w', encoding='utf-8') as f:
    f.write(content)
PY
  log "==> Agent configured for ${address}:${port} (crypto_method=aes)"
}

agent_enrolled() {
  [[ -f "$CLIENT_KEYS" ]] && grep -q '[[:graph:]]' "$CLIENT_KEYS"
}

agent_connected() {
  [[ -f "$OSSEC_LOG" ]] || return 1
  tail -n 40 "$OSSEC_LOG" | grep -qE 'Connected to the server|Connected to .*1514' \
    && ! tail -n 40 "$OSSEC_LOG" | grep -qE 'Invalid password|Unable to connect to any server'
}

ensure_wazuh_repo() {
  apt-get install -y gnupg apt-transport-https curl ca-certificates

  if [[ ! -f /usr/share/keyrings/wazuh.gpg ]]; then
    curl -fsSL https://packages.wazuh.com/key/GPG-KEY-WAZUH \
      | gpg --no-default-keyring --keyring gnupg-ring:/usr/share/keyrings/wazuh.gpg --import
    chmod 644 /usr/share/keyrings/wazuh.gpg
  fi

  if [[ ! -f /etc/apt/sources.list.d/wazuh.list ]]; then
    echo "deb [signed-by=/usr/share/keyrings/wazuh.gpg] https://packages.wazuh.com/4.x/apt/ stable main" \
      > /etc/apt/sources.list.d/wazuh.list
  fi

  apt-get update
}

agent_installed() {
  dpkg-query -W -f='${Status}' wazuh-agent 2>/dev/null | grep -q 'install ok installed'
}

invoke_agent_enrollment() {
  local manager="$1" agent_name="$2" agent_group="$3" reg_pass="$4"

  [[ -x "$AGENT_AUTH" ]] || die 'agent-auth not found'

  if [[ -f "$CLIENT_KEYS" ]]; then
    rm -f "$CLIENT_KEYS"
  fi

  log "==> Enrolling agent: ${agent_name} -> ${manager}:${REGISTRATION_PORT}"
  if ! "$AGENT_AUTH" \
      -m "$manager" \
      -p "$REGISTRATION_PORT" \
      -A "$agent_name" \
      -G "$agent_group" \
      -P "$reg_pass"; then
    die "Enrollment failed. From this host, TCP to ${manager}:${REGISTRATION_PORT} must succeed. If the log shows 'Duplicate IP', the manager already has this host registered - contact SoftSol IT to remove the stale entry, then re-run this script."
  fi

  agent_enrolled || die 'agent-auth exited 0 but client.keys was not created - contact SoftSol IT.'
}

install_softsol_martyria_agent() {
  local agent_name manager reg_pass needs_enrollment=1 log_path

  require_apt
  mkdir -p "$STAGE_DIR"
  log_path="${STAGE_DIR}/install.log"
  exec > >(tee -a "$log_path") 2>&1

  log ''
  log "  ${PRODUCT_NAME} - Wazuh endpoint agent installer (Ubuntu/Debian)"
  resolve_martyria_endpoint
  manager="$MANAGER"
  agent_name="$(default_agent_name)"
  reg_pass="${REGISTRATION_PASSWORD#"${REGISTRATION_PASSWORD%%[![:space:]]*}"}"
  reg_pass="${reg_pass%"${reg_pass##*[![:space:]]}"}"
  [[ -n "$reg_pass" ]] || die 'Registration password is not configured in this script.'

  log "  Manager: ${manager}"
  log ''

  if agent_installed; then
    log '==> Wazuh agent already installed - skipping package install (avoids wiping enrollment keys)'
    systemctl stop wazuh-agent 2>/dev/null || true

    if agent_enrolled; then
      log '==> Existing enrollment keys found - keeping keys (will update config and restart)'
      needs_enrollment=0
    fi
  else
    ensure_wazuh_repo
    log '==> Installing Wazuh agent package'
    DEBIAN_FRONTEND=noninteractive \
      WAZUH_MANAGER="$manager" \
      WAZUH_AGENT_NAME="$agent_name" \
      WAZUH_AGENT_GROUP="$AGENT_GROUP" \
      WAZUH_REGISTRATION_PASSWORD="$reg_pass" \
      apt-get install -y "wazuh-agent=${VERSION}-1"
    if agent_enrolled; then
      needs_enrollment=0
    fi
  fi

  if [[ "$needs_enrollment" -eq 1 ]]; then
    invoke_agent_enrollment "$manager" "$agent_name" "$AGENT_GROUP" "$reg_pass"
  fi

  set_agent_remote_config "$manager" "$EVENT_PORT"

  systemctl daemon-reload
  systemctl enable wazuh-agent
  sleep 2
  systemctl restart wazuh-agent
  sleep 15

  log ''
  log "==> ${PRODUCT_NAME} agent installed and started."
  log "    Name:    ${agent_name}"
  log "    Manager: ${manager}:${EVENT_PORT}"
  log "    Group:   ${AGENT_GROUP}"
  log ''
  if agent_connected; then
    log '==> Agent connected to manager.'
  else
    log '==> Agent service started. If it is not Active on the dashboard within 2 minutes, check:'
    log "    sudo tail -50 ${OSSEC_LOG}"
  fi
  log "Install log: ${log_path}"
}

require_root "$@"
install_softsol_martyria_agent