zoraxy-waf/cmd/mockwallarm/main.go

95 lines
2.1 KiB
Go

package main
import (
"encoding/json"
"fmt"
"log"
"net/http"
"os"
"time"
)
// MockWallarm is a test server that simulates a Wallarm node.
// Usage: go run mock_wallarm.go [--mode=allow|block|slow|flaky]
//
// Modes:
// allow — always return "allow" verdict
// block — always return "block" verdict
// slow — delay 5s then allow (tests timeout)
// flaky — succeed 3 times, then fail 3 times (tests circuit breaker)
type Verdict struct {
Verdict string `json:"verdict"`
Reason string `json:"reason,omitempty"`
Score int `json:"score,omitempty"`
}
var (
requestCount int
mode string
)
func main() {
mode = "allow"
if len(os.Args) > 1 {
mode = os.Args[1]
}
// Also support --mode=block style
if len(mode) > 7 && mode[:7] == "--mode=" {
mode = mode[7:]
}
http.HandleFunc("/inspect", handleInspect)
http.HandleFunc("/health", handleHealth)
port := "8085"
if p := os.Getenv("PORT"); p != "" {
port = p
}
log.Printf("Mock Wallarm node starting on :%s (mode=%s)", port, mode)
if err := http.ListenAndServe(":"+port, nil); err != nil {
log.Fatal(err)
}
}
func handleInspect(w http.ResponseWriter, r *http.Request) {
requestCount++
switch mode {
case "block":
writeVerdict(w, "block", "SQL injection detected", 95)
case "slow":
time.Sleep(5 * time.Second)
writeVerdict(w, "allow", "clean", 0)
case "flaky":
if requestCount <= 3 {
writeVerdict(w, "allow", "clean", 0)
} else {
// Simulate failure — return 500
http.Error(w, "internal error", 500)
}
if requestCount >= 6 {
requestCount = 0 // reset cycle
}
default: // allow
writeVerdict(w, "allow", fmt.Sprintf("clean (request #%d)", requestCount), 0)
}
}
func handleHealth(w http.ResponseWriter, r *http.Request) {
writeJSON(w, http.StatusOK, map[string]string{"status": "ok"})
}
func writeVerdict(w http.ResponseWriter, verdict, reason string, score int) {
writeJSON(w, http.StatusOK, Verdict{Verdict: verdict, Reason: reason, Score: score})
}
func writeJSON(w http.ResponseWriter, status int, v interface{}) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(status)
json.NewEncoder(w).Encode(v)
}