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 TestMassAssignmentSpec_BuildsDifferentialRecipe(t *testing.T) {
	cmd := `httpreq --method POST --url https://api.t/users --header "Authorization: Bearer tok" --body '{"name":"x","role":"admin"}'`
	spec, err := MassAssignmentSpec(cmd, nil)
	if err != nil {
		t.Fatal(err)
	}
	if spec.Class != verify.ProbeDifferential || !spec.Valid() {
		t.Fatalf("not a valid differential spec: %+v", spec)
	}
	// Attack keeps the privileged field; control drops it.
	if !strings.Contains(spec.AttackPayload, `"role":"admin"`) {
		t.Errorf("attack payload should keep the privileged field: %q", spec.AttackPayload)
	}
	if strings.Contains(spec.ControlPayload, "role") {
		t.Errorf("control payload must drop the privileged field: %q", spec.ControlPayload)
	}
	if !strings.Contains(spec.ControlPayload, `"name":"x"`) {
		t.Errorf("control payload should keep benign fields: %q", spec.ControlPayload)
	}
	if spec.Indicator != "admin" {
		t.Errorf("distinctive string value should be the indicator, got %q", spec.Indicator)
	}
	if !strings.Contains(spec.BaseRequest, "Authorization: Bearer tok") {
		t.Error("base request should carry the session credential")
	}
	if !strings.HasPrefix(spec.BaseRequest, "POST /users HTTP/1.1") {
		t.Errorf("unexpected request line: %q", spec.BaseRequest)
	}
	if !strings.HasSuffix(spec.BaseRequest, "\r\n\r\n"+verify.DefaultMarker) {
		t.Errorf("body should be the injection marker: %q", spec.BaseRequest)
	}
}

func TestMassAssignmentSpec_ResolvesVars(t *testing.T) {
	cmd := `httpreq --method PUT --url https://api.t/profile --header "Authorization: Bearer {{tok}}" --body '{"email":"a@b.c","is_admin":true}'`
	spec, err := MassAssignmentSpec(cmd, map[string]string{"tok": "SECRET123"})
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(spec.BaseRequest, "Authorization: Bearer SECRET123") {
		t.Errorf("{{tok}} should be substituted: %q", spec.BaseRequest)
	}
	// Boolean value → no indicator (falls back to divergence).
	if spec.Indicator != "" {
		t.Errorf("boolean privileged value should not become an indicator, got %q", spec.Indicator)
	}
	if strings.Contains(spec.ControlPayload, "is_admin") {
		t.Errorf("control must drop is_admin: %q", spec.ControlPayload)
	}
}

func TestMassAssignmentSpec_Rejections(t *testing.T) {
	cases := []struct{ name, cmd string }{
		{"GET not body-mutating", `httpreq --method GET --url https://t/users`},
		{"no body", `httpreq --method POST --url https://t/users`},
		{"body not JSON object", `httpreq --method POST --url https://t/users --body 'not json'`},
		{"no privileged field", `httpreq --method POST --url https://t/users --body '{"name":"x","email":"a@b.c"}'`},
		{"not httpreq", `curl https://t/users`},
	}
	for _, c := range cases {
		if _, err := MassAssignmentSpec(c.cmd, nil); err == nil {
			t.Errorf("%s: expected rejection, got nil error", c.name)
		}
	}
}

// End to end: the produced recipe actually re-proves a mass-assignment against a
// live endpoint that reflects a role only when the client supplies it.
func TestMassAssignmentSpec_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: the server trusts a client-supplied role (mass assignment).
		resp := map[string]any{"id": 1, "name": body["name"]}
		if role, ok := body["role"].(string); ok {
			resp["role"] = role
		} else {
			resp["role"] = "user" // default — the control gets the safe default
		}
		w.WriteHeader(200)
		_ = json.NewEncoder(w).Encode(resp)
	}))
	defer srv.Close()

	cmd := `httpreq --method POST --url ` + srv.URL + `/users --body '{"name":"x","role":"admin"}'`
	spec, err := MassAssignmentSpec(cmd, nil)
	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)
	}
	if !strings.Contains(proof, "admin") {
		t.Errorf("proof should cite the reflected indicator: %q", proof)
	}
}

// A non-vulnerable endpoint that ignores the client-supplied role collapses the
// differential — the oracle must Disprove, not falsely confirm.
func TestMassAssignmentSpec_DisprovesWhenFieldIgnored(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)
		// Secure: role is server-assigned, never from the body.
		w.WriteHeader(200)
		_ = json.NewEncoder(w).Encode(map[string]any{"id": 1, "name": body["name"], "role": "user"})
	}))
	defer srv.Close()

	cmd := `httpreq --method POST --url ` + srv.URL + `/users --body '{"name":"x","role":"admin"}'`
	spec, err := MassAssignmentSpec(cmd, nil)
	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.Disproven {
		t.Fatalf("verdict=%v err=%v proof=%q, want disproven", v, err, proof)
	}
}