#!/bin/bash

# Đảm bảo script chạy với quyền root
if [ "$EUID" -ne 0 ]; then
    echo "Vui lòng chạy script với quyền root!"
    exit 1
fi

# Thêm 103.206.218.201 broker01 vào /etc/hosts nếu chưa có
HOSTS_FILE="/etc/hosts"
BROKER_ENTRY="103.206.218.201 broker01"

if ! grep -q "$BROKER_ENTRY" "$HOSTS_FILE"; then
    echo "Đang thêm $BROKER_ENTRY vào $HOSTS_FILE"
    echo "$BROKER_ENTRY" >> "$HOSTS_FILE"
else
    echo "$BROKER_ENTRY đã tồn tại trong $HOSTS_FILE"
fi

# Chọn ngôn ngữ (English hoặc Vietnamese)
echo "Select language / Chọn ngôn ngữ:"
echo "1. English"
echo "2. Tiếng Việt"
read -p "Enter choice / Nhập lựa chọn: " LANG_CHOICE

if [[ "$LANG_CHOICE" == "1" ]]; then
    LANG="en"
    MSG_ROOT="Please run as root!"
    MSG_ENTER_SERVER_NAME="Enter server name (Ex: htn217217): "
    MSG_ENTER_TOPIC_NAME="Enter Kafka topic name (Ex: audit_logs): "
    MSG_SERVER_NAME_EMPTY="Server name cannot be empty!"
    MSG_TOPIC_NAME_EMPTY="Kafka topic name cannot be empty!"
    MSG_INSTALLING="Installing necessary packages and configuring audit log..."
    MSG_UNSUPPORTED_OS="Unsupported OS"
    MSG_SETUP_COMPLETE="Audit and log forwarding setup completed. Status of services:"
elif [[ "$LANG_CHOICE" == "2" ]]; then
    LANG="vi"
    MSG_ROOT="Vui lòng chạy script với quyền root!"
    MSG_ENTER_SERVER_NAME="Nhập tên server (ví dụ: htn217217): "
    MSG_ENTER_TOPIC_NAME="Nhập Kafka topic name (ví dụ: audit_logs): "
    MSG_SERVER_NAME_EMPTY="Tên server không được để trống!"
    MSG_TOPIC_NAME_EMPTY="Kafka topic name không được để trống!"
    MSG_INSTALLING="Đang cài đặt các gói cần thiết và cấu hình log audit..."
    MSG_UNSUPPORTED_OS="Hệ điều hành không được hỗ trợ"
    MSG_SETUP_COMPLETE="Hoàn tất cấu hình ghi log và đẩy log. Trạng thái dịch vụ:"
else
    echo "Invalid choice / Lựa chọn không hợp lệ"
    exit 1
fi

# Kafka server configuration
KAFKA_SERVER="broker01"
KAFKA_PORT="9092"

# Yêu cầu người dùng nhập tên server
read -p "$MSG_ENTER_SERVER_NAME" CUSTOM_SERVER_NAME
if [[ -z "$CUSTOM_SERVER_NAME" ]]; then
    echo "$MSG_SERVER_NAME_EMPTY"
    exit 1
fi

# Yêu cầu người dùng nhập Kafka topic name
read -p "$MSG_ENTER_TOPIC_NAME" KAFKA_TOPIC
if [[ -z "$KAFKA_TOPIC" ]]; then
    echo "$MSG_TOPIC_NAME_EMPTY"
    exit 1
fi

# Kiểm tra hệ điều hành
if [ -f /etc/os-release ]; then
    . /etc/os-release
    OS=$ID
else
    echo "$MSG_UNSUPPORTED_OS"
    exit 1
fi

# Cài đặt gói cần thiết và cấu hình audit log
echo "$MSG_INSTALLING"
if [[ "$OS" == "centos" ]]; then
    yum install -y audit auditd rsyslog rsyslog-kafka
    systemctl enable auditd
    systemctl enable rsyslog
    systemctl start auditd
    systemctl start rsyslog

elif [[ "$OS" == "ubuntu" ]]; then
    apt update
    apt install -y auditd audispd-plugins rsyslog rsyslog-kafka
    systemctl enable auditd
    systemctl enable rsyslog
    systemctl start auditd
    systemctl start rsyslog
else
    echo "$MSG_UNSUPPORTED_OS"
    exit 1
fi

# Cấu hình auditd ghi log dưới dạng JSON
cat << EOF > /etc/audit/auditd.conf
log_file = /var/log/audit/audit.log
log_format = ENRICHED
max_log_file = 8
max_log_file_action = ROTATE
enable_krb5 = no
declare_error_on_failure = no
EOF
systemctl restart auditd

# Tích hợp JSON log vào rsyslog
cat << EOF > /etc/rsyslog.d/kafka-audit.conf
module(load="imfile")
module(load="omkafka")

# Đọc log JSON từ file audit.log
input (
    type="imfile"
    File="/var/log/audit/audit.log"
    Tag="$KAFKA_TOPIC"
    freshStartTail="on"
    reopenOnTruncate="on"
    Severity="error"
)

# Template định dạng JSON gửi tới Kafka
template (
    name="JSONFormat"
    type="string"
    string="server:$CUSTOM_SERVER_NAME, log: %msg%"
)

# Gửi log JSON tới Kafka
if (\$syslogtag == '$KAFKA_TOPIC') then {
    action(
        name="$KAFKA_TOPIC"
        broker=["$KAFKA_SERVER:$KAFKA_PORT"]
        type="omkafka"
        topic="$KAFKA_TOPIC"
        partitions.auto="on"
        template="JSONFormat"
        confParam=[
            "socket.keepalive.enable=true",
            "log.connection.close=false",
            "queue.buffering.max.messages=400000"
        ]
        closeTimeout="60000"
        queue.size="500000"
        queue.saveonshutdown="on"
        action.resumeRetryCount="-1"
        action.resumeInterval="60"
    )
}
EOF
systemctl restart rsyslog

# Kiểm tra trạng thái dịch vụ
echo "$MSG_SETUP_COMPLETE"
systemctl status auditd | grep Active
systemctl status rsyslog | grep Active

# Kiểm tra log JSON
echo "Kiểm tra log JSON trong /var/log/audit/audit.log"
tail -n 5 /var/log/audit/audit.log

