package exploit

import (
	"context"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os/exec"
	"strings"
	"testing"

	"github.com/Armur-Ai/Pentest-Swarm-AI/internal/swarm/verify"
	"github.com/Armur-Ai/Pentest-Swarm-AI/internal/swarm/verify/oracle"
)

func TestCmdInjectionOOBSpec_Shape(t *testing.T) {
	spec, err := CmdInjectionOOBSpec("https://t/ping?host=1.1.1.1&count=2", "host", "Bearer tok")
	if err != nil {
		t.Fatal(err)
	}
	if spec.Class != verify.ProbeOOB || !spec.Valid() {
		t.Fatalf("not a valid oob spec: %+v", spec)
	}
	if !strings.Contains(spec.BaseRequest, "host="+verify.DefaultMarker) {
		t.Errorf("sink param not marked: %q", spec.BaseRequest)
	}
	if !strings.Contains(spec.BaseRequest, "count=2") {
		t.Errorf("other params should be preserved: %q", spec.BaseRequest)
	}
	if !strings.Contains(spec.BaseRequest, "Authorization: Bearer tok") {
		t.Error("credential should be carried")
	}
	// Shell metacharacters are query-encoded; the §OOB§ slot stays literal.
	if strings.ContainsAny(spec.OOBTemplate, "; `|") {
		t.Errorf("shell metachars not encoded in template: %q", spec.OOBTemplate)
	}
	if !strings.Contains(spec.OOBTemplate, verify.OOBMarker) {
		t.Errorf("OOB marker missing from template: %q", spec.OOBTemplate)
	}
}

func TestCmdInjectionOOBSpec_RejectsEmptyParam(t *testing.T) {
	if _, err := CmdInjectionOOBSpec("https://t/x?a=1", "", ""); err == nil {
		t.Error("empty parameter should be rejected")
	}
}

// End to end: a target vulnerable to command injection runs the decoded shell
// payload, which curls the OOB listener; the oracle observes the callback.
func TestCmdInjectionOOBSpec_ProvesViaOOBOracle(t *testing.T) {
	if _, err := exec.LookPath("sh"); err != nil {
		t.Skip("sh not available")
	}
	if _, err := exec.LookPath("curl"); err != nil {
		t.Skip("curl not available — the injected payload needs it to call back")
	}
	var hit bool
	oob := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hit = true
		w.WriteHeader(200)
	}))
	defer oob.Close()

	target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Vulnerable: the host param is interpolated into a shell command.
		raw := r.URL.Query().Get("host")
		decoded, _ := url.QueryUnescape(raw) // already decoded by r.URL.Query(); harmless
		_ = decoded
		// Emulate `ping <host>` with the raw (already-decoded) value in a shell.
		cmd := exec.Command("sh", "-c", "ping -c1 "+raw)
		_ = cmd.Run() // we don't care about ping; the injected ;curl …; is what matters
		w.WriteHeader(200)
	}))
	defer target.Close()

	spec, err := CmdInjectionOOBSpec(target.URL+"/ping?host=seed", "host", "")
	if err != nil {
		t.Fatal(err)
	}
	cand := verify.Candidate{Target: target.URL, Spec: spec}
	o := oracle.NewOOB(fakeOOBSession{url: oob.URL, hit: &hit, proof: "http from 127.0.0.1"},
		oracle.LiveDelivery(target.Client()))
	v, _, proof, err := o.Probe(context.Background(), cand)
	if err != nil || v != verify.Proven {
		t.Fatalf("verdict=%v err=%v proof=%q, want proven", v, err, proof)
	}
}
