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 TestBOLASpec_BuildsDifferentialRecipe(t *testing.T) {
	spec, err := BOLASpec("https://api.crapi.io/identity/api/v2/vehicle/{id}/location", "42", "Bearer tok")
	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 != "42" || spec.ControlPayload != "999999999" {
		t.Errorf("payloads = attack %q control %q", spec.AttackPayload, spec.ControlPayload)
	}
	if !strings.Contains(spec.BaseRequest, "Authorization: Bearer tok") {
		t.Error("base request should carry the session credential")
	}
	if !strings.Contains(spec.BaseRequest, "/vehicle/"+verify.DefaultMarker+"/location") {
		t.Errorf("id slot not marked: %q", spec.BaseRequest)
	}
}

func TestBOLASpec_UUIDControlShape(t *testing.T) {
	spec, err := BOLASpec("https://t/api/user/{id}", "3fa85f64-5717-4562-b3fc-2c963f66afa6", "")
	if err != nil {
		t.Fatal(err)
	}
	if spec.ControlPayload != "00000000-0000-0000-0000-000000000000" {
		t.Errorf("uuid control = %q, want zero-uuid", spec.ControlPayload)
	}
}

func TestBOLASpec_NoIDSlotStillValidIfMarkerAbsent(t *testing.T) {
	// A template with no {id} produces a base request without the marker, which
	// NewDifferentialSpec must reject.
	_, err := BOLASpec("https://t/api/users", "42", "")
	if err == nil {
		t.Error("template with no {id} slot should be rejected")
	}
}

// The produced recipe actually re-proves a BOLA end to end: a live endpoint
// authorizes victim id 42 but denies the control id.
func TestBOLASpec_ProvesViaDifferentialOracle(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.Contains(r.URL.Path, "/42/") {
			w.WriteHeader(200)
			_, _ = w.Write([]byte(`{"owner":"someone-else","vin":"X"}`))
			return
		}
		w.WriteHeader(403)
	}))
	defer srv.Close()

	spec, err := BOLASpec(srv.URL+"/vehicle/{id}/loc", "42", "")
	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)
	}
}

func TestBOLATemplateAndID(t *testing.T) {
	cases := []struct {
		url, wantTmpl, wantID string
		wantOK                bool
	}{
		{"https://t/api/vehicle/42/loc", "https://t/api/vehicle/{id}/loc", "42", true},
		{"https://t/api/user/3fa85f64-5717-4562-b3fc-2c963f66afa6", "https://t/api/user/{id}", "3fa85f64-5717-4562-b3fc-2c963f66afa6", true},
		{"https://t/api/users", "", "", false}, // no id-shaped segment
	}
	for _, c := range cases {
		tmpl, id, ok := BOLATemplateAndID(c.url)
		if ok != c.wantOK || tmpl != c.wantTmpl || id != c.wantID {
			t.Errorf("BOLATemplateAndID(%q) = %q,%q,%v; want %q,%q,%v", c.url, tmpl, id, ok, c.wantTmpl, c.wantID, c.wantOK)
		}
	}
}

// Round-trip: a hit URL → template+id → spec that re-proves it.
func TestBOLATemplateAndID_FeedsSpec(t *testing.T) {
	tmpl, id, ok := BOLATemplateAndID("https://api/v2/vehicle/77/location")
	if !ok {
		t.Fatal("expected a templatable url")
	}
	spec, err := BOLASpec(tmpl, id, "Bearer x")
	if err != nil {
		t.Fatal(err)
	}
	if spec.AttackPayload != "77" {
		t.Errorf("attack payload = %q, want 77", spec.AttackPayload)
	}
}