auth-proxy/auth.go

package main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"log/slog"
	"net"
	"net/http"
	"strconv"
	"strings"
	"time"
)

var startTime = time.Now()

func CheckAuth(db *sql.DB, r *http.Request) bool {
	ip := r.Header.Get("X-Real-IP")
	if ip == "" {
		return false
	}
	ok, err := IsPermanentWhitelisted(db, ip)
	if err == nil && ok {
		return true
	}
	ok, err = IsTempWhitelisted(db, ip)
	if err == nil && ok {
		return true
	}
	return false
}

func authHandler(cfg Config, db *sql.DB) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		ip := r.Header.Get("X-Real-IP")
		if ip == "" {
			slog.Warn("auth denied: no X-Real-IP", "remote_addr", r.RemoteAddr)
			w.WriteHeader(http.StatusUnauthorized)
			return
		}
		if !CheckAuth(db, r) {
			slog.Warn("auth denied", "ip", ip)
			w.WriteHeader(http.StatusUnauthorized)
			return
		}
		slog.Info("auth allowed", "ip", ip)
		w.WriteHeader(http.StatusNoContent)
	}
}

func apiKeyMiddleware(next http.HandlerFunc, token string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if !verifyAPIKey(r, token) {
			slog.Warn("unauthorized API access")
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		next(w, r)
	}
}

func 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
		}
		if !IsIP(req.IP) {
			http.Error(w, "invalid IP address", http.StatusBadRequest)
			return
		}
		clientIP := getClientIP(r)
		if err := AddTempEntry(db, req.IP, req.TTLSeconds, req.Reason, clientIP); err != nil {
			slog.Error("failed to add temp entry", "error", err)
			http.Error(w, "internal error", http.StatusInternalServerError)
			return
		}
		slog.Info("temp whitelist added", "ip", req.IP, "ttl", req.TTLSeconds)
		w.WriteHeader(http.StatusNoContent)
	}
}

func whitelistListHandler(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
		}
		limit, offset, err := parsePagination(r)
		if err != nil {
			http.Error(w, "bad request", http.StatusBadRequest)
			return
		}
		entries, err := ListTempEntriesPaginated(db, limit, offset)
		if err != nil {
			http.Error(w, "database error", http.StatusInternalServerError)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(entries)
	}
}

func permWhitelistListHandler(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
		}
		limit, offset, err := parsePagination(r)
		if err != nil {
			http.Error(w, "bad request", http.StatusBadRequest)
			return
		}
		entries, err := ListPermEntriesPaginated(db, limit, offset)
		if err != nil {
			http.Error(w, "database error", http.StatusInternalServerError)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(entries)
	}
}

func 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
		}
		clientIP := getClientIP(r)
		if err := DeleteTempEntry(db, ip, clientIP); err != nil {
			slog.Error("failed to delete temp entry", "error", err)
			http.Error(w, "internal error", http.StatusInternalServerError)
			return
		}
		slog.Info("temp 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
		}

		query := `SELECT timestamp, action, ip, cidr, ttl_seconds, reason, api_client_ip FROM audit_log`
		var args []any

		filterAction := r.URL.Query().Get("action")
		filterIP := r.URL.Query().Get("ip")
		limitStr := r.URL.Query().Get("limit")
		offsetStr := r.URL.Query().Get("offset")

		var conditions []string
		if filterAction != "" {
			conditions = append(conditions, `action = ?`)
			args = append(args, filterAction)
		}
		if filterIP != "" {
			conditions = append(conditions, `ip = ?`)
			args = append(args, filterIP)
		}

		if len(conditions) > 0 {
			query += " WHERE " + strings.Join(conditions, " AND ")
		}
		query += ` ORDER BY timestamp DESC`

		limit := 200
		offset := 0
		if limitStr != "" {
			if parsed, err := strconv.Atoi(limitStr); err == nil && parsed > 0 && parsed <= 1000 {
				limit = parsed
			}
		}
		if offsetStr != "" {
			if parsed, err := strconv.Atoi(offsetStr); err == nil && parsed >= 0 {
				offset = parsed
			}
		}
		query += fmt.Sprintf(` LIMIT %d OFFSET %d`, limit, offset)

		rows, err := db.Query(query, args...)
		if err != nil {
			slog.Error("failed to query logs", "error", err)
			http.Error(w, "database error", http.StatusInternalServerError)
			return
		}
		defer rows.Close()

		logs := make([]map[string]interface{}, 0)
		for rows.Next() {
			var ts time.Time
			var action, apiClientIP string
			var ip, cidr, reason sql.NullString
			var ttlSeconds sql.NullInt64
			if err := rows.Scan(&ts, &action, &ip, &cidr, &ttlSeconds, &reason, &apiClientIP); err != nil {
				continue
			}
			entry := map[string]interface{}{
				"timestamp":     ts,
				"action":        action,
				"api_client_ip": apiClientIP,
				"reason":        reason.String,
			}
			if ip.Valid {
				entry["ip"] = ip.String
			}
			if cidr.Valid {
				entry["cidr"] = cidr.String
			}
			if ttlSeconds.Valid {
				entry["ttl_seconds"] = ttlSeconds.Int64
			}
			logs = append(logs, entry)
		}
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(logs)
	}
}

func permWhitelistAddHandler(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
		}
		type request struct {
			Entry string `json:"entry"`
		}
		var req request
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			http.Error(w, "bad request", http.StatusBadRequest)
			return
		}
		if req.Entry == "" {
			http.Error(w, "entry required", http.StatusBadRequest)
			return
		}
		if !IsIP(req.Entry) && !IsCIDR(req.Entry) {
			http.Error(w, "invalid entry (must be valid IP or CIDR range)", http.StatusBadRequest)
			return
		}
		ok, err := AddPermanentEntry(db, req.Entry, getClientIP(r))
		if err != nil {
			slog.Error("failed to add perm entry", "error", err)
			http.Error(w, "internal error", http.StatusInternalServerError)
			return
		}
		if !ok {
			http.Error(w, "entry already exists", http.StatusConflict)
			return
		}
		w.WriteHeader(http.StatusNoContent)
	}
}

func permWhitelistDeleteHandler(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
		}
		entry := r.PathValue("entry")
		if entry == "" {
			http.Error(w, "entry required", http.StatusBadRequest)
			return
		}
		if err := DeletePermanentEntry(db, entry, getClientIP(r)); err != nil {
			slog.Error("failed to delete perm entry", "error", err)
			http.Error(w, "internal error", http.StatusInternalServerError)
			return
		}
		w.WriteHeader(http.StatusNoContent)
	}
}

func getClientIP(r *http.Request) string {
	if ip := r.Header.Get("X-Real-IP"); ip != "" {
		return ip
	}
	ip, _, _ := net.SplitHostPort(r.RemoteAddr)
	return ip
}

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

func statusHandler(db *sql.DB) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		var permCount int
		var tempCount int
		var dbOK bool

		if err := db.QueryRow(`SELECT COUNT(*) FROM perm_whitelist`).Scan(&permCount); err == nil {
			dbOK = true
		}
		if err := db.QueryRow(`SELECT COUNT(*) FROM temp_whitelist WHERE expires_at > ?`, time.Now().Format(time.RFC3339)).Scan(&tempCount); err == nil {
			dbOK = true
		} else {
			dbOK = false
		}

		health := map[string]interface{}{
			"status":               "ok",
			"uptime":               time.Since(startTime).String(),
			"perm_whitelist_count": permCount,
			"temp_whitelist_count": tempCount,
			"db_ok":                dbOK,
		}

		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(health)
	}
}

func parsePagination(r *http.Request) (int, int, error) {
	limitStr := r.URL.Query().Get("limit")
	offsetStr := r.URL.Query().Get("offset")

	limit := 200
	offset := 0

	if limitStr != "" {
		if parsed, err := strconv.Atoi(limitStr); err == nil && parsed > 0 && parsed <= 1000 {
			limit = parsed
		}
	}
	if offsetStr != "" {
		if parsed, err := strconv.Atoi(offsetStr); err == nil && parsed >= 0 {
			offset = parsed
		}
	}
	return limit, offset, nil
}