package verify

import (
	"strings"
	"testing"
)

func TestProbeSpec_MarkerToken(t *testing.T) {
	if (&ProbeSpec{}).MarkerToken() != DefaultMarker {
		t.Error("empty marker should default")
	}
	if (&ProbeSpec{Marker: "X"}).MarkerToken() != "X" {
		t.Error("explicit marker should win")
	}
	var nilSpec *ProbeSpec
	if nilSpec.MarkerToken() != DefaultMarker {
		t.Error("nil spec should default")
	}
}

func TestProbeSpec_Valid(t *testing.T) {
	base := "GET /api?q=" + DefaultMarker + " HTTP/1.1\r\nHost: t\r\n\r\n"
	cases := []struct {
		name string
		spec *ProbeSpec
		want bool
	}{
		{"nil", nil, false},
		{"no base", &ProbeSpec{Class: ProbeDifferential, AttackPayload: "x"}, false},
		{"base missing marker", &ProbeSpec{Class: ProbeDifferential, BaseRequest: "GET / HTTP/1.1\r\n\r\n", AttackPayload: "x"}, false},
		{"differential ok", &ProbeSpec{Class: ProbeDifferential, BaseRequest: base, AttackPayload: "2"}, true},
		{"oob missing template marker", &ProbeSpec{Class: ProbeOOB, BaseRequest: base, OOBTemplate: "http://x"}, false},
		{"oob ok", &ProbeSpec{Class: ProbeOOB, BaseRequest: base, OOBTemplate: "' AND LOAD_FILE('//" + OOBMarker + "')"}, true},
		{"timing missing delay marker", &ProbeSpec{Class: ProbeTiming, BaseRequest: base, DelayTemplate: "SLEEP(5)"}, false},
		{"timing ok", &ProbeSpec{Class: ProbeTiming, BaseRequest: base, DelayTemplate: "' AND SLEEP(" + DelayMarker + ")-- -"}, true},
		{"unknown class", &ProbeSpec{Class: "nope", BaseRequest: base}, false},
	}
	for _, tc := range cases {
		if got := tc.spec.Valid(); got != tc.want {
			t.Errorf("%s: Valid = %v, want %v", tc.name, got, tc.want)
		}
	}
}

func TestProbeSpec_Inject(t *testing.T) {
	s := &ProbeSpec{BaseRequest: "GET /?q=" + DefaultMarker + "&n=" + DefaultMarker + " HTTP/1.1\r\n\r\n"}
	got := s.Inject("PWN")
	if strings.Count(got, "PWN") != 2 || strings.Contains(got, DefaultMarker) {
		t.Errorf("Inject should replace all markers: %q", got)
	}
}

func TestNewSpecConstructors(t *testing.T) {
	base := "GET /api?id=" + DefaultMarker + " HTTP/1.1\r\nHost: t\r\n\r\n"

	if _, err := NewDifferentialSpec(base, "", "2", "999", "victim"); err != nil {
		t.Errorf("valid differential spec rejected: %v", err)
	}
	if _, err := NewTimingSpec(base, "", "' AND SLEEP("+DelayMarker+")-- -"); err != nil {
		t.Errorf("valid timing spec rejected: %v", err)
	}
	if _, err := NewOOBSpec(base, "", "http://"+OOBMarker+"/"); err != nil {
		t.Errorf("valid oob spec rejected: %v", err)
	}

	// Invalid: base without the marker.
	if _, err := NewDifferentialSpec("GET / HTTP/1.1\r\n\r\n", "", "2", "999", ""); err == nil {
		t.Error("spec with no marker in base should be rejected")
	}
	// Invalid: timing template missing the delay marker.
	if _, err := NewTimingSpec(base, "", "SLEEP(5)"); err == nil {
		t.Error("timing template without §DELAY§ should be rejected")
	}
	// Invalid: oob template missing the OOB marker.
	if _, err := NewOOBSpec(base, "", "http://static/"); err == nil {
		t.Error("oob template without §OOB§ should be rejected")
	}
}

func TestNewSpec_CustomMarker(t *testing.T) {
	base := "GET /api?id=INJECT HTTP/1.1\r\nHost: t\r\n\r\n"
	s, err := NewDifferentialSpec(base, "INJECT", "2", "999", "")
	if err != nil {
		t.Fatal(err)
	}
	if s.MarkerToken() != "INJECT" || !strings.Contains(s.Inject("X"), "id=X") {
		t.Errorf("custom marker not honoured: %q", s.Inject("X"))
	}
}

func TestCandidate_CanProbe(t *testing.T) {
	base := "GET /api?id=" + DefaultMarker + " HTTP/1.1\r\nHost: t\r\n\r\n"
	spec := &ProbeSpec{Class: ProbeDifferential, BaseRequest: base, AttackPayload: "2", ControlPayload: "999999"}

	c := Candidate{Spec: spec}
	if !c.CanProbe(ProbeDifferential) {
		t.Error("should be probeable for its class")
	}
	if c.CanProbe(ProbeTiming) {
		t.Error("wrong class must not probe")
	}
	if (Candidate{}).CanProbe(ProbeDifferential) {
		t.Error("no spec must not probe")
	}
}