warpbox/lib/security/guard.go

package security

import (
	"sort"
	"strings"
	"sync"
	"time"
)

type Config struct {
	IPWhitelist         string
	AdminIPWhitelist    string
	LoginWindowSeconds  int64
	LoginMaxAttempts    int
	BanSeconds          int64
	ScanWindowSeconds   int64
	ScanMaxAttempts     int
	UploadWindowSeconds int64
	UploadMaxRequests   int
	UploadMaxBytes      int64
}

type Guard struct {
	mu             sync.Mutex
	failedLogins   map[string][]time.Time
	scanAttempts   map[string][]time.Time
	uploadEvents   map[string][]uploadEvent
	bannedUntil    map[string]time.Time
	ipWhitelist    map[string]bool
	adminWhitelist map[string]bool
}

type uploadEvent struct {
	at    time.Time
	bytes int64
}

type BanEntry struct {
	IP    string    `json:"ip"`
	Until time.Time `json:"until"`
}

func NewGuard() *Guard {
	return &Guard{
		failedLogins:   map[string][]time.Time{},
		scanAttempts:   map[string][]time.Time{},
		uploadEvents:   map[string][]uploadEvent{},
		bannedUntil:    map[string]time.Time{},
		ipWhitelist:    map[string]bool{},
		adminWhitelist: map[string]bool{},
	}
}

func (g *Guard) Reload(cfg Config) {
	g.mu.Lock()
	defer g.mu.Unlock()
	g.ipWhitelist = parseList(cfg.IPWhitelist)
	g.adminWhitelist = parseList(cfg.AdminIPWhitelist)
}

func (g *Guard) IsWhitelisted(ip string) bool {
	g.mu.Lock()
	defer g.mu.Unlock()
	return g.ipWhitelist[ip]
}

func (g *Guard) IsAdminWhitelisted(ip string) bool {
	g.mu.Lock()
	defer g.mu.Unlock()
	return g.adminWhitelist[ip] || g.ipWhitelist[ip]
}

func (g *Guard) IsBanned(ip string) bool {
	g.mu.Lock()
	defer g.mu.Unlock()
	until, ok := g.bannedUntil[ip]
	if !ok {
		return false
	}
	if time.Now().UTC().After(until) {
		delete(g.bannedUntil, ip)
		return false
	}
	return true
}

func (g *Guard) Ban(ip string, seconds int64) {
	if seconds <= 0 || ip == "" {
		return
	}
	g.mu.Lock()
	defer g.mu.Unlock()
	g.bannedUntil[ip] = time.Now().UTC().Add(time.Duration(seconds) * time.Second)
}

func (g *Guard) BanUntil(ip string, until time.Time) {
	if ip == "" || until.IsZero() {
		return
	}
	g.mu.Lock()
	defer g.mu.Unlock()
	g.bannedUntil[ip] = until.UTC()
}

func (g *Guard) Unban(ip string) {
	if ip == "" {
		return
	}
	g.mu.Lock()
	defer g.mu.Unlock()
	delete(g.bannedUntil, ip)
}

func (g *Guard) BanList() []BanEntry {
	g.mu.Lock()
	defer g.mu.Unlock()
	now := time.Now().UTC()
	out := make([]BanEntry, 0, len(g.bannedUntil))
	for ip, until := range g.bannedUntil {
		if now.After(until) {
			delete(g.bannedUntil, ip)
			continue
		}
		out = append(out, BanEntry{IP: ip, Until: until})
	}
	sort.Slice(out, func(i, j int) bool {
		return out[i].Until.Before(out[j].Until)
	})
	return out
}

func (g *Guard) RegisterFailedLogin(ip string, windowSeconds int64, maxAttempts int, banSeconds int64) (bool, int) {
	if ip == "" || maxAttempts <= 0 || windowSeconds <= 0 {
		return false, 0
	}
	g.mu.Lock()
	defer g.mu.Unlock()
	now := time.Now().UTC()
	cutoff := now.Add(-time.Duration(windowSeconds) * time.Second)
	attempts := pruneTimes(g.failedLogins[ip], cutoff)
	attempts = append(attempts, now)
	g.failedLogins[ip] = attempts
	if len(attempts) >= maxAttempts {
		g.bannedUntil[ip] = now.Add(time.Duration(banSeconds) * time.Second)
		return true, len(attempts)
	}
	return false, len(attempts)
}

func (g *Guard) RegisterScanAttempt(ip string, windowSeconds int64, maxAttempts int, banSeconds int64) (bool, int) {
	if ip == "" || maxAttempts <= 0 || windowSeconds <= 0 {
		return false, 0
	}
	g.mu.Lock()
	defer g.mu.Unlock()
	now := time.Now().UTC()
	cutoff := now.Add(-time.Duration(windowSeconds) * time.Second)
	attempts := pruneTimes(g.scanAttempts[ip], cutoff)
	attempts = append(attempts, now)
	g.scanAttempts[ip] = attempts
	if len(attempts) >= maxAttempts {
		g.bannedUntil[ip] = now.Add(time.Duration(banSeconds) * time.Second)
		return true, len(attempts)
	}
	return false, len(attempts)
}

func (g *Guard) AllowUpload(ip string, size int64, windowSeconds int64, maxRequests int, maxBytes int64) (bool, int, int64) {
	if ip == "" || windowSeconds <= 0 || maxRequests <= 0 {
		return true, 0, 0
	}
	g.mu.Lock()
	defer g.mu.Unlock()
	now := time.Now().UTC()
	cutoff := now.Add(-time.Duration(windowSeconds) * time.Second)
	events := g.uploadEvents[ip]
	kept := make([]uploadEvent, 0, len(events)+1)
	totalBytes := int64(0)
	for _, event := range events {
		if event.at.After(cutoff) {
			kept = append(kept, event)
			totalBytes += event.bytes
		}
	}
	nextCount := len(kept) + 1
	nextBytes := totalBytes + size
	if nextCount > maxRequests {
		return false, nextCount, nextBytes
	}
	if maxBytes > 0 && nextBytes > maxBytes {
		return false, nextCount, nextBytes
	}
	kept = append(kept, uploadEvent{at: now, bytes: size})
	g.uploadEvents[ip] = kept
	return true, nextCount, nextBytes
}

func pruneTimes(values []time.Time, cutoff time.Time) []time.Time {
	kept := make([]time.Time, 0, len(values))
	for _, value := range values {
		if value.After(cutoff) {
			kept = append(kept, value)
		}
	}
	return kept
}

func parseList(raw string) map[string]bool {
	out := map[string]bool{}
	for _, chunk := range strings.Split(raw, ",") {
		value := strings.TrimSpace(chunk)
		if value != "" {
			out[value] = true
		}
	}
	return out
}