Install and configure Loki for centralized log aggregation with Grafana integration

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

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# Configuration
LOKI_VERSION="2.9.3"
LOKI_USER="loki"
LOKI_CONFIG_DIR="/etc/loki"
LOKI_DATA_DIR="/var/lib/loki"
LOKI_LOG_DIR="/var/log/loki"

# Function to print colored output
print_status() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1" >&2
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

# Usage function
usage() {
    echo "Usage: $0"
    echo "Install and configure Loki for centralized log aggregation"
    exit 1
}

# Cleanup function for rollback
cleanup() {
    print_error "Installation failed. Cleaning up..."
    systemctl stop loki promtail 2>/dev/null || true
    systemctl disable loki promtail 2>/dev/null || true
    rm -f /etc/systemd/system/loki.service /etc/systemd/system/promtail.service
    rm -f /usr/local/bin/loki /usr/local/bin/promtail
    userdel $LOKI_USER 2>/dev/null || true
    rm -rf $LOKI_CONFIG_DIR $LOKI_DATA_DIR $LOKI_LOG_DIR
    systemctl daemon-reload
}

trap cleanup ERR

# Check if running as root or with sudo
check_privileges() {
    if [[ $EUID -ne 0 ]]; then
        print_error "This script must be run as root or with sudo"
        exit 1
    fi
}

# Auto-detect distribution
detect_distro() {
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        case "$ID" in
            ubuntu|debian) 
                PKG_MGR="apt"
                PKG_INSTALL="apt install -y"
                PKG_UPDATE="apt update && apt upgrade -y"
                ;;
            almalinux|rocky|centos|rhel|ol|fedora) 
                PKG_MGR="dnf"
                PKG_INSTALL="dnf install -y"
                PKG_UPDATE="dnf update -y"
                ;;
            amzn) 
                PKG_MGR="yum"
                PKG_INSTALL="yum install -y"
                PKG_UPDATE="yum update -y"
                ;;
            *) 
                print_error "Unsupported distribution: $ID"
                exit 1
                ;;
        esac
    else
        print_error "Cannot detect distribution. /etc/os-release not found."
        exit 1
    fi
}

# Update system packages
update_system() {
    echo "[1/8] Updating system packages..."
    $PKG_UPDATE
    print_status "System packages updated"
}

# Install dependencies
install_dependencies() {
    echo "[2/8] Installing required dependencies..."
    $PKG_INSTALL wget unzip curl
    print_status "Dependencies installed"
}

# Create Loki user and directories
create_user_and_dirs() {
    echo "[3/8] Creating Loki user and directories..."
    
    # Create system user
    if ! id "$LOKI_USER" &>/dev/null; then
        useradd --system --no-create-home --shell /bin/false $LOKI_USER
    fi
    
    # Create directories
    mkdir -p $LOKI_CONFIG_DIR $LOKI_DATA_DIR $LOKI_LOG_DIR
    mkdir -p $LOKI_DATA_DIR/{chunks,rules,compactor}
    
    # Set ownership and permissions
    chown -R $LOKI_USER:$LOKI_USER $LOKI_DATA_DIR $LOKI_LOG_DIR
    chmod 755 $LOKI_CONFIG_DIR $LOKI_DATA_DIR $LOKI_LOG_DIR
    
    print_status "User and directories created"
}

# Download and install Loki
install_loki() {
    echo "[4/8] Downloading and installing Loki..."
    
    cd /tmp
    wget -q "https://github.com/grafana/loki/releases/download/v${LOKI_VERSION}/loki-linux-amd64.zip"
    unzip -q loki-linux-amd64.zip
    mv loki-linux-amd64 /usr/local/bin/loki
    chmod 755 /usr/local/bin/loki
    rm loki-linux-amd64.zip
    
    print_status "Loki binary installed"
}

# Download and install Promtail
install_promtail() {
    echo "[5/8] Downloading and installing Promtail..."
    
    cd /tmp
    wget -q "https://github.com/grafana/loki/releases/download/v${LOKI_VERSION}/promtail-linux-amd64.zip"
    unzip -q promtail-linux-amd64.zip
    mv promtail-linux-amd64 /usr/local/bin/promtail
    chmod 755 /usr/local/bin/promtail
    rm promtail-linux-amd64.zip
    
    print_status "Promtail binary installed"
}

# Configure Loki
configure_loki() {
    echo "[6/8] Configuring Loki..."
    
    cat > $LOKI_CONFIG_DIR/loki.yml << 'EOF'
auth_enabled: false

server:
  http_listen_port: 3100
  grpc_listen_port: 9096
  log_level: info

common:
  path_prefix: /var/lib/loki
  storage:
    filesystem:
      chunks_directory: /var/lib/loki/chunks
      rules_directory: /var/lib/loki/rules
  replication_factor: 1
  ring:
    instance_addr: 127.0.0.1
    kvstore:
      store: inmemory

query_range:
  results_cache:
    cache:
      embedded_cache:
        enabled: true
        max_size_mb: 100

schema_config:
  configs:
    - from: 2020-10-24
      store: boltdb-shipper
      object_store: filesystem
      schema: v11
      index:
        prefix: index_
        period: 24h

limits_config:
  enforce_metric_name: false
  reject_old_samples: true
  reject_old_samples_max_age: 168h
  ingestion_rate_mb: 16
  ingestion_burst_size_mb: 32
  max_query_parallelism: 32
  retention_period: 744h

compactor:
  working_directory: /var/lib/loki/compactor
  shared_store: filesystem
  compaction_interval: 10m
  retention_enabled: true
  retention_delete_delay: 2h
  retention_delete_worker_count: 150
EOF

    cat > $LOKI_CONFIG_DIR/promtail.yml << 'EOF'
server:
  http_listen_port: 9080
  grpc_listen_port: 0

positions:
  filename: /var/lib/loki/positions.yaml

clients:
  - url: http://localhost:3100/loki/api/v1/push

scrape_configs:
  - job_name: system
    static_configs:
      - targets:
          - localhost
        labels:
          job: varlogs
          __path__: /var/log/*log
          
  - job_name: syslog
    static_configs:
      - targets:
          - localhost
        labels:
          job: syslog
          __path__: /var/log/syslog
          
  - job_name: auth
    static_configs:
      - targets:
          - localhost
        labels:
          job: auth
          __path__: /var/log/auth.log
EOF

    chmod 644 $LOKI_CONFIG_DIR/*.yml
    print_status "Loki and Promtail configured"
}

# Create systemd services
create_systemd_services() {
    echo "[7/8] Creating systemd services..."
    
    # Loki service
    cat > /etc/systemd/system/loki.service << EOF
[Unit]
Description=Loki service
After=network.target

[Service]
Type=simple
User=$LOKI_USER
Group=$LOKI_USER
ExecStart=/usr/local/bin/loki -config.file $LOKI_CONFIG_DIR/loki.yml
Restart=on-failure
RestartSec=20
StandardOutput=journal
StandardError=journal
SyslogIdentifier=loki
KillMode=mixed
KillSignal=SIGINT

[Install]
WantedBy=multi-user.target
EOF

    # Promtail service
    cat > /etc/systemd/system/promtail.service << EOF
[Unit]
Description=Promtail service
After=network.target

[Service]
Type=simple
User=root
ExecStart=/usr/local/bin/promtail -config.file $LOKI_CONFIG_DIR/promtail.yml
Restart=on-failure
RestartSec=20
StandardOutput=journal
StandardError=journal
SyslogIdentifier=promtail
KillMode=mixed
KillSignal=SIGINT

[Install]
WantedBy=multi-user.target
EOF

    systemctl daemon-reload
    systemctl enable loki promtail
    systemctl start loki promtail
    
    print_status "Systemd services created and started"
}

# Verify installation
verify_installation() {
    echo "[8/8] Verifying installation..."
    
    # Check if services are running
    if systemctl is-active --quiet loki; then
        print_status "Loki service is running"
    else
        print_error "Loki service is not running"
        exit 1
    fi
    
    if systemctl is-active --quiet promtail; then
        print_status "Promtail service is running"
    else
        print_error "Promtail service is not running"
        exit 1
    fi
    
    # Check if Loki is responding
    sleep 5
    if curl -s http://localhost:3100/ready | grep -q "ready"; then
        print_status "Loki API is responding"
    else
        print_warning "Loki API may not be ready yet"
    fi
    
    print_status "Installation completed successfully!"
    echo ""
    echo "Loki is now running on http://localhost:3100"
    echo "Promtail is collecting logs and sending them to Loki"
    echo ""
    echo "To integrate with Grafana, add Loki as a data source:"
    echo "URL: http://localhost:3100"
}

# Main execution
main() {
    print_status "Starting Loki installation..."
    
    check_privileges
    detect_distro
    update_system
    install_dependencies
    create_user_and_dirs
    install_loki
    install_promtail
    configure_loki
    create_systemd_services
    verify_installation
}

# Run main function
main "$@"