oubliette/internal/storage/memstore.go

package storage

import (
	"context"
	"sort"
	"sync"
	"time"

	"github.com/google/uuid"
)

// MemoryStore is an in-memory implementation of Store for use in tests.
type MemoryStore struct {
	mu            sync.Mutex
	LoginAttempts []LoginAttempt
	Sessions      map[string]*Session
	SessionLogs   []SessionLog
	SessionEvents []SessionEvent
}

// NewMemoryStore returns a new empty MemoryStore.
func NewMemoryStore() *MemoryStore {
	return &MemoryStore{
		Sessions: make(map[string]*Session),
	}
}

func (m *MemoryStore) RecordLoginAttempt(_ context.Context, username, password, ip, country string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	now := time.Now().UTC()
	for i := range m.LoginAttempts {
		a := &m.LoginAttempts[i]
		if a.Username == username && a.Password == password && a.IP == ip {
			a.Count++
			a.LastSeen = now
			a.Country = country
			return nil
		}
	}

	m.LoginAttempts = append(m.LoginAttempts, LoginAttempt{
		ID:        int64(len(m.LoginAttempts) + 1),
		Username:  username,
		Password:  password,
		IP:        ip,
		Country:   country,
		Count:     1,
		FirstSeen: now,
		LastSeen:  now,
	})
	return nil
}

func (m *MemoryStore) CreateSession(_ context.Context, ip, username, shellName, country string) (string, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	id := uuid.New().String()
	now := time.Now().UTC()
	m.Sessions[id] = &Session{
		ID:          id,
		IP:          ip,
		Country:     country,
		Username:    username,
		ShellName:   shellName,
		ConnectedAt: now,
	}
	return id, nil
}

func (m *MemoryStore) EndSession(_ context.Context, sessionID string, disconnectedAt time.Time) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if s, ok := m.Sessions[sessionID]; ok {
		t := disconnectedAt.UTC()
		s.DisconnectedAt = &t
	}
	return nil
}

func (m *MemoryStore) UpdateHumanScore(_ context.Context, sessionID string, score float64) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if s, ok := m.Sessions[sessionID]; ok {
		s.HumanScore = &score
	}
	return nil
}

func (m *MemoryStore) SetExecCommand(_ context.Context, sessionID string, command string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if s, ok := m.Sessions[sessionID]; ok {
		s.ExecCommand = &command
	}
	return nil
}

func (m *MemoryStore) AppendSessionLog(_ context.Context, sessionID, input, output string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	m.SessionLogs = append(m.SessionLogs, SessionLog{
		ID:        int64(len(m.SessionLogs) + 1),
		SessionID: sessionID,
		Timestamp: time.Now().UTC(),
		Input:     input,
		Output:    output,
	})
	return nil
}

func (m *MemoryStore) GetSession(_ context.Context, sessionID string) (*Session, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	s, ok := m.Sessions[sessionID]
	if !ok {
		return nil, nil
	}
	copy := *s
	return &copy, nil
}

func (m *MemoryStore) GetSessionLogs(_ context.Context, sessionID string) ([]SessionLog, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	var logs []SessionLog
	for _, l := range m.SessionLogs {
		if l.SessionID == sessionID {
			logs = append(logs, l)
		}
	}
	sort.Slice(logs, func(i, j int) bool {
		return logs[i].Timestamp.Before(logs[j].Timestamp)
	})
	return logs, nil
}

func (m *MemoryStore) AppendSessionEvents(_ context.Context, events []SessionEvent) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	m.SessionEvents = append(m.SessionEvents, events...)
	return nil
}

func (m *MemoryStore) GetSessionEvents(_ context.Context, sessionID string) ([]SessionEvent, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	var events []SessionEvent
	for _, e := range m.SessionEvents {
		if e.SessionID == sessionID {
			events = append(events, e)
		}
	}
	return events, nil
}

func (m *MemoryStore) DeleteRecordsBefore(_ context.Context, cutoff time.Time) (int64, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	var total int64

	// Delete old login attempts.
	kept := m.LoginAttempts[:0]
	for _, a := range m.LoginAttempts {
		if a.LastSeen.Before(cutoff) {
			total++
		} else {
			kept = append(kept, a)
		}
	}
	m.LoginAttempts = kept

	// Delete old sessions and their logs.
	for id, s := range m.Sessions {
		if s.ConnectedAt.Before(cutoff) {
			delete(m.Sessions, id)
			total++
		}
	}

	keptLogs := m.SessionLogs[:0]
	for _, l := range m.SessionLogs {
		if _, ok := m.Sessions[l.SessionID]; ok {
			keptLogs = append(keptLogs, l)
		} else {
			total++
		}
	}
	m.SessionLogs = keptLogs

	keptEvents := m.SessionEvents[:0]
	for _, e := range m.SessionEvents {
		if _, ok := m.Sessions[e.SessionID]; ok {
			keptEvents = append(keptEvents, e)
		} else {
			total++
		}
	}
	m.SessionEvents = keptEvents

	return total, nil
}

func (m *MemoryStore) GetDashboardStats(_ context.Context) (*DashboardStats, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	stats := &DashboardStats{}
	ips := make(map[string]struct{})
	for _, a := range m.LoginAttempts {
		stats.TotalAttempts += int64(a.Count)
		ips[a.IP] = struct{}{}
	}
	stats.UniqueIPs = int64(len(ips))
	stats.TotalSessions = int64(len(m.Sessions))
	for _, s := range m.Sessions {
		if s.DisconnectedAt == nil {
			stats.ActiveSessions++
		}
	}
	return stats, nil
}

func (m *MemoryStore) GetTopUsernames(_ context.Context, limit int) ([]TopEntry, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.topN("username", limit), nil
}

func (m *MemoryStore) GetTopPasswords(_ context.Context, limit int) ([]TopEntry, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.topN("password", limit), nil
}

func (m *MemoryStore) GetTopIPs(_ context.Context, limit int) ([]TopEntry, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	type ipInfo struct {
		count   int64
		country string
	}
	agg := make(map[string]*ipInfo)
	for _, a := range m.LoginAttempts {
		info, ok := agg[a.IP]
		if !ok {
			info = &ipInfo{}
			agg[a.IP] = info
		}
		info.count += int64(a.Count)
		if a.Country != "" {
			info.country = a.Country
		}
	}

	entries := make([]TopEntry, 0, len(agg))
	for ip, info := range agg {
		entries = append(entries, TopEntry{Value: ip, Country: info.country, Count: info.count})
	}
	sort.Slice(entries, func(i, j int) bool {
		return entries[i].Count > entries[j].Count
	})
	if limit > 0 && len(entries) > limit {
		entries = entries[:limit]
	}
	return entries, nil
}

func (m *MemoryStore) GetTopCountries(_ context.Context, limit int) ([]TopEntry, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	counts := make(map[string]int64)
	for _, a := range m.LoginAttempts {
		if a.Country == "" {
			continue
		}
		counts[a.Country] += int64(a.Count)
	}

	entries := make([]TopEntry, 0, len(counts))
	for k, v := range counts {
		entries = append(entries, TopEntry{Value: k, Count: v})
	}
	sort.Slice(entries, func(i, j int) bool {
		return entries[i].Count > entries[j].Count
	})
	if limit > 0 && len(entries) > limit {
		entries = entries[:limit]
	}
	return entries, nil
}

// topN aggregates login attempts by the given field and returns the top N. Must be called with m.mu held.
func (m *MemoryStore) topN(field string, limit int) []TopEntry {
	counts := make(map[string]int64)
	for _, a := range m.LoginAttempts {
		var key string
		switch field {
		case "username":
			key = a.Username
		case "password":
			key = a.Password
		case "ip":
			key = a.IP
		}
		counts[key] += int64(a.Count)
	}

	entries := make([]TopEntry, 0, len(counts))
	for k, v := range counts {
		entries = append(entries, TopEntry{Value: k, Count: v})
	}
	sort.Slice(entries, func(i, j int) bool {
		return entries[i].Count > entries[j].Count
	})
	if limit > 0 && len(entries) > limit {
		entries = entries[:limit]
	}
	return entries
}

func (m *MemoryStore) GetRecentSessions(_ context.Context, limit int, activeOnly bool) ([]Session, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	return m.collectSessions(limit, activeOnly, DashboardFilter{}), nil
}

func (m *MemoryStore) GetFilteredSessions(_ context.Context, limit int, activeOnly bool, f DashboardFilter) ([]Session, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	return m.collectSessions(limit, activeOnly, f), nil
}

// collectSessions gathers sessions matching filter criteria. Must be called with m.mu held.
func (m *MemoryStore) collectSessions(limit int, activeOnly bool, f DashboardFilter) []Session {
	// Compute event counts and input bytes per session.
	eventCounts := make(map[string]int)
	inputBytes := make(map[string]int64)
	for _, e := range m.SessionEvents {
		eventCounts[e.SessionID]++
		if e.Direction == 0 {
			inputBytes[e.SessionID] += int64(len(e.Data))
		}
	}

	var sessions []Session
	for _, s := range m.Sessions {
		if activeOnly && s.DisconnectedAt != nil {
			continue
		}
		if !matchesSessionFilter(s, f) {
			continue
		}
		sess := *s
		sess.EventCount = eventCounts[s.ID]
		sess.InputBytes = inputBytes[s.ID]
		sessions = append(sessions, sess)
	}

	if f.SortBy == "input_bytes" {
		sort.Slice(sessions, func(i, j int) bool {
			return sessions[i].InputBytes > sessions[j].InputBytes
		})
	} else {
		sort.Slice(sessions, func(i, j int) bool {
			return sessions[i].ConnectedAt.After(sessions[j].ConnectedAt)
		})
	}

	if limit > 0 && len(sessions) > limit {
		sessions = sessions[:limit]
	}
	return sessions
}

// matchesSessionFilter returns true if the session matches the given filter.
func matchesSessionFilter(s *Session, f DashboardFilter) bool {
	if f.Since != nil && s.ConnectedAt.Before(*f.Since) {
		return false
	}
	if f.Until != nil && s.ConnectedAt.After(*f.Until) {
		return false
	}
	if f.IP != "" && s.IP != f.IP {
		return false
	}
	if f.Country != "" && s.Country != f.Country {
		return false
	}
	if f.Username != "" && s.Username != f.Username {
		return false
	}
	if f.HumanScoreAboveZero {
		if s.HumanScore == nil || *s.HumanScore <= 0 {
			return false
		}
	}
	return true
}

func (m *MemoryStore) GetTopExecCommands(_ context.Context, limit int) ([]TopEntry, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	counts := make(map[string]int64)
	for _, s := range m.Sessions {
		if s.ExecCommand != nil {
			counts[*s.ExecCommand]++
		}
	}

	entries := make([]TopEntry, 0, len(counts))
	for k, v := range counts {
		entries = append(entries, TopEntry{Value: k, Count: v})
	}
	sort.Slice(entries, func(i, j int) bool {
		return entries[i].Count > entries[j].Count
	})
	if limit > 0 && len(entries) > limit {
		entries = entries[:limit]
	}
	return entries, nil
}

func (m *MemoryStore) CloseActiveSessions(_ context.Context, disconnectedAt time.Time) (int64, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	var count int64
	t := disconnectedAt.UTC()
	for _, s := range m.Sessions {
		if s.DisconnectedAt == nil {
			s.DisconnectedAt = &t
			count++
		}
	}
	return count, nil
}

func (m *MemoryStore) GetAttemptsOverTime(_ context.Context, days int, since, until *time.Time) ([]TimeSeriesPoint, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	var cutoff time.Time
	if since != nil {
		cutoff = *since
	} else {
		cutoff = time.Now().UTC().AddDate(0, 0, -days)
	}

	counts := make(map[string]int64)
	for _, a := range m.LoginAttempts {
		if a.LastSeen.Before(cutoff) {
			continue
		}
		if until != nil && a.LastSeen.After(*until) {
			continue
		}
		day := a.LastSeen.Format("2006-01-02")
		counts[day] += int64(a.Count)
	}

	points := make([]TimeSeriesPoint, 0, len(counts))
	for day, count := range counts {
		t, _ := time.Parse("2006-01-02", day)
		points = append(points, TimeSeriesPoint{Timestamp: t, Count: count})
	}
	sort.Slice(points, func(i, j int) bool {
		return points[i].Timestamp.Before(points[j].Timestamp)
	})
	return points, nil
}

func (m *MemoryStore) GetHourlyPattern(_ context.Context, since, until *time.Time) ([]HourlyCount, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	hourCounts := make(map[int]int64)
	for _, a := range m.LoginAttempts {
		if since != nil && a.LastSeen.Before(*since) {
			continue
		}
		if until != nil && a.LastSeen.After(*until) {
			continue
		}
		hour := a.LastSeen.Hour()
		hourCounts[hour] += int64(a.Count)
	}

	counts := make([]HourlyCount, 0, len(hourCounts))
	for h, c := range hourCounts {
		counts = append(counts, HourlyCount{Hour: h, Count: c})
	}
	sort.Slice(counts, func(i, j int) bool {
		return counts[i].Hour < counts[j].Hour
	})
	return counts, nil
}

func (m *MemoryStore) GetCountryStats(_ context.Context) ([]CountryCount, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	counts := make(map[string]int64)
	for _, a := range m.LoginAttempts {
		if a.Country == "" {
			continue
		}
		counts[a.Country] += int64(a.Count)
	}

	result := make([]CountryCount, 0, len(counts))
	for country, count := range counts {
		result = append(result, CountryCount{Country: country, Count: count})
	}
	sort.Slice(result, func(i, j int) bool {
		return result[i].Count > result[j].Count
	})
	return result, nil
}

// matchesFilter returns true if the login attempt matches the given filter. Must be called with m.mu held.
func matchesFilter(a *LoginAttempt, f DashboardFilter) bool {
	if f.Since != nil && a.LastSeen.Before(*f.Since) {
		return false
	}
	if f.Until != nil && a.LastSeen.After(*f.Until) {
		return false
	}
	if f.IP != "" && a.IP != f.IP {
		return false
	}
	if f.Country != "" && a.Country != f.Country {
		return false
	}
	if f.Username != "" && a.Username != f.Username {
		return false
	}
	return true
}

func (m *MemoryStore) GetFilteredDashboardStats(_ context.Context, f DashboardFilter) (*DashboardStats, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	stats := &DashboardStats{}
	ips := make(map[string]struct{})
	for i := range m.LoginAttempts {
		a := &m.LoginAttempts[i]
		if !matchesFilter(a, f) {
			continue
		}
		stats.TotalAttempts += int64(a.Count)
		ips[a.IP] = struct{}{}
	}
	stats.UniqueIPs = int64(len(ips))

	for _, s := range m.Sessions {
		if f.Since != nil && s.ConnectedAt.Before(*f.Since) {
			continue
		}
		if f.Until != nil && s.ConnectedAt.After(*f.Until) {
			continue
		}
		if f.IP != "" && s.IP != f.IP {
			continue
		}
		if f.Country != "" && s.Country != f.Country {
			continue
		}
		stats.TotalSessions++
		if s.DisconnectedAt == nil {
			stats.ActiveSessions++
		}
	}

	return stats, nil
}

func (m *MemoryStore) GetFilteredTopUsernames(_ context.Context, limit int, f DashboardFilter) ([]TopEntry, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.filteredTopN("username", limit, f), nil
}

func (m *MemoryStore) GetFilteredTopPasswords(_ context.Context, limit int, f DashboardFilter) ([]TopEntry, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.filteredTopN("password", limit, f), nil
}

func (m *MemoryStore) GetFilteredTopIPs(_ context.Context, limit int, f DashboardFilter) ([]TopEntry, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	type ipInfo struct {
		count   int64
		country string
	}
	agg := make(map[string]*ipInfo)
	for i := range m.LoginAttempts {
		a := &m.LoginAttempts[i]
		if !matchesFilter(a, f) {
			continue
		}
		info, ok := agg[a.IP]
		if !ok {
			info = &ipInfo{}
			agg[a.IP] = info
		}
		info.count += int64(a.Count)
		if a.Country != "" {
			info.country = a.Country
		}
	}

	entries := make([]TopEntry, 0, len(agg))
	for ip, info := range agg {
		entries = append(entries, TopEntry{Value: ip, Country: info.country, Count: info.count})
	}
	sort.Slice(entries, func(i, j int) bool {
		return entries[i].Count > entries[j].Count
	})
	if limit > 0 && len(entries) > limit {
		entries = entries[:limit]
	}
	return entries, nil
}

func (m *MemoryStore) GetFilteredTopCountries(_ context.Context, limit int, f DashboardFilter) ([]TopEntry, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	counts := make(map[string]int64)
	for i := range m.LoginAttempts {
		a := &m.LoginAttempts[i]
		if a.Country == "" {
			continue
		}
		if !matchesFilter(a, f) {
			continue
		}
		counts[a.Country] += int64(a.Count)
	}

	entries := make([]TopEntry, 0, len(counts))
	for k, v := range counts {
		entries = append(entries, TopEntry{Value: k, Count: v})
	}
	sort.Slice(entries, func(i, j int) bool {
		return entries[i].Count > entries[j].Count
	})
	if limit > 0 && len(entries) > limit {
		entries = entries[:limit]
	}
	return entries, nil
}

// filteredTopN aggregates login attempts by the given field with filter applied and returns the top N. Must be called with m.mu held.
func (m *MemoryStore) filteredTopN(field string, limit int, f DashboardFilter) []TopEntry {
	counts := make(map[string]int64)
	for i := range m.LoginAttempts {
		a := &m.LoginAttempts[i]
		if !matchesFilter(a, f) {
			continue
		}
		var key string
		switch field {
		case "username":
			key = a.Username
		case "password":
			key = a.Password
		case "ip":
			key = a.IP
		}
		counts[key] += int64(a.Count)
	}

	entries := make([]TopEntry, 0, len(counts))
	for k, v := range counts {
		entries = append(entries, TopEntry{Value: k, Count: v})
	}
	sort.Slice(entries, func(i, j int) bool {
		return entries[i].Count > entries[j].Count
	})
	if limit > 0 && len(entries) > limit {
		entries = entries[:limit]
	}
	return entries
}

func (m *MemoryStore) Close() error {
	return nil
}