diff --git a/db.go b/db.go new file mode 100644 index 0000000..7e8ced1 --- /dev/null +++ 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" +} \ No newline at end of file