package exploit

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

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

func TestXXEOOBSpec_Shape(t *testing.T) {
	spec, err := XXEOOBSpec("https://t/api/import/xml", "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.HasPrefix(spec.BaseRequest, "POST /api/import/xml HTTP/1.1") {
		t.Errorf("unexpected request line: %q", spec.BaseRequest)
	}
	if !strings.Contains(spec.BaseRequest, "Content-Type: application/xml") {
		t.Error("should send an XML content type")
	}
	if !strings.Contains(spec.OOBTemplate, "<!ENTITY") || !strings.Contains(spec.OOBTemplate, verify.OOBMarker) {
		t.Errorf("payload should declare an external entity with the OOB marker: %q", spec.OOBTemplate)
	}
}

// End to end: a target whose XML parser resolves the external entity fetches the
// callback; the OOB oracle observes it and proves XXE.
func TestXXEOOBSpec_ProvesViaOOBOracle(t *testing.T) {
	var hit bool
	oob := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { hit = true; w.WriteHeader(200) }))
	defer oob.Close()

	sysRe := regexp.MustCompile(`SYSTEM\s+"([^"]+)"`)
	target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, _ := io.ReadAll(r.Body)
		if m := sysRe.FindStringSubmatch(string(body)); m != nil { // resolve the external entity
			if req, err := http.NewRequest(http.MethodGet, m[1], nil); err == nil {
				if resp, err := ssrfTestClient.Do(req); err == nil {
					resp.Body.Close()
				}
			}
		}
		w.WriteHeader(200)
	}))
	defer target.Close()

	spec, err := XXEOOBSpec(target.URL+"/import", "")
	if err != nil {
		t.Fatal(err)
	}
	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(), verify.Candidate{Target: target.URL, Spec: spec})
	if err != nil || v != verify.Proven {
		t.Fatalf("verdict=%v err=%v proof=%q, want proven", v, err, proof)
	}
}