package exploit

import (
	"context"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strconv"
	"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 TestSQLiTimingSpec_BuildsTimingRecipe(t *testing.T) {
	spec, err := SQLiTimingSpec("https://t/items?id=1&sort=asc",
		SQLiInjection{Place: "GET", Parameter: "id", DBMS: "MySQL >= 5.0.12"}, "Bearer tok")
	if err != nil {
		t.Fatal(err)
	}
	if spec.Class != verify.ProbeTiming || !spec.Valid() {
		t.Fatalf("not a valid timing spec: %+v", spec)
	}
	if !strings.Contains(decodeTemplate(spec.DelayTemplate), "SLEEP("+verify.DelayMarker+")") {
		t.Errorf("delay template missing MySQL sleep: %q", spec.DelayTemplate)
	}
	if !strings.Contains(spec.BaseRequest, "id="+verify.DefaultMarker) {
		t.Errorf("injectable param not marked: %q", spec.BaseRequest)
	}
	if !strings.Contains(spec.BaseRequest, "sort=asc") {
		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")
	}
}

// decodeTemplate reverses the query/form percent-encoding the producer applies
// to a delay payload, leaving the §DELAY§ marker literal, so tests can assert on
// the readable SQL the oracle will send.
func decodeTemplate(tmpl string) string {
	parts := strings.Split(tmpl, verify.DelayMarker)
	for i, p := range parts {
		if dec, err := url.QueryUnescape(p); err == nil {
			parts[i] = dec
		}
	}
	return strings.Join(parts, verify.DelayMarker)
}

func TestSQLiTimingSpec_DBMSPayloads(t *testing.T) {
	cases := []struct{ dbms, want string }{
		{"MySQL", "SLEEP("},
		{"PostgreSQL 13", "PG_SLEEP("},
		{"Microsoft SQL Server 2019", "WAITFOR DELAY"},
		{"Oracle", "DBMS_PIPE.RECEIVE_MESSAGE"},
	}
	for _, c := range cases {
		spec, err := SQLiTimingSpec("https://t/x?p=1", SQLiInjection{Parameter: "p", DBMS: c.dbms}, "")
		if err != nil {
			t.Fatalf("%s: %v", c.dbms, err)
		}
		if !strings.Contains(decodeTemplate(spec.DelayTemplate), c.want) {
			t.Errorf("%s: template %q missing %q", c.dbms, spec.DelayTemplate, c.want)
		}
	}
}

func TestSQLiTimingSpec_RejectsUnknownDBMS(t *testing.T) {
	if _, err := SQLiTimingSpec("https://t/x?p=1", SQLiInjection{Parameter: "p", DBMS: "CouchDB"}, ""); err == nil {
		t.Error("unknown DBMS should be rejected, not silently non-firing")
	}
	// SQLite has no linear time-based primitive and must be refused.
	if _, err := SQLiTimingSpec("https://t/x?p=1", SQLiInjection{Parameter: "p", DBMS: "SQLite"}, ""); err == nil {
		t.Error("SQLite should be rejected")
	}
}

func TestSQLiTimingSpec_POSTBody(t *testing.T) {
	spec, err := SQLiTimingSpec("https://t/login", SQLiInjection{Place: "POST", Parameter: "user", DBMS: "MySQL"}, "")
	if err != nil {
		t.Fatal(err)
	}
	if !strings.HasPrefix(spec.BaseRequest, "POST /login HTTP/1.1") {
		t.Errorf("expected POST request line: %q", spec.BaseRequest)
	}
	if !strings.Contains(spec.BaseRequest, "user="+verify.DefaultMarker) {
		t.Errorf("POST body should mark the param: %q", spec.BaseRequest)
	}
	if !strings.Contains(spec.BaseRequest, "Content-Type: application/x-www-form-urlencoded") {
		t.Errorf("POST should set a form content type: %q", spec.BaseRequest)
	}
}

func TestSQLiTimingSpec_RejectsUnsupportedPlace(t *testing.T) {
	if _, err := SQLiTimingSpec("https://t/x", SQLiInjection{Place: "COOKIE", Parameter: "sid", DBMS: "MySQL"}, ""); err == nil {
		t.Error("unsupported place should be rejected")
	}
}

func TestSQLiInjectionPoints_FromSqlmapData(t *testing.T) {
	// A representative sqlmap /scan/data shape: injection objects in a "value"
	// list, DBMS both top-level and (for the second) only in the technique data.
	parsed := []map[string]any{
		{"type": float64(1), "value": []any{
			map[string]any{
				"place": "GET", "parameter": "id", "dbms": "MySQL",
			},
			map[string]any{
				"place": "GET", "parameter": "id", "dbms": "MySQL", // duplicate
			},
			map[string]any{
				"place": "POST", "parameter": "q",
				"data": map[string]any{
					"5": map[string]any{"title": "PostgreSQL > 8.1 AND time-based blind", "payload": "q=1 AND PG_SLEEP(5)"},
				},
			},
			map[string]any{"place": "GET"}, // no parameter → skipped
		}},
	}
	pts := SQLiInjectionPoints(parsed)
	if len(pts) != 2 {
		t.Fatalf("got %d points, want 2: %+v", len(pts), pts)
	}
	if pts[0].Parameter != "id" || normalizeDBMS(pts[0].DBMS) != "mysql" {
		t.Errorf("point 0 = %+v", pts[0])
	}
	if pts[1].Parameter != "q" || normalizeDBMS(pts[1].DBMS) != "postgresql" {
		t.Errorf("point 1 DBMS should come from technique title: %+v", pts[1])
	}
}

func TestSQLiInjectionPoints_Empty(t *testing.T) {
	if pts := SQLiInjectionPoints(nil); pts != nil {
		t.Errorf("nil input should yield no points, got %+v", pts)
	}
	if pts := SQLiInjectionPoints([]map[string]any{{"status": float64(1)}}); len(pts) != 0 {
		t.Errorf("non-injection entry should yield no points, got %+v", pts)
	}
}

// End to end: the produced recipe actually re-proves a blind SQLi against a
// live endpoint that sleeps for the injected number of seconds.
func TestSQLiTimingSpec_ProvesViaTimingOracle(t *testing.T) {
	if testing.Short() {
		t.Skip("timing oracle sends real delayed requests")
	}
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Emulate a time-based injection: §PAYLOAD§ carries SLEEP(n), where n is
		// the seconds the oracle injected at §DELAY§. Sleep that real duration so
		// latency tracks the delay (0→0, and scaling for the ladder test). The
		// oracle is configured with sub-second delays so the test stays fast.
		val := r.URL.Query().Get("id")
		if i := strings.Index(val, "SLEEP("); i >= 0 {
			rest := val[i+len("SLEEP("):]
			if j := strings.IndexByte(rest, ')'); j >= 0 {
				if secs, err := strconv.ParseFloat(strings.TrimSpace(rest[:j]), 64); err == nil && secs > 0 {
					time.Sleep(time.Duration(secs * float64(time.Second)))
				}
			}
		}
		w.WriteHeader(200)
	}))
	defer srv.Close()

	spec, err := SQLiTimingSpec(srv.URL+"/items?id=1", SQLiInjection{Place: "GET", Parameter: "id", DBMS: "MySQL"}, "")
	if err != nil {
		t.Fatal(err)
	}
	cand := verify.Candidate{Target: srv.URL, Spec: spec}
	// Sub-second delays + ladder mode: the injected number (0.2 / 0.4) lands in
	// SLEEP(), the server sleeps that long, and the induced latency must scale
	// with the injected delay — exactly as a real time-based injection does.
	o := oracle.NewTiming(oracle.LiveTimeProbe(srv.Client()), oracle.WithTimingConfig(oracle.TimingConfig{
		Delay:           200 * time.Millisecond,
		BaselineSamples: 5,
		AttackSamples:   3,
		Ladder:          true,
	}))
	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)
	}
}