package main import ( "database/sql" "log/slog" "net" "strings" "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 perm_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 perm_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 perm_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 perm_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{} = make([]map[string]interface{}, 0) 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 ListPermEntries(db *sql.DB) ([]map[string]interface{}, error) { rows, err := db.Query(`SELECT entry FROM perm_whitelist ORDER BY created_at`) if err != nil { return nil, err } defer rows.Close() var list []map[string]interface{} = make([]map[string]interface{}, 0) for rows.Next() { var entry string if err := rows.Scan(&entry); err != nil { continue } list = append(list, map[string]interface{}{ "entry": entry, }) } return list, nil } func ipMatchesEntry(ip, entry string) (bool, error) { 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 strings.Contains(err.Error(), "UNIQUE constraint failed") }