#!/bin/bash

# Detect OS Type using /etc/os-release
if grep -q "Amazon Linux" /etc/os-release; then
    if grep -q "2" /etc/os-release; then
        OS_TYPE="AmazonLinux2"
    elif grep -q "2023" /etc/os-release; then
        OS_TYPE="AmazonLinux2023"
    else
        echo "Unsupported Amazon Linux version. Exiting."
        exit 1
    fi
elif grep -q "Ubuntu 22.04" /etc/os-release; then
    OS_TYPE="UbuntuJammy"
elif grep -q "Ubuntu 23.04" /etc/os-release; then
    OS_TYPE="UbuntuNoble"
elif grep -q "Ubuntu 20.04" /etc/os-release; then
    OS_TYPE="UbuntuFocal"
else
    echo "Unsupported OS. Exiting."
    exit 1
fi

# Function to retrieve Instance ID using IMDSv2, fallback to IMDSv1 if needed
get_instance_id() {
    # Try to use IMDSv2 first
    TOKEN=$(curl -X PUT -H "X-aws-ec2-metadata-token-ttl-seconds: 21600" -s "http://169.254.169.254/latest/api/token")
    if [ -n "$TOKEN" ]; then
        INSTANCE_ID=$(curl -H "X-aws-ec2-metadata-token: $TOKEN" -s "http://169.254.169.254/latest/meta-data/instance-id")
    fi

    # If IMDSv2 fails, fall back to IMDSv1
    if [ -z "$INSTANCE_ID" ]; then
        INSTANCE_ID=$(curl -s "http://169.254.169.254/latest/meta-data/instance-id")
    fi

    # If both methods fail, display an error
    if [ -z "$INSTANCE_ID" ]; then
        echo "Unable to retrieve instance ID. Verify that the script is running on an EC2 instance with metadata access enabled."
        exit 1
    fi
}

# Retrieve Instance ID
get_instance_id

# Install Fluentd based on OS Type
case "$OS_TYPE" in
    "AmazonLinux2")
        curl -fsSL https://toolbelt.treasuredata.com/sh/install-amazon2-fluent-package5-lts.sh | sh
        ;;
    "AmazonLinux2023")
        curl -fsSL https://toolbelt.treasuredata.com/sh/install-amazon2023-fluent-package5-lts.sh | sh
        ;;
    "UbuntuJammy")
        curl -fsSL https://toolbelt.treasuredata.com/sh/install-ubuntu-jammy-fluent-package5-lts.sh | sh
        ;;
    "UbuntuNoble")
        curl -fsSL https://toolbelt.treasuredata.com/sh/install-ubuntu-noble-fluent-package5-lts.sh | sh
        ;;
    "UbuntuFocal")
        curl -fsSL https://toolbelt.treasuredata.com/sh/install-ubuntu-focal-fluent-package5-lts.sh | sh
        ;;
    *)
        echo "Unsupported OS. Exiting."
        exit 1
        ;;
esac

# Install the fluent-plugin-ec2_metadata plugin
# Start Fluentd Service
sudo systemctl start fluentd.service

# Enable Fluentd service to start on boot
sudo systemctl enable fluentd.service

# Check status of Fluentd service
sudo systemctl status fluentd.service

# Update the fluentd.conf file content (without removing it)
cat <<EOF | sudo tee /etc/fluent/fluentd.conf > /dev/null
<source>
  @type tail
  path /var/log/auth.log,/var/log/syslog,/var/log/secure,/var/log/messages  # Paths to auth and system logs
  pos_file /var/log/fluentd/pos/linux_logs.pos  # Position file to track the last read line
  tag linux.server  # Tag for filtering logs
  format none  # Use "none" if logs are raw text; use a parser if required
  read_from_head true  # Start reading from the beginning of the file
</source>

<match linux.server>
  @type s3
  s3_bucket <s3-bucket-name>             # Replace with your S3 bucket name
  s3_region ap-south-1
  path linux/$INSTANCE_ID/%Y/%m/%d/  # Use dynamic values from the record to organize logs by instance
  time_slice_format %Y%m%d%H               # Log chunk format: year, month, day, hour
  time_slice_wait 5m                      # Wait before finalizing a log chunk
  buffer_path /var/log/fluentd/buffer/s3   # Buffer directory for Fluentd
  <format>
    @type json                             # Store logs in JSON format
  </format>
  <buffer>
    @type file                             # File-based buffer
    path /var/log/fluentd/buffer/s3
    flush_mode interval                    # Flush logs at intervals
    flush_interval 1m                      # Flush every 1 minute
    retry_max_times 10                     # Retry up to 10 times on failure
    retry_wait 30   
    expire_interval 1h                    # Delete buffer files older than 5 minutes
  </buffer>
</match>
EOF

# Create necessary directories and set permissions
sudo mkdir -p /var/log/fluentd/buffer/s3
sudo chown -R fluentd:fluentd /var/log/fluentd
sudo chmod -R 777 /var/log/

# Restart Fluentd service
sudo systemctl restart fluentd.service

# Check status of Fluentd service again
sudo systemctl status fluentd.service