package exploit

import (
	"context"
	"net/http"
	"net/http/httptest"
	"reflect"
	"strings"
	"testing"
	"time"

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

func TestURLShapedValue(t *testing.T) {
	cases := []struct {
		v    string
		want bool
	}{
		{"http://internal/", true},
		{"https://example.com/x", true},
		{"//evil.com", true},
		{"example.com", true},
		{"example.com:8080/path", true},
		{"169.254.169.254", false}, // numeric, no letter → not treated as host by shape alone
		{"3.14", false},
		{"", false},
		{"hello world", false},
		{"42", false},
		{"localhost", false}, // no dot
	}
	for _, c := range cases {
		if got := URLShapedValue(c.v); got != c.want {
			t.Errorf("URLShapedValue(%q) = %v, want %v", c.v, got, c.want)
		}
	}
}

func TestDetectSSRFParams(t *testing.T) {
	cases := []struct {
		url  string
		want []string
	}{
		{"https://t/fetch?url=http://x", []string{"url"}},                      // by name
		{"https://t/p?q=search&next=/home", []string{"next"}},                  // by name
		{"https://t/p?thumb=https://cdn.example.com/a.png", []string{"thumb"}}, // by URL-shaped value
		{"https://t/p?id=42&name=bob", nil},                                    // nothing SSRF-ish
		{"https://t/p?url=a&callback=b", []string{"url", "callback"}},          // multiple, in order
		{"https://t/nopath", nil},                                              // no query
	}
	for _, c := range cases {
		got := DetectSSRFParams(c.url)
		if !reflect.DeepEqual(got, c.want) {
			t.Errorf("DetectSSRFParams(%q) = %v, want %v", c.url, got, c.want)
		}
	}
}

func TestDetectSSRFParams_DedupesRepeatedName(t *testing.T) {
	got := DetectSSRFParams("https://t/p?url=a&url=b")
	if !reflect.DeepEqual(got, []string{"url"}) {
		t.Errorf("repeated param should dedupe, got %v", got)
	}
}

func TestSSRFOOBSpec_BuildsOOBRecipe(t *testing.T) {
	spec, err := SSRFOOBSpec("https://t/fetch?url=http://old&x=1", "url", "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, "url="+verify.DefaultMarker) {
		t.Errorf("sink param not marked: %q", spec.BaseRequest)
	}
	if !strings.Contains(spec.BaseRequest, "x=1") {
		t.Errorf("other params should be preserved: %q", spec.BaseRequest)
	}
	if !strings.Contains(spec.BaseRequest, "Authorization: Bearer tok") {
		t.Error("base request should carry the session credential")
	}
	if spec.OOBTemplate != verify.OOBMarker {
		t.Errorf("OOB template = %q, want the bare marker", spec.OOBTemplate)
	}
}

func TestSSRFOOBSpec_RejectsEmptyParam(t *testing.T) {
	if _, err := SSRFOOBSpec("https://t/fetch?url=x", "", ""); err == nil {
		t.Error("empty parameter should be rejected")
	}
}

// fakeOOBSession records whether a callback arrived at the reserved URL — enough
// to drive the OOB oracle end to end without a live interactsh listener.
type fakeOOBSession struct {
	url   string
	hit   *bool
	proof string
}

func (f fakeOOBSession) Payload(context.Context) (string, string, error) {
	return f.url, "tok", nil
}
func (f fakeOOBSession) Observed(context.Context, string, time.Duration) (bool, string, error) {
	return *f.hit, f.proof, nil
}

// End to end: the produced recipe actually re-proves an SSRF. The "target"
// dereferences the url parameter (simulating the vulnerability), calling back to
// the OOB listener, which the oracle then observes as Proven.
func TestSSRFOOBSpec_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()

	target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Vulnerable: the server fetches whatever the url parameter points at.
		if dest := r.URL.Query().Get("url"); dest != "" {
			req, _ := http.NewRequest(http.MethodGet, dest, nil)
			if resp, err := ssrfTestClient.Do(req); err == nil {
				resp.Body.Close()
			}
		}
		w.WriteHeader(200)
	}))
	defer target.Close()

	spec, err := SSRFOOBSpec(target.URL+"/fetch?url=seed", "url", "")
	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()), oracle.WithWindow(2*time.Second))
	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)
	}
}

var ssrfTestClient = &http.Client{Timeout: 5 * time.Second}