auth-proxy/auth.go

package main

import (
	"database/sql"
	"encoding/json"
	"log/slog"
	"net/http"
	"sort"
	"sync"
	"time"

	_ "modernc.org/sqlite"
)

type tempEntry struct {
	IP      string    `json:"ip"`
	Reason  string    `json:"reason,omitempty"`
	Expires time.Time `json:"expires"`
}

type tempWhitelist struct {
	mu      sync.RWMutex
	entries map[string]tempEntry
}

func newTempWhitelist() *tempWhitelist {
	return &tempWhitelist{entries: make(map[string]tempEntry)}
}

type PermanentWhitelist struct {
	mu      sync.RWMutex
	entries map[string]struct{}
}

func newPermanentWhitelist() *PermanentWhitelist {
	return &PermanentWhitelist{entries: make(map[string]struct{})}
}

func (tw *tempWhitelist) CheckAuth(r *http.Request, perm *PermanentWhitelist) bool {
	ip := r.Header.Get("X-Real-IP")
	if ip == "" {
		return false
	}
	perm.mu.RLock()
	_, ok := perm.entries[ip]
	perm.mu.RUnlock()
	if ok {
		return true
	}
	tw.mu.RLock()
	if e, exists := tw.entries[ip]; exists && e.Expires.After(time.Now()) {
		tw.mu.RUnlock()
		return true
	}
	tw.mu.RUnlock()
	return false
}

func authHandler(cfg Config, tw *tempWhitelist, perm *PermanentWhitelist) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if !tw.CheckAuth(r, perm) {
			w.WriteHeader(http.StatusUnauthorized)
			return
		}
		w.WriteHeader(http.StatusNoContent)
	}
}

func (tw *tempWhitelist) whitelistTempHandler(cfg Config, db *sql.DB) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if !verifyAPIKey(r, cfg.APIToken) {
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		if r.Method != http.MethodPost {
			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
			return
		}
		type request struct {
			IP         string `json:"ip"`
			TTLSeconds int    `json:"ttl_seconds"`
			Reason     string `json:"reason,omitempty"`
		}
		var req request
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			http.Error(w, "bad request", http.StatusBadRequest)
			return
		}
		if req.IP == "" || req.TTLSeconds <= 0 {
			http.Error(w, "ip and ttl_seconds required", http.StatusBadRequest)
			return
		}
		expires := time.Now().Add(time.Duration(req.TTLSeconds) * time.Second)

		tw.mu.Lock()
		tw.entries[req.IP] = tempEntry{IP: req.IP, Reason: req.Reason, Expires: expires}
		tw.mu.Unlock()

		go logWhitelistEvent(db, "add_temp", req.IP, req.TTLSeconds, req.Reason)

		slog.Info("whitelist added", "ip", req.IP, "ttl", req.TTLSeconds, "reason", req.Reason)
		w.WriteHeader(http.StatusNoContent)
	}
}

func (tw *tempWhitelist) whitelistListHandler(cfg Config) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if !verifyAPIKey(r, cfg.APIToken) {
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		tw.mu.RLock()
		entries := make([]tempEntry, 0, len(tw.entries))
		for _, e := range tw.entries {
			entries = append(entries, e)
		}
		tw.mu.RUnlock()
		sort.Slice(entries, func(i, j int) bool {
			return entries[i].Expires.Before(entries[j].Expires)
		})
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(entries)
	}
}

func (tw *tempWhitelist) whitelistDeleteHandler(cfg Config, db *sql.DB) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if !verifyAPIKey(r, cfg.APIToken) {
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		ip := r.PathValue("ip")
		if ip == "" {
			http.Error(w, "ip required", http.StatusBadRequest)
			return
		}
		tw.mu.Lock()
		delete(tw.entries, ip)
		tw.mu.Unlock()

		go logWhitelistEvent(db, "delete_temp", ip, 0, "", getAPIUser(r))
		slog.Info("whitelist deleted", "ip", ip)
		w.WriteHeader(http.StatusNoContent)
	}
}

func logsHandler(cfg Config, db *sql.DB) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if !verifyAPIKey(r, cfg.APIToken) {
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		rows, err := db.Query(`SELECT timestamp, action, ip, ttl_seconds, reason, api_user 
			FROM whitelist_audit ORDER BY timestamp DESC LIMIT 100`)
		if err != nil {
			http.Error(w, "database error", http.StatusInternalServerError)
			return
		}
		defer rows.Close()
		var logs []map[string]interface{}
		for rows.Next() {
			var ts time.Time
			var action, ip, reason, apiUser string
			var ttl int
			if err := rows.Scan(&ts, &action, &ip, &ttl, &reason); err != nil {
				continue
			}
			logs = append(logs, map[string]interface{}{
				"timestamp": ts,
				"action":    action,
				"ip":        ip,
				"ttl":       ttl,
				"reason":    reason,
			})
		}
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(logs)
	}
}

func statusHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain")
	w.Write([]byte("ok\n"))
}

func verifyAPIKey(r *http.Request, expectedToken string) bool {
	return r.Header.Get("Authorization") == "Bearer "+expectedToken
}