auth-proxy/db.go

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
	}
	_, err = db.Exec(`INSERT INTO audit_log(action, cidr, api_client_ip) VALUES('add_perm', ?, ?)`, entry, apiClientIP)
	if err != nil {
		return true, err
	}
	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 {
		_, err = db.Exec(`INSERT INTO audit_log(action, cidr, api_client_ip) VALUES('delete_perm', ?, ?)`, entry, apiClientIP)
		if err != nil {
			return err
		}
	}
	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
	}
	_, err = db.Exec(`INSERT INTO audit_log(action, ip, ttl_seconds, reason, api_client_ip) VALUES('add_temp', ?, ?, ?, ?)`,
		ip, ttlSeconds, reason, apiClientIP)
	if err != nil {
		return err
	}
	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 {
		_, err = db.Exec(`INSERT INTO audit_log(action, ip, api_client_ip) VALUES('delete_temp', ?, ?)`, ip, apiClientIP)
		if err != nil {
			return err
		}
	}
	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 {
		_, err = db.Exec(`INSERT INTO audit_log(action, ip) VALUES('expire_temp', ?)`, ip)
		if err != nil {
			slog.Warn("failed to log expired temp entry", "ip", ip, "error", err)
		}
		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")
}