refactor: added ton of sqlite helpers and used sqlite for all the whitelists

2026-05-04 08:29:36 +03:30
parent 536d701641
commit d82ab73336
1 changed files with 225 additions and 0 deletions
--- a/db.go
+++ b/db.go
@@ -0,0 +1,225 @@
+package main
+
+import (
+	"database/sql"
+	"log/slog"
+	"time"
+
+	_ "modernc.org/sqlite"
+)
+
+func InitDB(dbPath string) (*sql.DB, error) {
+	db, err := sql.Open("sqlite", dbPath)
+	if err != nil {
+		return nil, err
+	}
+
+	_, err = db.Exec(`PRAGMA journal_mode=WAL; PRAGMA busy_timeout=5000;`)
+	if err != nil {
+		return nil, err
+	}
+
+	schema := `
+	CREATE TABLE IF NOT EXISTS permanent_whitelist (
+		id INTEGER PRIMARY KEY AUTOINCREMENT,
+		entry TEXT NOT NULL UNIQUE,
+		created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
+	);
+
+	CREATE TABLE IF NOT EXISTS temp_whitelist (
+		id INTEGER PRIMARY KEY AUTOINCREMENT,
+		ip TEXT NOT NULL UNIQUE,
+		expires_at DATETIME NOT NULL,
+		reason TEXT,
+		created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
+	);
+	CREATE INDEX IF NOT EXISTS idx_temp_expires ON temp_whitelist(expires_at);
+
+	CREATE TABLE IF NOT EXISTS audit_log (
+		id INTEGER PRIMARY KEY AUTOINCREMENT,
+		timestamp DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
+		action TEXT NOT NULL,
+		ip TEXT,
+		cidr TEXT,
+		ttl_seconds INTEGER,
+		reason TEXT,
+		api_client_ip TEXT
+	);
+	`
+	_, err = db.Exec(schema)
+	if err != nil {
+		return nil, err
+	}
+
+	return db, nil
+}
+
+func AddPermanentEntry(db *sql.DB, entry, apiClientIP string) (bool, error) {
+	_, err := db.Exec(`INSERT INTO permanent_whitelist(entry) VALUES(?)`, entry)
+	if err != nil {
+		if isUniqueConstraintError(err) {
+			return false, nil
+		}
+		return false, err
+	}
+	_, _ = db.Exec(`INSERT INTO audit_log(action, cidr, api_client_ip) VALUES('add_perm', ?, ?)`, entry, apiClientIP)
+	return true, nil
+}
+
+func DeletePermanentEntry(db *sql.DB, entry, apiClientIP string) error {
+	res, err := db.Exec(`DELETE FROM permanent_whitelist WHERE entry = ?`, entry)
+	if err != nil {
+		return err
+	}
+	rows, _ := res.RowsAffected()
+	if rows > 0 {
+		_, _ = db.Exec(`INSERT INTO audit_log(action, cidr, api_client_ip) VALUES('delete_perm', ?, ?)`, entry, apiClientIP)
+	}
+	return nil
+}
+
+func IsPermanentWhitelisted(db *sql.DB, ip string) (bool, error) {
+	rows, err := db.Query(`SELECT entry FROM permanent_whitelist`)
+	if err != nil {
+		return false, err
+	}
+	defer rows.Close()
+
+	for rows.Next() {
+		var entry string
+		if err := rows.Scan(&entry); err != nil {
+			continue
+		}
+		matches, err := ipMatchesEntry(ip, entry)
+		if err != nil {
+			continue
+		}
+		if matches {
+			return true, nil
+		}
+	}
+	return false, nil
+}
+
+func AddTempEntry(db *sql.DB, ip string, ttlSeconds int, reason, apiClientIP string) error {
+	expiresAt := time.Now().Add(time.Duration(ttlSeconds) * time.Second)
+	_, err := db.Exec(`
+		INSERT INTO temp_whitelist(ip, expires_at, reason)
+		VALUES(?, ?, ?)
+		ON CONFLICT(ip) DO UPDATE SET expires_at=excluded.expires_at, reason=excluded.reason
+	`, ip, expiresAt, reason)
+	if err != nil {
+		return err
+	}
+	_, _ = db.Exec(`INSERT INTO audit_log(action, ip, ttl_seconds, reason, api_client_ip) VALUES('add_temp', ?, ?, ?, ?)`,
+		ip, ttlSeconds, reason, apiClientIP)
+	return nil
+}
+
+func DeleteTempEntry(db *sql.DB, ip, apiClientIP string) error {
+	res, err := db.Exec(`DELETE FROM temp_whitelist WHERE ip = ?`, ip)
+	if err != nil {
+		return err
+	}
+	rows, _ := res.RowsAffected()
+	if rows > 0 {
+		_, _ = db.Exec(`INSERT INTO audit_log(action, ip, api_client_ip) VALUES('delete_temp', ?, ?)`, ip, apiClientIP)
+	}
+	return nil
+}
+
+func IsTempWhitelisted(db *sql.DB, ip string) (bool, error) {
+	var expiresAt time.Time
+	err := db.QueryRow(`SELECT expires_at FROM temp_whitelist WHERE ip = ?`, ip).Scan(&expiresAt)
+	if err == sql.ErrNoRows {
+		return false, nil
+	}
+	if err != nil {
+		return false, err
+	}
+	return expiresAt.After(time.Now()), nil
+}
+
+func CleanupExpiredTemp(db *sql.DB) error {
+	rows, err := db.Query(`SELECT ip FROM temp_whitelist WHERE expires_at <= ?`, time.Now())
+	if err != nil {
+		return err
+	}
+	var expiredIPs []string
+	for rows.Next() {
+		var ip string
+		if err := rows.Scan(&ip); err == nil {
+			expiredIPs = append(expiredIPs, ip)
+		}
+	}
+	rows.Close()
+
+	_, err = db.Exec(`DELETE FROM temp_whitelist WHERE expires_at <= ?`, time.Now())
+	if err != nil {
+		return err
+	}
+
+	for _, ip := range expiredIPs {
+		_, _ = db.Exec(`INSERT INTO audit_log(action, ip) VALUES('expire_temp', ?)`, ip)
+		slog.Info("expired temporary whitelist", "ip", ip)
+	}
+	return nil
+}
+
+func ListTempEntries(db *sql.DB) ([]map[string]interface{}, error) {
+	rows, err := db.Query(`SELECT ip, expires_at, reason FROM temp_whitelist WHERE expires_at > ? ORDER BY expires_at`, time.Now())
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	var list []map[string]interface{}
+	for rows.Next() {
+		var ip, reason string
+		var expiresAt time.Time
+		if err := rows.Scan(&ip, &expiresAt, &reason); err != nil {
+			continue
+		}
+		list = append(list, map[string]interface{}{
+			"ip":       ip,
+			"expires":  expiresAt,
+			"reason":   reason,
+		})
+	}
+	return list, nil
+}
+
+func ipMatchesEntry(ip, entry string) (bool, error) {
+	// Simple case: exact IP match
+	if entry == ip {
+		return true, nil
+	}
+
+	if len(entry) > 0 && entry[0] != '/' && containsSlash(entry) {
+		_, ipnet, err := net.ParseCIDR(entry)
+		if err != nil {
+			return false, err
+		}
+		parsedIP := net.ParseIP(ip)
+		if parsedIP == nil {
+			return false, nil
+		}
+		return ipnet.Contains(parsedIP), nil
+	}
+	return false, nil
+}
+
+func containsSlash(s string) bool {
+	for i := 0; i < len(s); i++ {
+		if s[i] == '/' {
+			return true
+		}
+	}
+	return false
+}
+
+func isUniqueConstraintError(err error) bool {
+	if err == nil {
+		return false
+	}
+	return err.Error() == "UNIQUE constraint failed: permanent_whitelist.entry"
+}