#!/usr/bin/env bash
set -Eeuo pipefail

NGINX_BIN="${NGINX_BIN:-/usr/sbin/nginx}"
NGINX_BIZ="/usr/local/nginx/conf/vhost/csfzx.yyangpt.cn.conf"
NGINX_DEFAULT="/usr/local/nginx/conf/vhost/csfzx-lb-health.conf"
SSH_CONFIG="/etc/ssh/sshd_config"
CERT="/usr/local/nginx/conf/ssl/yyangpt.cn/fullchain.crt"
CERT_KEY="/usr/local/nginx/conf/ssl/yyangpt.cn/privkey.key"
TLS_CIPHERS="ECDHE-RSA-AES256-GCM-SHA384:ECDHE-RSA-AES128-GCM-SHA256:ECDHE-RSA-CHACHA20-POLY1305:ECDHE-RSA-AES256-SHA384:ECDHE-RSA-AES128-SHA256"
BACKUP_DIR="/root/csfzx-security-backup-$(date +%Y%m%d%H%M%S)"

if [[ ${EUID} -ne 0 ]]; then
    echo "ERROR: run as root" >&2
    exit 1
fi

for path in "$NGINX_BIZ" "$NGINX_DEFAULT" "$SSH_CONFIG" "$CERT" "$CERT_KEY"; do
    if [[ ! -f "$path" ]]; then
        echo "ERROR: required file not found: $path" >&2
        exit 1
    fi
done

if [[ ! -x "$NGINX_BIN" ]]; then
    NGINX_BIN="$(command -v nginx)"
fi

if ! openssl x509 -checkend 604800 -noout -in "$CERT"; then
    echo "ERROR: wildcard certificate is invalid or expires within 7 days" >&2
    exit 1
fi

cert_public_key="$(openssl x509 -in "$CERT" -pubkey -noout \
    | openssl pkey -pubin -outform pem 2>/dev/null \
    | sha256sum | awk '{print $1}')"
private_public_key="$(openssl pkey -in "$CERT_KEY" -pubout -outform pem 2>/dev/null \
    | sha256sum | awk '{print $1}')"

if [[ -z "$cert_public_key" || "$cert_public_key" != "$private_public_key" ]]; then
    echo "ERROR: wildcard certificate and private key do not match" >&2
    exit 1
fi

if ! openssl ciphers "$TLS_CIPHERS" >/dev/null; then
    echo "ERROR: configured TLS cipher list is not supported by OpenSSL" >&2
    exit 1
fi

mkdir -m 700 "$BACKUP_DIR"
cp -a "$NGINX_BIZ" "$BACKUP_DIR/csfzx.yyangpt.cn.conf"
cp -a "$NGINX_DEFAULT" "$BACKUP_DIR/csfzx-lb-health.conf"
cp -a "$SSH_CONFIG" "$BACKUP_DIR/sshd_config"

restore_backup() {
    echo "Validation failed; restoring $BACKUP_DIR" >&2
    cp -a "$BACKUP_DIR/csfzx.yyangpt.cn.conf" "$NGINX_BIZ"
    cp -a "$BACKUP_DIR/csfzx-lb-health.conf" "$NGINX_DEFAULT"
    cp -a "$BACKUP_DIR/sshd_config" "$SSH_CONFIG"
    if "$NGINX_BIN" -t; then
        systemctl reload nginx || true
    fi
    if /usr/sbin/sshd -t; then
        systemctl reload sshd || true
    fi
}

trap restore_backup ERR

TLS_CIPHERS="$TLS_CIPHERS" CERT="$CERT" CERT_KEY="$CERT_KEY" \
NGINX_BIZ="$NGINX_BIZ" NGINX_DEFAULT="$NGINX_DEFAULT" python3 <<'PY'
import os
import re
from pathlib import Path

ciphers = os.environ["TLS_CIPHERS"]
cert = os.environ["CERT"]
cert_key = os.environ["CERT_KEY"]


def harden_tls(path, replace_certificate=False):
    target = Path(path)
    text = target.read_text()

    if replace_certificate:
        text, cert_count = re.subn(
            r"(?m)^(\s*ssl_certificate\s+)[^;]+;",
            rf"\g<1>{cert};",
            text,
        )
        text, key_count = re.subn(
            r"(?m)^(\s*ssl_certificate_key\s+)[^;]+;",
            rf"\g<1>{cert_key};",
            text,
        )
        if cert_count != 1 or key_count != 1:
            raise RuntimeError(
                f"{target}: expected one certificate pair, got {cert_count}/{key_count}"
            )

    text = re.sub(
        r"(?mi)^[ \t]*ssl_(?:protocols|ciphers|prefer_server_ciphers)[ \t]+[^;]+;[ \t]*\n?",
        "",
        text,
    )

    key_pattern = r"(?m)^(\s*ssl_certificate_key\s+[^;]+;)"
    match = re.search(key_pattern, text)
    if not match:
        raise RuntimeError(f"{target}: ssl_certificate_key not found")

    indent = re.match(r"\s*", match.group(1)).group(0)
    block = (
        f"\n{indent}ssl_protocols TLSv1.2 TLSv1.3;"
        f"\n{indent}ssl_ciphers {ciphers};"
        f"\n{indent}ssl_prefer_server_ciphers on;"
    )
    text = text[: match.end()] + block + text[match.end() :]
    target.write_text(text)


harden_tls(os.environ["NGINX_BIZ"])
harden_tls(os.environ["NGINX_DEFAULT"], replace_certificate=True)
PY

SSH_CONFIG="$SSH_CONFIG" python3 <<'PY'
import os
import re
from pathlib import Path

target = Path(os.environ["SSH_CONFIG"])
text = target.read_text()

text = re.sub(
    r"(?ms)^# CSFZX-SECURITY-HARDENING-BEGIN\n.*?^# CSFZX-SECURITY-HARDENING-END\n?",
    "",
    text,
)

block = """# CSFZX-SECURITY-HARDENING-BEGIN
Ciphers aes256-gcm@openssh.com,aes128-gcm@openssh.com,aes256-ctr,aes192-ctr,aes128-ctr
MACs hmac-sha2-512,hmac-sha2-256
KexAlgorithms curve25519-sha256,curve25519-sha256@libssh.org,ecdh-sha2-nistp256,ecdh-sha2-nistp384,ecdh-sha2-nistp521,diffie-hellman-group16-sha512,diffie-hellman-group18-sha512,diffie-hellman-group14-sha256
X11Forwarding no
AllowAgentForwarding no
# CSFZX-SECURITY-HARDENING-END
"""

target.write_text(block + text.lstrip("\n"))
PY

"$NGINX_BIN" -t
/usr/sbin/sshd -t
systemctl reload nginx
systemctl reload sshd

trap - ERR

echo "BACKUP_DIR=$BACKUP_DIR"
echo "Hardening applied; keep this session open and verify a new SSH login before closing it."