package agents

import (
	"context"
	"strings"
	"testing"

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

func TestVerifyBOLA_StampsWhenProven(t *testing.T) {
	a := &ExploitAgent{}
	var gotSpec *verify.ProbeSpec
	a.SetBOLAVerifier(func(_ context.Context, target string, spec *verify.ProbeSpec) (bool, string) {
		gotSpec = spec
		return true, "attack 200 vs control 403"
	})

	cf := pipeline.ClassifiedFinding{Title: "BOLA"}
	a.verifyBOLA(context.Background(), &cf, "https://t/api/vehicle/42/loc", "Bearer tok")

	if !cf.Verified || cf.VerifiedBy == "" {
		t.Fatalf("expected VERIFIED stamp, got %+v", cf)
	}
	if gotSpec == nil || gotSpec.AttackPayload != "42" {
		t.Errorf("verifier got spec %+v, want attack payload 42", gotSpec)
	}
}

func TestVerifyBOLA_NoStampWhenUnproven(t *testing.T) {
	a := &ExploitAgent{}
	a.SetBOLAVerifier(func(context.Context, string, *verify.ProbeSpec) (bool, string) { return false, "" })
	cf := pipeline.ClassifiedFinding{Title: "BOLA"}
	a.verifyBOLA(context.Background(), &cf, "https://t/api/vehicle/42/loc", "")
	if cf.Verified {
		t.Error("unproven verdict must not stamp VERIFIED")
	}
}

func TestVerifyBOLA_NilVerifierNoop(t *testing.T) {
	a := &ExploitAgent{} // no verifier
	cf := pipeline.ClassifiedFinding{Title: "BOLA"}
	a.verifyBOLA(context.Background(), &cf, "https://t/api/vehicle/42/loc", "")
	if cf.Verified {
		t.Error("no verifier must be a no-op")
	}
}

func TestVerifyBOLA_UntemplatableURLNoop(t *testing.T) {
	a := &ExploitAgent{}
	called := false
	a.SetBOLAVerifier(func(context.Context, string, *verify.ProbeSpec) (bool, string) { called = true; return true, "x" })
	cf := pipeline.ClassifiedFinding{Title: "BOLA"}
	a.verifyBOLA(context.Background(), &cf, "https://t/api/users", "") // no id segment
	if cf.Verified || called {
		t.Error("an un-templatable URL must not invoke the verifier or stamp")
	}
}

func maPath(cmd string) pipeline.AttackPath {
	return pipeline.AttackPath{Name: "Mass assignment", Steps: []pipeline.AttackStep{{Name: "escalate", Command: cmd}}}
}

func TestVerifyMassAssignment_StampsWhenProven(t *testing.T) {
	a := &ExploitAgent{}
	var gotSpec *verify.ProbeSpec
	var gotTarget string
	a.SetMassAssignVerifier(func(_ context.Context, target string, spec *verify.ProbeSpec) (bool, string) {
		gotSpec, gotTarget = spec, target
		return true, "indicator \"admin\" present in attack (200), absent in control (200)"
	})

	cf := pipeline.ClassifiedFinding{Title: "Mass assignment", Target: "https://api.t/users"}
	a.verifyMassAssignment(context.Background(), &cf,
		maPath(`httpreq --method POST --url https://api.t/users --body '{"name":"x","role":"admin"}'`), nil)

	if !cf.Verified || cf.VerifiedBy == "" {
		t.Fatalf("expected VERIFIED stamp, got %+v", cf)
	}
	if gotSpec == nil || gotSpec.Class != verify.ProbeDifferential {
		t.Errorf("verifier got spec %+v, want a differential spec", gotSpec)
	}
	if gotTarget != "https://api.t/users" {
		t.Errorf("oracle target = %q, want the step's own url", gotTarget)
	}
}

func TestVerifyMassAssignment_NoStampWhenUnproven(t *testing.T) {
	a := &ExploitAgent{}
	a.SetMassAssignVerifier(func(context.Context, string, *verify.ProbeSpec) (bool, string) { return false, "" })
	cf := pipeline.ClassifiedFinding{Title: "Mass assignment"}
	a.verifyMassAssignment(context.Background(), &cf,
		maPath(`httpreq --method POST --url https://api.t/users --body '{"role":"admin"}'`), nil)
	if cf.Verified {
		t.Error("unproven verdict must not stamp VERIFIED")
	}
}

func TestVerifyMassAssignment_NilVerifierNoop(t *testing.T) {
	a := &ExploitAgent{} // no verifier
	cf := pipeline.ClassifiedFinding{Title: "Mass assignment"}
	a.verifyMassAssignment(context.Background(), &cf,
		maPath(`httpreq --method POST --url https://api.t/users --body '{"role":"admin"}'`), nil)
	if cf.Verified {
		t.Error("no verifier must be a no-op")
	}
}

func TestVerifyMassAssignment_NonMassAssignChainNoop(t *testing.T) {
	a := &ExploitAgent{}
	called := false
	a.SetMassAssignVerifier(func(context.Context, string, *verify.ProbeSpec) (bool, string) { called = true; return true, "x" })
	cf := pipeline.ClassifiedFinding{Title: "BOLA"}
	// A GET chain carries no mass-assignment body — the spec never builds.
	a.verifyMassAssignment(context.Background(), &cf,
		maPath(`httpreq --method GET --url https://api.t/vehicle/42`), nil)
	if cf.Verified || called {
		t.Error("a non-mass-assignment chain must not invoke the verifier or stamp")
	}
}

func TestVerifySSRF_StampsWhenProven(t *testing.T) {
	a := &ExploitAgent{}
	var gotSpec *verify.ProbeSpec
	var gotTarget string
	a.SetSSRFVerifier(func(_ context.Context, target string, spec *verify.ProbeSpec) (bool, string) {
		gotSpec, gotTarget = spec, target
		return true, "http interaction from 10.0.0.5"
	})
	cf := pipeline.ClassifiedFinding{Title: "SSRF"}
	a.verifySSRF(context.Background(), &cf, "https://t/fetch?url=x", "url", "Bearer tok")
	if !cf.Verified || cf.VerifiedBy == "" {
		t.Fatalf("expected VERIFIED stamp, got %+v", cf)
	}
	if gotSpec == nil || gotSpec.Class != verify.ProbeOOB {
		t.Errorf("verifier got spec %+v, want an oob spec", gotSpec)
	}
	if gotTarget != "https://t/fetch?url=x" {
		t.Errorf("oracle target = %q, want the endpoint url", gotTarget)
	}
}

func TestVerifySSRF_NoStampWhenNoCallback(t *testing.T) {
	a := &ExploitAgent{}
	a.SetSSRFVerifier(func(context.Context, string, *verify.ProbeSpec) (bool, string) { return false, "" })
	cf := pipeline.ClassifiedFinding{Title: "SSRF"}
	a.verifySSRF(context.Background(), &cf, "https://t/fetch?url=x", "url", "")
	if cf.Verified {
		t.Error("no callback must not stamp VERIFIED (OOB never disproves, but silence is not proof)")
	}
}

func TestVerifySSRF_NilVerifierNoop(t *testing.T) {
	a := &ExploitAgent{}
	cf := pipeline.ClassifiedFinding{Title: "SSRF"}
	a.verifySSRF(context.Background(), &cf, "https://t/fetch?url=x", "url", "")
	if cf.Verified {
		t.Error("no verifier must be a no-op")
	}
}

func TestVerifySSRF_EmptyParamNoop(t *testing.T) {
	a := &ExploitAgent{}
	called := false
	a.SetSSRFVerifier(func(context.Context, string, *verify.ProbeSpec) (bool, string) { called = true; return true, "x" })
	cf := pipeline.ClassifiedFinding{Title: "SSRF"}
	a.verifySSRF(context.Background(), &cf, "https://t/fetch?url=x", "", "")
	if cf.Verified || called {
		t.Error("an empty sink param must not invoke the verifier or stamp")
	}
}

func TestVerifyAuthBypass_StampsWithTechnique(t *testing.T) {
	a := &ExploitAgent{}
	// Prove only the trailing-slash variant; others are denied.
	a.SetAuthBypassVerifier(func(_ context.Context, _ string, spec *verify.ProbeSpec) (bool, string) {
		if spec.AttackPayload == "/admin/" {
			return true, "attack 200 vs control 403 — authorization boundary crossed"
		}
		return false, ""
	})
	cf := pipeline.ClassifiedFinding{Title: "403 bypass"}
	tech, ok := a.verifyAuthBypass(context.Background(), &cf, "https://t/admin", "")
	if !ok || tech != "trailing-slash" {
		t.Fatalf("expected trailing-slash proven, got tech=%q ok=%v", tech, ok)
	}
	if !cf.Verified || !strings.Contains(cf.VerifiedBy, "trailing-slash") {
		t.Errorf("finding should record the winning technique: %+v", cf)
	}
}

func TestVerifyAuthBypass_NoStampWhenAllDenied(t *testing.T) {
	a := &ExploitAgent{}
	a.SetAuthBypassVerifier(func(context.Context, string, *verify.ProbeSpec) (bool, string) { return false, "" })
	cf := pipeline.ClassifiedFinding{Title: "403 bypass"}
	if _, ok := a.verifyAuthBypass(context.Background(), &cf, "https://t/admin", ""); ok || cf.Verified {
		t.Error("no working variant must not stamp VERIFIED")
	}
}

func TestVerifyAuthBypass_NilVerifierNoop(t *testing.T) {
	a := &ExploitAgent{}
	cf := pipeline.ClassifiedFinding{Title: "403 bypass"}
	if _, ok := a.verifyAuthBypass(context.Background(), &cf, "https://t/admin", ""); ok || cf.Verified {
		t.Error("no verifier must be a no-op")
	}
}

func TestRequestTarget(t *testing.T) {
	cases := map[string]string{
		"https://t/admin":  "/admin",
		"http://t/a/b?x=1": "/a/b?x=1",
		"https://t":        "/",
		"https://t/":       "/",
	}
	for in, want := range cases {
		if got := requestTarget(in); got != want {
			t.Errorf("requestTarget(%q) = %q, want %q", in, got, want)
		}
	}
}

func TestVerifyAuthBypass_CatchesHeaderTrickWhenPathsFail(t *testing.T) {
	a := &ExploitAgent{}
	a.SetAuthBypassVerifier(func(_ context.Context, _ string, spec *verify.ProbeSpec) (bool, string) {
		// Path variants all fail; only the X-Original-URL header trick works.
		if spec.AttackPayload == "/admin" && strings.Contains(spec.BaseRequest, "X-Original-URL") {
			return true, "attack 200 vs control 404"
		}
		return false, ""
	})
	cf := pipeline.ClassifiedFinding{Title: "403 bypass"}
	tech, ok := a.verifyAuthBypass(context.Background(), &cf, "https://t/admin", "")
	if !ok || tech != "header:X-Original-URL" {
		t.Fatalf("expected the header trick proven, got tech=%q ok=%v", tech, ok)
	}
	if !strings.Contains(cf.VerifiedBy, "X-Original-URL") {
		t.Errorf("finding should record the header technique: %+v", cf)
	}
}

func TestVerifySSTI_GatesOnPolyglotThenPinsEngine(t *testing.T) {
	a := &ExploitAgent{}
	var probes int
	a.SetSSTIVerifier(func(_ context.Context, _ string, spec *verify.ProbeSpec) (bool, string) {
		probes++
		// Polyglot gate proves; among per-engine specs, only Jinja proves.
		if strings.Contains(spec.AttackPayload, "%7B%7B") && strings.Contains(spec.AttackPayload, "%3C%25") {
			return true, "polyglot" // the combined payload (has both {{ and <% encoded)
		}
		if strings.HasPrefix(spec.AttackPayload, "%7B%7B") { // encoded "{{"
			return true, "indicator 36863 present in attack"
		}
		return false, ""
	})
	cf := pipeline.ClassifiedFinding{Title: "SSTI"}
	tech, ok := a.verifySSTI(context.Background(), &cf, "https://t/g?name=x", "name", "")
	if !ok || tech != "ssti:Jinja2/Twig/Nunjucks" {
		t.Fatalf("expected Jinja engine pinned, got tech=%q ok=%v", tech, ok)
	}
	if !cf.Verified {
		t.Error("finding should be VERIFIED")
	}
}

func TestVerifySSTI_NonTemplatingNoop(t *testing.T) {
	a := &ExploitAgent{}
	probes := 0
	a.SetSSTIVerifier(func(context.Context, string, *verify.ProbeSpec) (bool, string) {
		probes++
		return false, "" // polyglot gate fails
	})
	cf := pipeline.ClassifiedFinding{Title: "SSTI"}
	if _, ok := a.verifySSTI(context.Background(), &cf, "https://t/g?name=x", "name", ""); ok || cf.Verified {
		t.Error("a non-templating reflection must be a no-op")
	}
	if probes != 1 {
		t.Errorf("gate should cost exactly one polyglot probe, got %d", probes)
	}
}

func TestVerifySSTI_NilVerifierNoop(t *testing.T) {
	a := &ExploitAgent{}
	cf := pipeline.ClassifiedFinding{Title: "SSTI"}
	if _, ok := a.verifySSTI(context.Background(), &cf, "https://t/g?name=x", "name", ""); ok || cf.Verified {
		t.Error("no verifier must be a no-op")
	}
}

func TestVerifyCmdInjection_StampsWhenProven(t *testing.T) {
	a := &ExploitAgent{}
	var gotSpec *verify.ProbeSpec
	a.SetCmdInjectionVerifier(func(_ context.Context, _ string, spec *verify.ProbeSpec) (bool, string) {
		gotSpec = spec
		return true, "dns interaction from 10.0.0.9"
	})
	cf := pipeline.ClassifiedFinding{Title: "cmd injection"}
	a.verifyCmdInjection(context.Background(), &cf, "https://t/ping?host=x", "host", "")
	if !cf.Verified || cf.VerifiedBy == "" {
		t.Fatalf("expected VERIFIED stamp, got %+v", cf)
	}
	if gotSpec == nil || gotSpec.Class != verify.ProbeOOB {
		t.Errorf("verifier got spec %+v, want an oob spec", gotSpec)
	}
}

func TestVerifyCmdInjection_NoCallbackNoop(t *testing.T) {
	a := &ExploitAgent{}
	a.SetCmdInjectionVerifier(func(context.Context, string, *verify.ProbeSpec) (bool, string) { return false, "" })
	cf := pipeline.ClassifiedFinding{Title: "cmd injection"}
	a.verifyCmdInjection(context.Background(), &cf, "https://t/ping?host=x", "host", "")
	if cf.Verified {
		t.Error("no callback must not stamp VERIFIED")
	}
}

func TestVerifyCmdInjection_NilVerifierNoop(t *testing.T) {
	a := &ExploitAgent{}
	cf := pipeline.ClassifiedFinding{Title: "cmd injection"}
	a.verifyCmdInjection(context.Background(), &cf, "https://t/ping?host=x", "host", "")
	if cf.Verified {
		t.Error("no verifier must be a no-op")
	}
}