refactor: added ton of sqlite helpers and used sqlite for all the whitelists
This commit is contained in:
225
db.go
Normal file
225
db.go
Normal file
@@ -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"
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user