package exploit

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

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

func TestAuthBypassHeaderSpecs_Shapes(t *testing.T) {
	specs, err := AuthBypassHeaderSpecs("https://t/admin?x=1", "Bearer low")
	if err != nil {
		t.Fatal(err)
	}
	byTech := map[string]*verify.ProbeSpec{}
	for _, ts := range specs {
		if !ts.Spec.Valid() || ts.Spec.Class != verify.ProbeDifferential {
			t.Fatalf("%s: invalid differential spec", ts.Technique)
		}
		byTech[ts.Technique] = ts.Spec
	}
	// URL-override: base requests "/", header carries the path; attack = the
	// protected path+query, control = the bogus path.
	xou := byTech["header:X-Original-URL"]
	if xou == nil {
		t.Fatal("missing X-Original-URL technique")
	}
	if !strings.Contains(xou.BaseRequest, "GET / HTTP/1.1") {
		t.Errorf("url-override should request the public path: %q", xou.BaseRequest)
	}
	if !strings.Contains(xou.BaseRequest, "X-Original-URL: "+verify.DefaultMarker) {
		t.Errorf("header not marked: %q", xou.BaseRequest)
	}
	if xou.AttackPayload != "/admin?x=1" {
		t.Errorf("attack = %q, want the protected path+query", xou.AttackPayload)
	}
	if !strings.Contains(xou.BaseRequest, "Authorization: Bearer low") {
		t.Error("credential should be carried")
	}
	// IP-spoof: base requests the protected path, header carries the IP.
	xff := byTech["header:X-Forwarded-For"]
	if xff == nil {
		t.Fatal("missing X-Forwarded-For technique")
	}
	if !strings.Contains(xff.BaseRequest, "GET /admin?x=1 HTTP/1.1") {
		t.Errorf("ip-spoof should request the protected path: %q", xff.BaseRequest)
	}
	if xff.AttackPayload != "127.0.0.1" || xff.ControlPayload != "203.0.113.1" {
		t.Errorf("ip payloads = attack %q control %q", xff.AttackPayload, xff.ControlPayload)
	}
}

// End to end: a server that serves /admin only when X-Original-URL points at it
// is proven via the differential oracle.
func TestAuthBypassHeaderSpecs_ProvesURLOverride(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("X-Original-URL") == "/admin" {
			w.WriteHeader(200)
			_, _ = w.Write([]byte("admin panel"))
			return
		}
		w.WriteHeader(404)
	}))
	defer srv.Close()

	specs, err := AuthBypassHeaderSpecs(srv.URL+"/admin", "")
	if err != nil {
		t.Fatal(err)
	}
	proven := false
	d := oracle.NewDifferential(oracle.LiveDiffProbe(srv.Client()))
	for _, ts := range specs {
		if ts.Technique != "header:X-Original-URL" {
			continue
		}
		v, _, _, err := d.Probe(context.Background(), verify.Candidate{Target: srv.URL, Spec: ts.Spec})
		if err == nil && v == verify.Proven {
			proven = true
		}
	}
	if !proven {
		t.Error("X-Original-URL bypass should be proven")
	}
}

// End to end: a server that trusts X-Forwarded-For: 127.0.0.1 for /admin.
func TestAuthBypassHeaderSpecs_ProvesIPSpoof(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.Contains(r.Header.Get("X-Forwarded-For"), "127.0.0.1") {
			w.WriteHeader(200)
			_, _ = w.Write([]byte("localhost-only admin"))
			return
		}
		w.WriteHeader(403)
	}))
	defer srv.Close()

	specs, err := AuthBypassHeaderSpecs(srv.URL+"/admin", "")
	if err != nil {
		t.Fatal(err)
	}
	proven := false
	d := oracle.NewDifferential(oracle.LiveDiffProbe(srv.Client()))
	for _, ts := range specs {
		if ts.Technique != "header:X-Forwarded-For" {
			continue
		}
		v, _, _, err := d.Probe(context.Background(), verify.Candidate{Target: srv.URL, Spec: ts.Spec})
		if err == nil && v == verify.Proven {
			proven = true
		}
	}
	if !proven {
		t.Error("X-Forwarded-For bypass should be proven")
	}
}