go-logrotate/main.go

package main

import (
	"compress/gzip"
	"flag"
	"fmt"
	"io"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/spf13/viper"
)

// GlobalSettings represents global configuration settings
type GlobalSettings struct {
	Compress       bool   `mapstructure:"compress"`
	CompressFormat string `mapstructure:"compressFormat"`
	CreateNew      bool   `mapstructure:"createNew"`
	CreateMode     string `mapstructure:"createMode"`
}

// LogConfig represents configuration for a specific log file
type LogConfig struct {
	Path           string        `mapstructure:"path"`
	RotateOnSize   string        `mapstructure:"rotateOnSize"` // This will be parsed to get the size in bytes
	RetentionCount int           `mapstructure:"retentionCount"`
	RotateOnAge    time.Duration `mapstructure:"rotateOnAge"` // Assume the input is compatible with time.ParseDuration
	Compress       bool          `mapstructure:"compress"`    // Optional: per-log override
}

// Configuration holds the entire configuration
type Configuration struct {
	Settings GlobalSettings `mapstructure:"settings"`
	Logs     []LogConfig    `mapstructure:"logs"`
}

// ConvertSizeToBytes converts size strings like "10MB" to bytes
func ConvertSizeToBytes(sizeStr string) (int64, error) {
	sizeStr = strings.ToUpper(sizeStr)
	multiplier := int64(1)

	if strings.HasSuffix(sizeStr, "KB") {
		multiplier = 1024
		sizeStr = strings.TrimSuffix(sizeStr, "KB")
	} else if strings.HasSuffix(sizeStr, "MB") {
		multiplier = 1024 * 1024
		sizeStr = strings.TrimSuffix(sizeStr, "MB")
	} else if strings.HasSuffix(sizeStr, "GB") {
		multiplier = 1024 * 1024 * 1024
		sizeStr = strings.TrimSuffix(sizeStr, "GB")
	}

	size, err := strconv.ParseInt(sizeStr, 10, 64)
	if err != nil {
		return 0, err
	}

	return size * multiplier, nil
}

func compressFile(sourcePath, destPath string) error {
	inputFile, err := os.Open(sourcePath)
	if err != nil {
		return err
	}
	defer inputFile.Close() // Ensure the file is closed as soon as the function returns

	outputFile, err := os.Create(destPath)
	if err != nil {
		return err
	}
	defer outputFile.Close() // Ensure the file is closed as soon as the function returns

	gzWriter := gzip.NewWriter(outputFile)
	defer gzWriter.Close() // Ensure the gzip writer is closed

	_, err = io.Copy(gzWriter, inputFile)
	if err != nil {
		return err
	}

	// Ensure all writes are flushed and the writer is closed before deleting the source file
	if err := gzWriter.Close(); err != nil {
		return err
	}
	if err := outputFile.Close(); err != nil {
		return err
	}
	if err := inputFile.Close(); err != nil {
		return err
	}

	// After confirming the file has been compressed and closed, it's safe to remove the source file
	return os.Remove(sourcePath)
}

func shouldRotateBasedOnAge(filePath string, rotateAge time.Duration) (bool, error) {
	fileInfo, err := os.Stat(filePath)
	if err != nil {
		return false, err
	}

	// Calculate the file's age by subtracting its modification time from the current time
	fileAge := time.Since(fileInfo.ModTime())

	// Return true if the file's age exceeds the rotateAge threshold
	return fileAge >= rotateAge, nil
}

// rotateFile rotates the specified log file according to the retention policy
func rotateFile(config LogConfig, verbose bool) error {
	// First, check if the current log file exceeds the size threshold
	rotateSize, err := ConvertSizeToBytes(config.RotateOnSize)
	if err != nil {
		return fmt.Errorf("error converting rotation size: %w", err)
	}

	fileInfo, err := os.Stat(config.Path)
	if err != nil {
		return fmt.Errorf("error stating current log file: %w", err)
	}

	shouldRotate := fileInfo.Size() >= rotateSize

	// Determine if the file needs rotation based on age, if rotateOnAge is specified
	if config.RotateOnAge > 0 {
		rotateBasedOnAge, err := shouldRotateBasedOnAge(config.Path, config.RotateOnAge)
		if err != nil {
			return fmt.Errorf("error checking if log file should be rotated based on age: %w", err)
		}
		shouldRotate = shouldRotate || rotateBasedOnAge
	}

	if shouldRotate {

		// File needs rotation; handle retention for existing rotated files
		compressExtension := ".gz"
		shouldCompress := config.Compress || (!config.Compress && viper.GetBool("settings.compress"))

		// Rotate existing files, considering compression
		for i := config.RetentionCount; i > 0; i-- {
			currentFile := fmt.Sprintf("%s.%d", config.Path, i)
			nextFileNumber := i + 1
			nextFile := fmt.Sprintf("%s.%d", config.Path, nextFileNumber)

			if shouldCompress {
				currentFile += compressExtension
				nextFile += compressExtension
			}

			if _, err := os.Stat(currentFile); !os.IsNotExist(err) {
				if i == config.RetentionCount {
					// Remove the oldest file
					os.Remove(currentFile)
				} else {
					// Rename (rotate) the file to the next higher number
					os.Rename(currentFile, nextFile)
				}
			}
		}

		// Rotate the current log file
		rotatedFileName := config.Path + ".1"
		if shouldCompress {
			// If compression is enabled, compress the file
			compressedFileName := rotatedFileName + compressExtension
			if err := compressFile(config.Path, compressedFileName); err != nil {
				return fmt.Errorf("failed to compress log file: %w", err)
			}
		} else {
			// If not compressing, simply rename the current log file
			if err := os.Rename(config.Path, rotatedFileName); err != nil {
				return fmt.Errorf("failed to rotate log file: %w", err)
			}
		}

		// If createNew is true, create a new log file with the original log file's name
		if viper.GetBool("settings.createNew") {
			mode, _ := strconv.ParseUint(viper.GetString("settings.createMode"), 0, 32)
			newFile, err := os.OpenFile(config.Path, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, os.FileMode(mode))
			if err != nil {
				return fmt.Errorf("failed to create new log file: %w", err)
			}
			if err := newFile.Close(); err != nil {
				return fmt.Errorf("failed to close new log file: %w", err)
			}
		}
	}
	return nil
}

func main() {
	var verbose bool
	flag.BoolVar(&verbose, "verbose", false, "enable verbose output")
	flag.Parse()

	viper.SetConfigName("config.yaml")
	viper.SetConfigType("yaml")
	viper.AddConfigPath(".")

	var config Configuration

	if err := viper.ReadInConfig(); err != nil {
		fmt.Printf("Error reading config file, %s", err)
		os.Exit(1)
	}

	if err := viper.Unmarshal(&config); err != nil {
		fmt.Printf("Unable to decode into struct, %s", err)
		os.Exit(1)
	}

	if verbose {
		fmt.Printf("Global Compress: %t\n", config.Settings.Compress)
		for _, log := range config.Logs {
			fmt.Printf("Log Path: %s, RotateOnSize: %s, RotateOnAge: %s\n", log.Path, log.RotateOnSize, log.RotateOnAge)
		}
	}

	// Process each log file according to its configuration
	for _, logConfig := range config.Logs {
		if err := rotateFile(logConfig, verbose); err != nil {
			fmt.Printf("Error rotating file %s: %s\n", logConfig.Path, err)
			continue // Proceed to the next log file instead of stopping
		}

		if verbose {
			fmt.Printf("Processed log file: %s\n", logConfig.Path)
		}
	}
}