package proxy

import (
	"fmt"
	"reflect"
	"testing"
)

func TestParseMarkers(t *testing.T) {
	tests := []struct {
		name     string
		template string
		wantPos  []IntrudePosition
		wantOut  string
		wantErr  bool
	}{
		{
			name:     "single position",
			template: "GET /users/§123§ HTTP/1.1",
			wantPos:  []IntrudePosition{{Index: 0, Base: "123"}},
			wantOut:  "GET /users/123 HTTP/1.1",
		},
		{
			name:     "two positions",
			template: "GET /a/§1§/b/§2§ HTTP/1.1",
			wantPos:  []IntrudePosition{{Index: 0, Base: "1"}, {Index: 1, Base: "2"}},
			wantOut:  "GET /a/1/b/2 HTTP/1.1",
		},
		{
			name:     "no markers",
			template: "GET / HTTP/1.1",
			wantPos:  nil,
			wantOut:  "GET / HTTP/1.1",
		},
		{
			name:     "empty marker",
			template: "GET /§§ HTTP/1.1",
			wantPos:  []IntrudePosition{{Index: 0, Base: ""}},
			wantOut:  "GET / HTTP/1.1",
		},
		{
			name:     "unterminated marker",
			template: "GET /§broken HTTP/1.1",
			wantErr:  true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			pos, out, err := ParseMarkers([]byte(tt.template))
			if tt.wantErr {
				if err == nil {
					t.Fatalf("expected error, got nil")
				}
				return
			}
			if err != nil {
				t.Fatalf("unexpected error: %v", err)
			}
			if !reflect.DeepEqual(pos, tt.wantPos) {
				t.Errorf("positions = %+v, want %+v", pos, tt.wantPos)
			}
			if string(out) != tt.wantOut {
				t.Errorf("stripped = %q, want %q", out, tt.wantOut)
			}
		})
	}
}

func TestBuildRequestValues(t *testing.T) {
	template := "GET /a/§1§/b/§2§/c/§3§ HTTP/1.1"

	tests := []struct {
		values []string
		want   string
	}{
		{[]string{"PAYLOAD", "2", "3"}, "GET /a/PAYLOAD/b/2/c/3 HTTP/1.1"},
		{[]string{"1", "PAYLOAD", "3"}, "GET /a/1/b/PAYLOAD/c/3 HTTP/1.1"},
		{[]string{"1", "2", "PAYLOAD"}, "GET /a/1/b/2/c/PAYLOAD HTTP/1.1"},
	}

	for _, tt := range tests {
		got, err := buildRequestValues([]byte(template), tt.values)
		if err != nil {
			t.Fatalf("values=%v: unexpected error: %v", tt.values, err)
		}
		if string(got) != tt.want {
			t.Errorf("values=%v: got %q, want %q", tt.values, got, tt.want)
		}
	}
}

func TestBuildRequestValuesPayloadContainingMarkerChar(t *testing.T) {
	// A payload that itself contains the marker character must not be
	// reinterpreted as a marker - buildRequestValues splits the ORIGINAL
	// template, never the already-substituted result.
	template := "GET /§1§/§2§ HTTP/1.1"
	got, err := buildRequestValues([]byte(template), []string{"§injected§", "2"})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	want := "GET /§injected§/2 HTTP/1.1"
	if string(got) != want {
		t.Errorf("got %q, want %q", got, want)
	}
}

func TestIntrudeValuesSniper(t *testing.T) {
	positions := []IntrudePosition{{Index: 0, Base: "a"}, {Index: 1, Base: "b"}}
	got, err := intrudeValues(Sniper, positions, [][]string{{"1", "2"}})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	want := [][]string{
		{"1", "b"}, {"2", "b"}, // position 0 fuzzed, position 1 at base
		{"a", "1"}, {"a", "2"}, // position 1 fuzzed, position 0 at base
	}
	if !reflect.DeepEqual(got, want) {
		t.Errorf("got %v, want %v", got, want)
	}
}

func TestIntrudeValuesBatteringRam(t *testing.T) {
	positions := []IntrudePosition{{Index: 0, Base: "a"}, {Index: 1, Base: "b"}}
	got, err := intrudeValues(BatteringRam, positions, [][]string{{"1", "2"}})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	// Same payload lands in every position at once, unlike Sniper.
	want := [][]string{{"1", "1"}, {"2", "2"}}
	if !reflect.DeepEqual(got, want) {
		t.Errorf("got %v, want %v", got, want)
	}
}

func TestIntrudeValuesPitchfork(t *testing.T) {
	positions := []IntrudePosition{{Index: 0, Base: "a"}, {Index: 1, Base: "b"}}
	got, err := intrudeValues(Pitchfork, positions, [][]string{{"1", "2", "3"}, {"x", "y"}})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	// Walks both sets in lockstep; stops at the shorter set's length (2),
	// silently ignoring "3" from the longer one - Burp's own convention.
	want := [][]string{{"1", "x"}, {"2", "y"}}
	if !reflect.DeepEqual(got, want) {
		t.Errorf("got %v, want %v", got, want)
	}
}

func TestIntrudeValuesPitchforkRequiresOneSetPerPosition(t *testing.T) {
	positions := []IntrudePosition{{Index: 0, Base: "a"}, {Index: 1, Base: "b"}}
	_, err := intrudeValues(Pitchfork, positions, [][]string{{"1", "2"}})
	if err == nil {
		t.Fatal("expected error for one payload set across two positions")
	}
}

func TestIntrudeValuesClusterBomb(t *testing.T) {
	positions := []IntrudePosition{{Index: 0, Base: "a"}, {Index: 1, Base: "b"}}
	got, err := intrudeValues(ClusterBomb, positions, [][]string{{"1", "2"}, {"x", "y"}})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	// Every combination - the rightmost (last) position cycles fastest.
	want := [][]string{{"1", "x"}, {"1", "y"}, {"2", "x"}, {"2", "y"}}
	if !reflect.DeepEqual(got, want) {
		t.Errorf("got %v, want %v", got, want)
	}
}

func TestIntrudeValuesClusterBombOverCapRejected(t *testing.T) {
	positions := []IntrudePosition{{Index: 0, Base: "a"}, {Index: 1, Base: "b"}, {Index: 2, Base: "c"}}
	big := make([]string, 20)
	for i := range big {
		big[i] = fmt.Sprintf("v%d", i)
	}
	// 20 * 20 * 20 = 8000, comfortably over the 1000 cap.
	_, err := intrudeValues(ClusterBomb, positions, [][]string{big, big, big})
	if err == nil {
		t.Fatal("expected the request-count cap to reject this attack")
	}
}

func TestIntrudeValuesOverCapRejected(t *testing.T) {
	positions := []IntrudePosition{{Index: 0, Base: "a"}}
	big := make([]string, maxIntrudeRequests+1)
	for i := range big {
		big[i] = fmt.Sprintf("v%d", i)
	}
	if _, err := intrudeValues(Sniper, positions, [][]string{big}); err == nil {
		t.Error("Sniper: expected the request-count cap to reject this attack")
	}
	if _, err := intrudeValues(BatteringRam, positions, [][]string{big}); err == nil {
		t.Error("BatteringRam: expected the request-count cap to reject this attack")
	}
}

func TestIntrudeRequestCount(t *testing.T) {
	positions, _, err := ParseMarkers([]byte("GET /§a§/§b§ HTTP/1.1"))
	if err != nil {
		t.Fatal(err)
	}
	if len(positions) != 2 {
		t.Fatalf("expected 2 positions, got %d", len(positions))
	}
}

func TestFixContentLength(t *testing.T) {
	tests := []struct {
		name string
		raw  string
		want string
	}{
		{
			name: "recalculates a stale length",
			raw:  "POST / HTTP/1.1\r\nHost: x\r\nContent-Length: 3\r\n\r\nfuzzedvalue",
			want: "POST / HTTP/1.1\r\nHost: x\r\nContent-Length: 11\r\n\r\nfuzzedvalue",
		},
		{
			name: "case-insensitive header name",
			raw:  "POST / HTTP/1.1\r\nHost: x\r\ncontent-length: 1\r\n\r\nabc",
			want: "POST / HTTP/1.1\r\nHost: x\r\nContent-Length: 3\r\n\r\nabc",
		},
		{
			name: "no content-length header - unchanged",
			raw:  "GET /§1§ HTTP/1.1\r\nHost: x\r\n\r\n",
			want: "GET /§1§ HTTP/1.1\r\nHost: x\r\n\r\n",
		},
		{
			name: "no body boundary - unchanged",
			raw:  "GET / HTTP/1.1\r\nHost: x",
			want: "GET / HTTP/1.1\r\nHost: x",
		},
		{
			name: "two content-length headers - left alone, ambiguous smuggling case",
			raw:  "POST / HTTP/1.1\r\nContent-Length: 3\r\nContent-Length: 999\r\n\r\nabc",
			want: "POST / HTTP/1.1\r\nContent-Length: 3\r\nContent-Length: 999\r\n\r\nabc",
		},
		{
			name: "empty body recalculates to zero",
			raw:  "POST / HTTP/1.1\r\nContent-Length: 5\r\n\r\n",
			want: "POST / HTTP/1.1\r\nContent-Length: 0\r\n\r\n",
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := fixContentLength([]byte(tt.raw))
			if string(got) != tt.want {
				t.Errorf("fixContentLength(%q) = %q, want %q", tt.raw, got, tt.want)
			}
		})
	}
}