package exploit

import (
	"context"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/Armur-Ai/Pentest-Swarm-AI/internal/swarm/verify"
	"github.com/Armur-Ai/Pentest-Swarm-AI/internal/swarm/verify/oracle"
)

func TestAuthBypassVariants(t *testing.T) {
	got := AuthBypassVariants("/admin")
	techniques := map[string]string{}
	for _, v := range got {
		techniques[v.Technique] = v.Path
	}
	for _, tech := range []string{"trailing-slash", "dot-segment", "double-slash", "path-parameter", "trailing-dot", "case-fold-segment"} {
		if _, ok := techniques[tech]; !ok {
			t.Errorf("missing variant %q (got %v)", tech, techniques)
		}
	}
	if techniques["trailing-slash"] != "/admin/" {
		t.Errorf("trailing-slash = %q", techniques["trailing-slash"])
	}
	if techniques["case-fold-segment"] != "/Admin" {
		t.Errorf("case-fold = %q, want /Admin", techniques["case-fold-segment"])
	}
	if techniques["dot-segment"] != "/./admin" {
		t.Errorf("dot-segment = %q", techniques["dot-segment"])
	}
}

func TestAuthBypassVariants_PreservesQuery(t *testing.T) {
	for _, v := range AuthBypassVariants("/admin?debug=1") {
		if !strings.HasSuffix(v.Path, "?debug=1") {
			t.Errorf("%s dropped the query: %q", v.Technique, v.Path)
		}
	}
}

func TestAuthBypassSpec_BuildsDifferentialRecipe(t *testing.T) {
	spec, err := AuthBypassSpec("https://t/admin", AuthBypassVariant{Technique: "trailing-slash", Path: "/admin/"}, "Bearer low")
	if err != nil {
		t.Fatal(err)
	}
	if spec.Class != verify.ProbeDifferential || !spec.Valid() {
		t.Fatalf("not a valid differential spec: %+v", spec)
	}
	if spec.AttackPayload != "/admin/" || spec.ControlPayload != "/admin" {
		t.Errorf("payloads = attack %q control %q", spec.AttackPayload, spec.ControlPayload)
	}
	if !strings.Contains(spec.BaseRequest, "GET "+verify.DefaultMarker+" HTTP/1.1") {
		t.Errorf("request-target not marked: %q", spec.BaseRequest)
	}
	if !strings.Contains(spec.BaseRequest, "Authorization: Bearer low") {
		t.Error("base request should carry the session credential")
	}
}

// End to end: a server that denies /admin but serves /admin/ (a classic
// front-end-ACL vs back-end-routing mismatch) is proven via the differential
// oracle — attack 200 vs control 403.
func TestAuthBypassSpec_ProvesViaDifferentialOracle(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/admin/" {
			w.WriteHeader(200)
			_, _ = w.Write([]byte(`{"secret":"flag"}`))
			return
		}
		w.WriteHeader(403)
	}))
	defer srv.Close()

	spec, err := AuthBypassSpec(srv.URL+"/admin", AuthBypassVariant{Technique: "trailing-slash", Path: "/admin/"}, "")
	if err != nil {
		t.Fatal(err)
	}
	cand := verify.Candidate{Target: srv.URL, Spec: spec}
	d := oracle.NewDifferential(oracle.LiveDiffProbe(srv.Client()))
	v, _, proof, err := d.Probe(context.Background(), cand)
	if err != nil || v != verify.Proven {
		t.Fatalf("verdict=%v err=%v proof=%q, want proven", v, err, proof)
	}
}

// A uniformly-protected endpoint (both the variant and the original are denied)
// must not be reported as a bypass.
func TestAuthBypassSpec_DisprovesWhenUniformlyProtected(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(403)
	}))
	defer srv.Close()

	spec, err := AuthBypassSpec(srv.URL+"/admin", AuthBypassVariant{Technique: "trailing-slash", Path: "/admin/"}, "")
	if err != nil {
		t.Fatal(err)
	}
	cand := verify.Candidate{Target: srv.URL, Spec: spec}
	d := oracle.NewDifferential(oracle.LiveDiffProbe(srv.Client()))
	v, _, _, err := d.Probe(context.Background(), cand)
	if err != nil {
		t.Fatal(err)
	}
	if v == verify.Proven {
		t.Error("a uniformly-protected endpoint must not be proven a bypass")
	}
}