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 TestOpenRedirectSpec_Shape(t *testing.T) {
	spec, err := OpenRedirectSpec("https://t/go?next=/home", "next", "")
	if err != nil {
		t.Fatal(err)
	}
	if spec.AttackPayload != openRedirectSentinel {
		t.Errorf("attack payload = %q, want the sentinel", spec.AttackPayload)
	}
	if !strings.Contains(spec.BaseRequest, "next="+verify.DefaultMarker) {
		t.Errorf("param not marked: %q", spec.BaseRequest)
	}
}

// End to end: a handler that 302s to whatever ?next= says is proven open redirect;
// one that ignores it (relative redirect) is not.
func TestOpenRedirectSpec_ProvesViaRedirectOracle(t *testing.T) {
	vuln := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, r.URL.Query().Get("next"), http.StatusFound) // reflects attacker URL
	}))
	defer vuln.Close()
	safe := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, "/dashboard", http.StatusFound) // always same-site
	}))
	defer safe.Close()

	noRedirect := oracle.DefaultClientNoRedirect()

	spec, _ := OpenRedirectSpec(vuln.URL+"/go?next=x", "next", "")
	v, _, proof, err := oracle.NewRedirect(oracle.LiveRedirectProbe(noRedirect)).
		Probe(context.Background(), verify.Candidate{Target: vuln.URL, Spec: spec})
	if err != nil || v != verify.Proven {
		t.Fatalf("vuln: verdict=%v err=%v proof=%q, want proven", v, err, proof)
	}

	spec2, _ := OpenRedirectSpec(safe.URL+"/go?next=x", "next", "")
	v2, _, _, err := oracle.NewRedirect(oracle.LiveRedirectProbe(noRedirect)).
		Probe(context.Background(), verify.Candidate{Target: safe.URL, Spec: spec2})
	if err != nil {
		t.Fatal(err)
	}
	if v2 == verify.Proven {
		t.Error("safe same-site redirect must not be proven an open redirect")
	}
}