go-worm/lib/keylog/writer.go

package keylog

import (
	"os"
	"path/filepath"
	"time"
)

type Writer struct {
	dir       string
	file      *os.File
	bucket    string
	curWindow string
	curSource string
}

func NewWriter(dir string) (*Writer, error) {
	if err := os.MkdirAll(dir, 0o700); err != nil {
		return nil, err
	}
	return &Writer{dir: dir}, nil
}

func (w *Writer) Write(event Event) error {
	if event.Text == "" {
		return nil
	}
	bucket := HourBucket(event.Time)
	if bucket != w.bucket {
		if err := w.rotate(bucket); err != nil {
			return err
		}
	}

	window := normalizeWindow(event.Window)
	source := eventSource(event.Injected)
	if window != w.curWindow || source != w.curSource {
		if w.curWindow != "" {
			if _, err := w.file.WriteString("\n\n"); err != nil {
				return err
			}
		}
		if _, err := w.file.WriteString(sectionHeader(source, window)); err != nil {
			return err
		}
		w.curWindow = window
		w.curSource = source
	}

	_, err := w.file.WriteString(event.Text)
	return err
}

func (w *Writer) rotate(bucket string) error {
	if w.file != nil {
		if err := w.file.Close(); err != nil {
			return err
		}
		w.file = nil
	}
	path := filepath.Join(w.dir, LogFilename(bucket))
	file, err := os.OpenFile(path, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o600)
	if err != nil {
		return err
	}
	w.file = file
	w.bucket = bucket
	w.curWindow = ""
	w.curSource = ""
	return nil
}

func (w *Writer) Close() error {
	if w.file == nil {
		return nil
	}
	err := w.file.Close()
	w.file = nil
	w.bucket = ""
	w.curWindow = ""
	w.curSource = ""
	return err
}

func PruneOldLogs(dir string, retentionDays int) error {
	if retentionDays <= 0 {
		return nil
	}
	cutoff := time.Now().AddDate(0, 0, -retentionDays)
	entries, err := os.ReadDir(dir)
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return err
	}
	for _, entry := range entries {
		if entry.IsDir() || !ValidLogFilename(entry.Name()) {
			continue
		}
		info, err := entry.Info()
		if err != nil {
			continue
		}
		if info.ModTime().Before(cutoff) {
			_ = os.Remove(filepath.Join(dir, entry.Name()))
		}
	}
	return nil
}