package exploit

import (
	"context"
	"encoding/json"
	"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 TestNoSQLAuthBypassSpec_Shape(t *testing.T) {
	spec, err := NoSQLAuthBypassSpec("https://t/api/user/login", "user", "pass", "")
	if err != nil {
		t.Fatal(err)
	}
	if spec.Class != verify.ProbeDifferential || !spec.Valid() {
		t.Fatalf("not a valid differential spec: %+v", spec)
	}
	if !strings.Contains(spec.AttackPayload, `"$gt"`) {
		t.Errorf("attack payload should carry an operator object: %q", spec.AttackPayload)
	}
	if strings.Contains(spec.ControlPayload, "$gt") {
		t.Errorf("control payload must be plain creds, not an operator: %q", spec.ControlPayload)
	}
	if !strings.HasPrefix(spec.BaseRequest, "POST /api/user/login HTTP/1.1") {
		t.Errorf("unexpected request line: %q", spec.BaseRequest)
	}
	if !strings.Contains(spec.BaseRequest, "Content-Type: application/json") {
		t.Error("should send a JSON content type")
	}
}

func TestNoSQLAuthBypassSpec_RejectsEmptyFields(t *testing.T) {
	if _, err := NoSQLAuthBypassSpec("https://t/login", "", "pass", ""); err == nil {
		t.Error("empty field name should be rejected")
	}
}

func TestNoSQLAuthBypassSpecs_CoversCommonPairs(t *testing.T) {
	specs, err := NoSQLAuthBypassSpecs("https://t/login", "")
	if err != nil {
		t.Fatal(err)
	}
	techs := map[string]bool{}
	for _, ts := range specs {
		techs[ts.Technique] = true
	}
	for _, want := range []string{"nosql:user/pass", "nosql:username/password", "nosql:email/password"} {
		if !techs[want] {
			t.Errorf("missing field-pair spec %q", want)
		}
	}
}

// End to end: a Mongo-style login that accepts an operator object as a credential
// (returns a token) but rejects wrong string creds is proven via the differential
// oracle.
func TestNoSQLAuthBypassSpec_ProvesViaDifferentialOracle(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var body map[string]any
		_ = json.NewDecoder(r.Body).Decode(&body)
		// Vulnerable: if "pass" is an object (operator), "match any" → logged in.
		if _, isObj := body["pass"].(map[string]any); isObj {
			w.WriteHeader(200)
			_, _ = w.Write([]byte(`{"token":"jwt-xyz"}`))
			return
		}
		w.WriteHeader(401)
	}))
	defer srv.Close()

	spec, err := NoSQLAuthBypassSpec(srv.URL+"/api/user/login", "user", "pass", "")
	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 login that ignores operator objects (treats them as a failed match) must not
// be flagged — the differential collapses.
func TestNoSQLAuthBypassSpec_DisprovesWhenNotInjectable(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(401) // always denies
	}))
	defer srv.Close()
	spec, err := NoSQLAuthBypassSpec(srv.URL+"/login", "user", "pass", "")
	if err != nil {
		t.Fatal(err)
	}
	v, _, _, err := oracle.NewDifferential(oracle.LiveDiffProbe(srv.Client())).
		Probe(context.Background(), verify.Candidate{Target: srv.URL, Spec: spec})
	if err != nil {
		t.Fatal(err)
	}
	if v == verify.Proven {
		t.Error("a non-injectable login must not be proven")
	}
}