zoraxy-waf/main.go

package main

import (
	"embed"
	"encoding/json"
	"fmt"
	"log"
	"net/http"

	"zoraxy-firewall/zoraxy_plugin"
)

//go:embed web/*
var webFS embed.FS

func main() {
	loadConfig()

	spec := &zoraxy_plugin.IntroSpect{
		ID: "zoraxy-firewall",
		Name: "Firewall",
		Author:        "Zoraxy Community",
		AuthorContact: "",
		Description: "Web Application Firewall with circuit breaker and fail-open protection.",
		URL:           "",
		Type:          zoraxy_plugin.PluginType_Router, // Type 0 — intercepts traffic
		VersionMajor:  1,
		VersionMinor:  0,
		VersionPatch:  0,
		UIPath:        "/ui",

		// Intercept ALL proxied requests.
		StaticCapturePaths:   []zoraxy_plugin.StaticCaptureRule{{CapturePath: "/"}},
		StaticCaptureIngress: "/inspect",
	}

	config, err := zoraxy_plugin.ServeAndRecvSpec(spec)
	if err != nil {
		log.Fatalf("failed to receive config: %v", err)
	}

	mux := http.NewServeMux()

	// --- Static capture handler (traffic interception) ---
	mux.HandleFunc("/inspect", handleInspect)

	// --- UI + config API ---
	uiRouter := zoraxy_plugin.NewPluginEmbedUIRouter(
		spec.ID, &webFS, "web", spec.UIPath,
	)
	uiRouter.HandleFunc("/api/config", handleGetConfig, mux)
	uiRouter.HandleFunc("/api/config/save", handleSaveConfig, mux)
	uiRouter.HandleFunc("/api/stats", handleGetStats, mux)

	uiRouter.RegisterTerminateHandler(func() {
		log.Println("zoraxy-waf shutting down")
	}, mux)
	uiRouter.AttachHandlerToMux(mux)

	// --- Return port listener for inspection endpoint async callbacks (future) ---
	cfg := getConfig()
	if cfg.ReturnPort > 0 {
		go func() {
			rmux := http.NewServeMux()
			rmux.HandleFunc("/verdict", handleReturnVerdict)
			addr := fmt.Sprintf(":%d", cfg.ReturnPort)
			log.Printf("WAF return listener on %s", addr)
			if err := http.ListenAndServe(addr, rmux); err != nil {
				log.Printf("WAF return listener error: %v", err)
			}
		}()
	}

	// --- Main plugin listener ---
	addr := fmt.Sprintf("127.0.0.1:%d", config.Port)
	log.Printf("zoraxy-waf v%d.%d.%d listening on %s",
		spec.VersionMajor, spec.VersionMinor, spec.VersionPatch, addr)

	if err := http.ListenAndServe(addr, mux); err != nil {
		log.Fatalf("server error: %v", err)
	}
}

// --- Config API handlers ---

func handleGetConfig(w http.ResponseWriter, r *http.Request) {
	c := getConfig()
	writeJSON(w, http.StatusOK, c)
}

func handleSaveConfig(w http.ResponseWriter, r *http.Request) {
	if err := r.ParseForm(); err != nil {
		writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid form"})
		return
	}

	cfgMu.Lock()
	cfg.Enabled = r.FormValue("enabled") == "true"
	if u := r.FormValue("node_url"); u != "" {
		cfg.NodeURL = u
	}
	if p := r.FormValue("return_port"); p != "" {
		fmt.Sscanf(p, "%d", &cfg.ReturnPort)
	}
	if t := r.FormValue("timeout_ms"); t != "" {
		fmt.Sscanf(t, "%d", &cfg.TimeoutMs)
	}
	cfgMu.Unlock()

	if err := saveConfig(); err != nil {
		log.Printf("ERROR saving config: %v", err)
		writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "save failed"})
		return
	}
	writeJSON(w, http.StatusOK, map[string]string{"status": "saved"})
}

func handleGetStats(w http.ResponseWriter, r *http.Request) {
	writeJSON(w, http.StatusOK, metricsSnapshot())
}

func handleReturnVerdict(w http.ResponseWriter, r *http.Request) {
	// Future: inspection endpoint async callbacks arrive here.
	var resp InspectionResponse
	if err := json.NewDecoder(r.Body).Decode(&resp); err != nil {
		writeJSON(w, http.StatusBadRequest, map[string]string{"error": "bad verdict"})
		return
	}
	log.Printf("WAF async verdict: %s (reason=%s)", resp.Verdict, resp.Reason)
	writeJSON(w, http.StatusOK, map[string]string{"status": "ok"})
}

// --- Helpers ---

func writeJSON(w http.ResponseWriter, status int, v interface{}) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	json.NewEncoder(w).Encode(v)
}