refactor: move Go source files to src/ directory
Organize project structure by relocating all Go files to a src/ directory, following standard Go project conventions and improving code organization.
This commit is contained in:
db123-test
412
src/db.go
Normal file
412
src/db.go
Normal file
@@ -0,0 +1,412 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"net"
|
||||
"time"
|
||||
|
||||
sqlite "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,
|
||||
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) {
|
||||
_, err = db.Exec(`
|
||||
INSERT INTO audit_log(action, ip, api_client_ip)
|
||||
VALUES('add_perm_duplicate', ?, ?)
|
||||
`, entry, apiClientIP)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return false, nil
|
||||
}
|
||||
return false, err
|
||||
}
|
||||
_, err = db.Exec(`
|
||||
INSERT INTO audit_log(action, ip, 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 {
|
||||
_, err := db.Exec(`
|
||||
DELETE FROM perm_whitelist
|
||||
WHERE entry = ?
|
||||
`, entry)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = db.Exec(`
|
||||
INSERT INTO audit_log(action, ip, 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 {
|
||||
_, err := db.Exec(`
|
||||
DELETE FROM temp_whitelist
|
||||
WHERE ip = ?
|
||||
`, ip)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, 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
|
||||
}
|
||||
|
||||
var reason string
|
||||
reason = "expired"
|
||||
|
||||
for _, ip := range expiredIPs {
|
||||
_, err = db.Exec(`
|
||||
INSERT INTO audit_log(action, ip, reason)
|
||||
VALUES('expire_temp', ?, ?)
|
||||
`, ip, reason)
|
||||
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 string
|
||||
var expiresAt time.Time
|
||||
var reason sql.NullString
|
||||
if err := rows.Scan(&ip, &expiresAt, &reason); err != nil {
|
||||
continue
|
||||
}
|
||||
list = append(list, map[string]interface{}{
|
||||
"ip": ip,
|
||||
"expires": expiresAt,
|
||||
"reason": reason.String,
|
||||
})
|
||||
}
|
||||
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 ListTempEntriesPaginated(db *sql.DB, limit, offset int) ([]map[string]interface{}, error) {
|
||||
rows, err := db.Query(`
|
||||
SELECT ip, expires_at, reason
|
||||
FROM temp_whitelist
|
||||
WHERE expires_at > ?
|
||||
ORDER BY expires_at
|
||||
LIMIT ? OFFSET ?
|
||||
`, time.Now(), limit, offset)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var list []map[string]interface{} = make([]map[string]interface{}, 0)
|
||||
for rows.Next() {
|
||||
var ip string
|
||||
var expiresAt time.Time
|
||||
var reason sql.NullString
|
||||
if err := rows.Scan(&ip, &expiresAt, &reason); err != nil {
|
||||
continue
|
||||
}
|
||||
list = append(list, map[string]interface{}{
|
||||
"ip": ip,
|
||||
"expires": expiresAt,
|
||||
"reason": reason.String,
|
||||
})
|
||||
}
|
||||
return list, nil
|
||||
}
|
||||
|
||||
func ListPermEntriesPaginated(db *sql.DB, limit, offset int) ([]map[string]interface{}, error) {
|
||||
rows, err := db.Query(`
|
||||
SELECT entry
|
||||
FROM perm_whitelist
|
||||
ORDER BY created_at
|
||||
LIMIT ? OFFSET ?
|
||||
`, limit, offset)
|
||||
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
|
||||
}
|
||||
var sqliteErr *sqlite.Error
|
||||
if errors.As(err, &sqliteErr) {
|
||||
return sqliteErr.Code() == 19
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func IsIP(s string) bool {
|
||||
return net.ParseIP(s) != nil
|
||||
}
|
||||
|
||||
func IsCIDR(s string) bool {
|
||||
_, _, err := net.ParseCIDR(s)
|
||||
return err == nil
|
||||
}
|
||||
|
||||
func RotateAuditLog(db *sql.DB, maxEntries int) error {
|
||||
if maxEntries <= 0 {
|
||||
maxEntries = 200000
|
||||
}
|
||||
var count int
|
||||
err := db.QueryRow(`SELECT COUNT(*) FROM audit_log`).Scan(&count)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if count <= maxEntries {
|
||||
return nil
|
||||
}
|
||||
rows, err := db.Query(`
|
||||
SELECT id FROM audit_log
|
||||
ORDER BY timestamp ASC
|
||||
LIMIT ?
|
||||
`, count-maxEntries)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var ids []int64
|
||||
for rows.Next() {
|
||||
var id int64
|
||||
if err := rows.Scan(&id); err == nil {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
}
|
||||
rows.Close()
|
||||
_, err = db.Exec(`
|
||||
DELETE FROM audit_log
|
||||
WHERE id IN (
|
||||
SELECT id FROM audit_log
|
||||
ORDER BY timestamp ASC
|
||||
LIMIT ?
|
||||
)
|
||||
`, count-maxEntries)
|
||||
return err
|
||||
}
|
||||
Reference in New Issue
Block a user