package exploit

import (
	"context"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/google/uuid"

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

// loopbackScope allows the 127.0.0.1 address httptest servers bind to.
func loopbackScope() *scope.ScopeDefinition {
	return &scope.ScopeDefinition{AllowedCIDRs: []string{"127.0.0.1/32"}}
}

func TestHTTPReq_GETReturnsStatusAndBody(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte(`{"role":"admin"}`))
	}))
	defer srv.Close()

	e := NewExecutor(loopbackScope(), nil, false)
	step := pipeline.AttackStep{ID: uuid.New(), Name: "get", Command: "httpreq --method GET --url " + srv.URL + "/api/user"}

	res, err := e.Execute(context.Background(), step, uuid.New())
	if err != nil {
		t.Fatalf("httpreq GET errored: %v", err)
	}
	if !res.Success {
		t.Fatalf("expected success; got %+v", res)
	}
	if !strings.Contains(res.Output, "HTTP 200") || !strings.Contains(res.Output, `"role":"admin"`) {
		t.Fatalf("output missing status/body: %q", res.Output)
	}
	if !strings.Contains(res.Output, "Content-Type: application/json") {
		t.Errorf("output missing content-type header: %q", res.Output)
	}
}

func TestHTTPReq_SendsMethodHeaderAndBody(t *testing.T) {
	var gotMethod, gotAuth, gotBody string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotMethod = r.Method
		gotAuth = r.Header.Get("Authorization")
		buf := make([]byte, r.ContentLength)
		_, _ = r.Body.Read(buf)
		gotBody = string(buf)
		w.WriteHeader(http.StatusCreated)
	}))
	defer srv.Close()

	e := NewExecutor(loopbackScope(), nil, false)
	step := pipeline.AttackStep{
		ID:   uuid.New(),
		Name: "post",
		// header + body are single-quoted so the shell-safe parser keeps each as one token.
		Command: "httpreq --method POST --url " + srv.URL + `/api/order --header 'Authorization: Bearer tok123' --body '{"qty":1,"is_admin":true}'`,
	}

	res, err := e.Execute(context.Background(), step, uuid.New())
	if err != nil {
		t.Fatalf("httpreq POST errored: %v", err)
	}
	if gotMethod != "POST" {
		t.Errorf("method = %q, want POST", gotMethod)
	}
	if gotAuth != "Bearer tok123" {
		t.Errorf("auth header = %q, want 'Bearer tok123'", gotAuth)
	}
	if !strings.Contains(gotBody, `"is_admin":true`) {
		t.Errorf("body = %q, want mass-assignment field", gotBody)
	}
	if !strings.Contains(res.Output, "HTTP 201") {
		t.Errorf("output missing 201: %q", res.Output)
	}
}

// The whole point of httpreq: it works even when the binary allowlist is
// configured and does NOT list it, because it is an in-process built-in.
func TestHTTPReq_BypassesBinaryAllowlist(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	}))
	defer srv.Close()

	e := NewExecutor(loopbackScope(), nil, false).
		WithAllowedExecutables([]string{"nuclei", "sqlmap"}) // deliberately no httpreq
	step := pipeline.AttackStep{ID: uuid.New(), Name: "get", Command: "httpreq --url " + srv.URL}

	res, err := e.Execute(context.Background(), step, uuid.New())
	if err != nil {
		t.Fatalf("built-in should bypass allowlist; got %v", err)
	}
	if !strings.Contains(res.Output, "HTTP 200") {
		t.Fatalf("expected 200; got %q", res.Output)
	}

	// And AllowedToolNames advertises it to the planner.
	found := false
	for _, n := range e.AllowedToolNames() {
		if n == builtinHTTPReq {
			found = true
		}
	}
	if !found {
		t.Errorf("AllowedToolNames should include %q", builtinHTTPReq)
	}
}

func TestHTTPReq_OutOfScopeBlocked(t *testing.T) {
	// Scope allows only example.com; a loopback URL is out of scope.
	e := NewExecutor(&scope.ScopeDefinition{AllowedDomains: []string{"example.com"}}, nil, false)
	step := pipeline.AttackStep{ID: uuid.New(), Name: "get", Command: "httpreq --url http://192.0.2.5/secret"}

	res, err := e.Execute(context.Background(), step, uuid.New())
	if err == nil {
		t.Fatal("expected out-of-scope block")
	}
	if res == nil || res.Success {
		t.Fatalf("blocked step must be unsuccessful; got %+v", res)
	}
	if !strings.Contains(res.Output, "BLOCKED") {
		t.Errorf("output should mark BLOCKED: %q", res.Output)
	}
}

func TestHTTPReq_SafeModeBlocksDestructiveMethod(t *testing.T) {
	e := NewExecutor(loopbackScope(), nil, false).WithSafeMode(true)
	step := pipeline.AttackStep{ID: uuid.New(), Name: "del", Command: "httpreq --method DELETE --url http://127.0.0.1:9/x"}

	_, err := e.Execute(context.Background(), step, uuid.New())
	if err == nil {
		t.Fatal("safe-mode should block DELETE method")
	}
}

func TestHTTPReq_NoFollowRedirect(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, "/elsewhere", http.StatusFound)
	}))
	defer srv.Close()

	e := NewExecutor(loopbackScope(), nil, false)
	step := pipeline.AttackStep{ID: uuid.New(), Name: "redir", Command: "httpreq --url " + srv.URL + "/start"}

	res, err := e.Execute(context.Background(), step, uuid.New())
	if err != nil {
		t.Fatalf("errored: %v", err)
	}
	if !strings.Contains(res.Output, "HTTP 302") || !strings.Contains(res.Output, "Location: /elsewhere") {
		t.Fatalf("expected unfollowed 302 with Location; got %q", res.Output)
	}
}

func TestParseHTTPReqArgs_UnknownFlagErrors(t *testing.T) {
	if _, err := parseHTTPReqArgs([]string{"--bogus", "x"}); err == nil {
		t.Fatal("expected error on unknown flag")
	}
	if _, err := parseHTTPReqArgs([]string{"--url"}); err == nil {
		t.Fatal("expected error on flag missing value")
	}
}

// The marquee test: a stateful chain. Step 1 logs in and captures the JWT
// from the JSON response; step 2 replays it in an Authorization header to reach
// a protected resource — exactly the BOLA / broken-auth pattern.
func TestExecuteChain_CaptureTokenThenReplay(t *testing.T) {
	const secretData = `{"vehicle":"victim-car","vin":"1HGBH41JXMN109186"}`
	var replayedAuth string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/identity/api/auth/login":
			w.Header().Set("Content-Type", "application/json")
			_, _ = w.Write([]byte(`{"token":"jwt-abc-123","type":"Bearer"}`))
		case "/identity/api/v2/vehicle/42/location":
			replayedAuth = r.Header.Get("Authorization")
			if replayedAuth != "Bearer jwt-abc-123" {
				w.WriteHeader(http.StatusUnauthorized)
				return
			}
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write([]byte(secretData))
		default:
			w.WriteHeader(http.StatusNotFound)
		}
	}))
	defer srv.Close()

	e := NewExecutor(loopbackScope(), nil, false)
	steps := []pipeline.AttackStep{
		{
			ID:      uuid.New(),
			Name:    "login",
			Command: "httpreq --method POST --url " + srv.URL + `/identity/api/auth/login --body '{"email":"a@b.c","password":"x"}' --capture tok=$.token`,
		},
		{
			ID:      uuid.New(),
			Name:    "bola",
			Command: "httpreq --url " + srv.URL + `/identity/api/v2/vehicle/42/location --header 'Authorization: Bearer {{tok}}'`,
		},
	}

	results, err := e.ExecuteChain(context.Background(), steps, uuid.New())
	if err != nil {
		t.Fatalf("chain errored: %v", err)
	}
	if len(results) != 2 {
		t.Fatalf("want 2 results, got %d", len(results))
	}
	if replayedAuth != "Bearer jwt-abc-123" {
		t.Fatalf("captured token not replayed; server saw auth=%q", replayedAuth)
	}
	if !strings.Contains(results[1].Output, "HTTP 200") || !strings.Contains(results[1].Output, "victim-car") {
		t.Fatalf("replay step did not reach protected data: %q", results[1].Output)
	}
}

func TestHTTPReq_DefaultsJSONContentTypeForBody(t *testing.T) {
	var gotCT string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotCT = r.Header.Get("Content-Type")
		w.WriteHeader(http.StatusOK)
	}))
	defer srv.Close()

	e := NewExecutor(loopbackScope(), nil, false)
	step := pipeline.AttackStep{
		ID:      uuid.New(),
		Command: "httpreq --method POST --url " + srv.URL + `/login --body '{"email":"a@b.c"}'`,
	}
	if _, err := e.Execute(context.Background(), step, uuid.New()); err != nil {
		t.Fatalf("execute: %v", err)
	}
	if gotCT != "application/json" {
		t.Fatalf("expected default Content-Type application/json for a bodied request, got %q", gotCT)
	}
}

func TestHTTPReq_ExplicitContentTypeWins(t *testing.T) {
	var gotCT string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotCT = r.Header.Get("Content-Type")
		w.WriteHeader(http.StatusOK)
	}))
	defer srv.Close()

	e := NewExecutor(loopbackScope(), nil, false)
	step := pipeline.AttackStep{
		ID:      uuid.New(),
		Command: "httpreq --method POST --url " + srv.URL + `/x --header 'Content-Type: application/xml' --body '<a/>'`,
	}
	if _, err := e.Execute(context.Background(), step, uuid.New()); err != nil {
		t.Fatalf("execute: %v", err)
	}
	if gotCT != "application/xml" {
		t.Fatalf("explicit Content-Type should win, got %q", gotCT)
	}
}

func TestExecuteChain_SeedsNonceIntoBody(t *testing.T) {
	var gotBody string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		b, _ := io.ReadAll(r.Body)
		gotBody = string(b)
		w.WriteHeader(http.StatusOK)
	}))
	defer srv.Close()

	e := NewExecutor(loopbackScope(), nil, false)
	steps := []pipeline.AttackStep{{
		ID:      uuid.New(),
		Name:    "signup",
		Command: "httpreq --method POST --url " + srv.URL + `/signup --body '{"email":"atk_{{nonce}}@example.com"}'`,
	}}

	if _, err := e.ExecuteChain(context.Background(), steps, uuid.New()); err != nil {
		t.Fatalf("chain errored: %v", err)
	}
	if strings.Contains(gotBody, "{{nonce}}") {
		t.Fatalf("nonce placeholder was not substituted: %q", gotBody)
	}
	if !strings.Contains(gotBody, "atk_") || !strings.Contains(gotBody, "@example.com") {
		t.Fatalf("body did not carry the substituted email: %q", gotBody)
	}
}

func TestSubstituteVars(t *testing.T) {
	vars := map[string]string{"tok": "abc123"}
	got := substituteVars(`--header "Authorization: Bearer {{tok}}"`, vars)
	if !strings.Contains(got, "Bearer abc123") {
		t.Errorf("substitution failed: %q", got)
	}
	// Unknown placeholders are left intact, not blanked.
	if out := substituteVars("{{missing}}", vars); out != "{{missing}}" {
		t.Errorf("unknown placeholder should be left as-is, got %q", out)
	}
}

func TestCaptureValue_JSONPathAndHeader(t *testing.T) {
	resp := &http.Response{Header: http.Header{}}
	resp.Header.Set("Location", "/next")
	body := []byte(`{"data":{"items":[{"id":"first"},{"id":"second"}]},"n":7,"ok":true}`)

	cases := []struct {
		sel, want string
		ok        bool
	}{
		{"$.data.items.0.id", "first", true},
		{"data.items.1.id", "second", true},
		{"$.n", "7", true},
		{"$.ok", "true", true},
		{"header:Location", "/next", true},
		{"$.missing", "", false},
		{"$.data.items.9.id", "", false},
	}
	for _, c := range cases {
		got, ok := captureValue(c.sel, resp, body)
		if ok != c.ok || (ok && got != c.want) {
			t.Errorf("captureValue(%q) = %q,%v; want %q,%v", c.sel, got, ok, c.want, c.ok)
		}
	}
}

func TestSplitHeader(t *testing.T) {
	k, v, ok := splitHeader("Authorization: Bearer abc")
	if !ok || k != "Authorization" || v != "Bearer abc" {
		t.Fatalf("splitHeader = %q/%q/%v", k, v, ok)
	}
	if _, _, ok := splitHeader("no-colon"); ok {
		t.Error("expected failure on header without colon")
	}
}