package exploit

import (
	"context"
	"testing"

	"github.com/google/uuid"

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

// TestAllowlist_NotConfiguredPreservesBehaviour confirms that an Executor
// constructed without WithAllowedExecutables behaves exactly as before
// (no blocking). Important so the existing test suite — and any callers
// that don't opt in — keep working.
func TestAllowlist_NotConfiguredPreservesBehaviour(t *testing.T) {
	e := NewExecutor(&scope.ScopeDefinition{AllowedDomains: []string{"example.com"}}, nil, true)
	step := pipeline.AttackStep{ID: uuid.New(), Name: "t", Command: "anything-goes-here https://example.com"}
	res, err := e.Execute(context.Background(), step, uuid.New())
	if err != nil {
		t.Fatalf("no-allowlist exec failed: %v", err)
	}
	if res == nil || !res.Success {
		t.Fatalf("dry-run with no allowlist should succeed; got %+v", res)
	}
}

// TestAllowlist_BlocksShellWrapper confirms the gate works for the
// canonical prompt-injection attack: model emits `bash -lc "..."` to
// bridge from a tool invocation into arbitrary shell execution.
func TestAllowlist_BlocksShellWrapper(t *testing.T) {
	e := NewExecutor(&scope.ScopeDefinition{AllowedDomains: []string{"example.com"}}, nil, true).
		WithAllowedExecutables([]string{"curl", "nuclei", "sqlmap"})

	step := pipeline.AttackStep{ID: uuid.New(), Name: "t", Command: "bash"}
	_, err := e.Execute(context.Background(), step, uuid.New())
	if err == nil {
		t.Fatal("expected bash to be blocked by allowlist")
	}
}

// TestAllowlist_AllowsRegisteredTool confirms a registered tool passes
// through. Dry-run keeps the test hermetic (no subprocess fired).
func TestAllowlist_AllowsRegisteredTool(t *testing.T) {
	e := NewExecutor(&scope.ScopeDefinition{AllowedDomains: []string{"example.com"}}, nil, true).
		WithAllowedExecutables([]string{"curl", "nuclei"})

	step := pipeline.AttackStep{ID: uuid.New(), Name: "t", Command: "nuclei -u https://example.com"}
	res, err := e.Execute(context.Background(), step, uuid.New())
	if err != nil {
		t.Fatalf("allowed tool should pass; got %v", err)
	}
	if res == nil || !res.Success {
		t.Fatalf("dry-run of allowed tool should succeed; got %+v", res)
	}
}

// TestAllowlist_BlocksCleanupBeforeRegister confirms the gate runs
// against the cleanup command too, *before* it's registered. A
// prompt-injected cleanup containing `rm -rf` would otherwise sit in
// the registry waiting to fire at campaign end.
func TestAllowlist_BlocksCleanupBeforeRegister(t *testing.T) {
	e := NewExecutor(&scope.ScopeDefinition{AllowedDomains: []string{"example.com"}}, nil, true).
		WithAllowedExecutables([]string{"curl"})

	step := pipeline.AttackStep{
		ID:             uuid.New(),
		Name:           "t",
		Command:        "curl https://example.com",
		CleanupCommand: "rm -rf /tmp/owned",
	}
	_, err := e.Execute(context.Background(), step, uuid.New())
	if err == nil {
		t.Fatal("expected dangerous cleanup command to be blocked")
	}
}

// TestAllowlist_PathPrefixIgnored confirms the allowlist matches on
// basename, so /usr/local/bin/nuclei is treated the same as nuclei.
// Important because LLM-authored commands sometimes include full paths.
func TestAllowlist_PathPrefixIgnored(t *testing.T) {
	e := NewExecutor(&scope.ScopeDefinition{AllowedDomains: []string{"example.com"}}, nil, true).
		WithAllowedExecutables([]string{"nuclei"})

	step := pipeline.AttackStep{ID: uuid.New(), Name: "t", Command: "/usr/local/bin/nuclei -u https://example.com"}
	res, err := e.Execute(context.Background(), step, uuid.New())
	if err != nil {
		t.Fatalf("path-prefixed allowed tool should pass; got %v", err)
	}
	if res == nil || !res.Success {
		t.Fatalf("expected success; got %+v", res)
	}
}

// TestAllowlist_CaseInsensitive — allowlist matching is case-insensitive
// to forgive minor inconsistencies (e.g. "Nuclei" vs "nuclei"). Useful
// for hand-written tests and unusual platforms.
func TestAllowlist_CaseInsensitive(t *testing.T) {
	e := NewExecutor(&scope.ScopeDefinition{AllowedDomains: []string{"example.com"}}, nil, true).
		WithAllowedExecutables([]string{"NUCLEI"})

	step := pipeline.AttackStep{ID: uuid.New(), Name: "t", Command: "nuclei -u https://example.com"}
	if _, err := e.Execute(context.Background(), step, uuid.New()); err != nil {
		t.Fatalf("case-insensitive match failed: %v", err)
	}
}