// Intruder-equivalent: mark positions in a raw request template with §
// (Burp's own marker character, so anyone who's used Burp already knows
// the syntax) and fuzz them across four attack modes - Sniper, Battering
// ram, Pitchfork, and Cluster bomb - matching Burp's own naming and
// semantics for how positions and payload sets combine.
package proxy

import (
	"bytes"
	"context"
	"fmt"

	"mitmux/internal/store"
)

const marker = "§"

// maxIntrudeRequests caps the number of requests one attack can send - a
// safety limit against an accidental huge wordlist (or, for Cluster bomb,
// a payload-set product) turning into an unbounded flood, not a tuned
// production value.
const maxIntrudeRequests = 1000

// AttackMode selects how payload sets combine across marked positions,
// matching Burp's own four attack types.
type AttackMode string

const (
	// Sniper fuzzes one position at a time through a single shared
	// payload set; every other marked position holds its base value.
	// Requests: positions × len(payloads).
	Sniper AttackMode = "sniper"
	// BatteringRam sends the same payload, from a single shared payload
	// set, into every marked position at once. Requests: len(payloads).
	BatteringRam AttackMode = "battering_ram"
	// Pitchfork walks one payload set per position in lockstep - request
	// i takes payload i from every set. Requests: the shortest set's
	// length (Burp's own convention when sets are uneven).
	Pitchfork AttackMode = "pitchfork"
	// ClusterBomb tries every combination of one payload set per
	// position. Requests: the product of every set's length.
	ClusterBomb AttackMode = "cluster_bomb"
)

// IntrudePosition is one marked, resolved insertion point.
type IntrudePosition struct {
	Index int    // 0-based, in order of appearance
	Base  string // the text between its markers
}

// ParseMarkers finds every §base§ pair in template and returns the
// resolved positions plus template with the markers stripped out (the
// form actually used as the base request when no position is being
// fuzzed). An odd number of § markers is a user error - unterminated
// marker - reported rather than guessed at.
func ParseMarkers(template []byte) (positions []IntrudePosition, stripped []byte, err error) {
	parts := bytes.Split(template, []byte(marker))
	if len(parts)%2 != 1 {
		return nil, nil, fmt.Errorf("unterminated %s marker - markers must come in pairs", marker)
	}
	if len(parts) == 1 {
		return nil, template, nil
	}

	var buf bytes.Buffer
	for i, part := range parts {
		if i%2 == 1 {
			positions = append(positions, IntrudePosition{Index: len(positions), Base: string(part)})
		}
		buf.Write(part)
	}
	return positions, buf.Bytes(), nil
}

// buildRequestValues re-derives offsets from template's ORIGINAL marker
// layout (rather than operating on already-stripped bytes) and substitutes
// values[i] for the i-th marked position, in order. len(values) must equal
// the number of marked positions in template.
func buildRequestValues(template []byte, values []string) ([]byte, error) {
	parts := bytes.Split(template, []byte(marker))
	if len(parts)%2 != 1 {
		return nil, fmt.Errorf("unterminated %s marker", marker)
	}
	var buf bytes.Buffer
	pos := 0
	for i, part := range parts {
		if i%2 == 1 {
			if pos >= len(values) {
				return nil, fmt.Errorf("position %d has no value", pos)
			}
			buf.WriteString(values[pos])
			pos++
			continue
		}
		buf.Write(part)
	}
	return buf.Bytes(), nil
}

// intrudeValues computes, for one full attack, every position-substitution
// set to send - one []string per request (indexed by position, in send
// order) - according to mode. Pure and side-effect free, so the request
// count can be validated against maxIntrudeRequests before anything is
// dispatched, and so it's testable without a live target.
//
// payloadSets[0] is the shared payload set for Sniper and BatteringRam,
// which only ever need one. Pitchfork and ClusterBomb are inherently
// per-position - theirs is the whole point of the two modes - so they
// require exactly len(positions) sets, one per marked position in order.
func intrudeValues(mode AttackMode, positions []IntrudePosition, payloadSets [][]string) ([][]string, error) {
	if len(positions) == 0 {
		return nil, fmt.Errorf("no %s-marked positions in the request template", marker)
	}
	if len(payloadSets) == 0 || len(payloadSets[0]) == 0 {
		return nil, fmt.Errorf("no payloads")
	}

	bases := make([]string, len(positions))
	for i, p := range positions {
		bases[i] = p.Base
	}

	var out [][]string
	switch mode {
	case "", Sniper:
		payloads := payloadSets[0]
		if total := len(positions) * len(payloads); total > maxIntrudeRequests {
			return nil, fmt.Errorf("attack would send %d requests (%d positions × %d payloads), over the %d limit",
				total, len(positions), len(payloads), maxIntrudeRequests)
		}
		for posIdx := range positions {
			for _, payload := range payloads {
				values := append([]string(nil), bases...)
				values[posIdx] = payload
				out = append(out, values)
			}
		}

	case BatteringRam:
		payloads := payloadSets[0]
		if len(payloads) > maxIntrudeRequests {
			return nil, fmt.Errorf("attack would send %d requests, over the %d limit", len(payloads), maxIntrudeRequests)
		}
		for _, payload := range payloads {
			values := make([]string, len(positions))
			for i := range values {
				values[i] = payload
			}
			out = append(out, values)
		}

	case Pitchfork:
		if len(payloadSets) != len(positions) {
			return nil, fmt.Errorf("pitchfork needs one payload set per position (%d positions, %d payload sets given)",
				len(positions), len(payloadSets))
		}
		n := len(payloadSets[0])
		for _, set := range payloadSets {
			if len(set) == 0 {
				return nil, fmt.Errorf("no payloads")
			}
			if len(set) < n {
				n = len(set)
			}
		}
		if n > maxIntrudeRequests {
			return nil, fmt.Errorf("attack would send %d requests, over the %d limit", n, maxIntrudeRequests)
		}
		for i := 0; i < n; i++ {
			values := make([]string, len(positions))
			for p := range positions {
				values[p] = payloadSets[p][i]
			}
			out = append(out, values)
		}

	case ClusterBomb:
		if len(payloadSets) != len(positions) {
			return nil, fmt.Errorf("cluster bomb needs one payload set per position (%d positions, %d payload sets given)",
				len(positions), len(payloadSets))
		}
		// Checked incrementally, one set at a time, so a pathological
		// product (e.g. three sets of 10000) bails out before ever
		// trying to enumerate it, not after.
		total := 1
		for _, set := range payloadSets {
			if len(set) == 0 {
				return nil, fmt.Errorf("no payloads")
			}
			total *= len(set)
			if total > maxIntrudeRequests {
				return nil, fmt.Errorf("attack would send at least %d requests, over the %d limit", total, maxIntrudeRequests)
			}
		}
		idx := make([]int, len(positions))
		for {
			values := make([]string, len(positions))
			for p := range positions {
				values[p] = payloadSets[p][idx[p]]
			}
			out = append(out, values)

			// Odometer increment, rightmost (last) position fastest -
			// matches Burp's own cluster-bomb iteration order.
			p := len(positions) - 1
			for p >= 0 {
				idx[p]++
				if idx[p] < len(payloadSets[p]) {
					break
				}
				idx[p] = 0
				p--
			}
			if p < 0 {
				break
			}
		}

	default:
		return nil, fmt.Errorf("unknown attack mode %q", mode)
	}
	return out, nil
}

// Intrude runs one attack of the given mode over a §marked§ request
// template. onResult is called synchronously after each request completes
// - with a 0-based iteration index, the values substituted into each
// marked position for that request (indexed by position, same order as
// ParseMarkers), the resulting entry (nil if sendErr is set), and any send
// error - so a caller can stream progress, and stops the attack early if
// it returns false.
func (s *Server) Intrude(ctx context.Context, scheme, host string, template []byte, mode AttackMode, payloadSets [][]string,
	onResult func(iteration int, values []string, entry *store.Entry, sendErr error) bool) error {
	positions, _, err := ParseMarkers(template)
	if err != nil {
		return err
	}

	requests, err := intrudeValues(mode, positions, payloadSets)
	if err != nil {
		return err
	}

	for i, values := range requests {
		raw, err := buildRequestValues(template, values)
		if err != nil {
			return err
		}
		raw = fixContentLength(raw)

		// sendRaw is already self-bounding (dialForRepeat's own dial
		// timeout, then conn.SetDeadline for the rest), so ctx here
		// only needs to carry cancellation - e.g. the IPC connection
		// driving this attack closing mid-run.
		e, sendErr := s.sendRaw(ctx, scheme, host, raw, "intruder")

		if !onResult(i, values, e, sendErr) {
			return nil
		}
	}
	return nil
}

// fixContentLength recalculates an existing Content-Length header to
// match raw's actual body length after marker substitution. A fuzzed
// payload routinely differs in length from the base value it replaces;
// left as-is, a Content-Length carried over unchanged from the original
// captured request makes the target server wait for bytes that will
// never arrive, hanging that request for the full upstream timeout -
// confirmed: identical attacks against a URL-only marker (no body
// length change) completed in single-digit milliseconds per payload,
// the same attack with the marker inside a body parameter took 60s per
// payload. This is Intruder-specific, not something Repeater does to
// what's typed: a fuzzed value's length is a side effect of automated
// substitution, whereas a Repeater edit is deliberate and Repeater's own
// "no auto-fixed Content-Length" behavior is unchanged.
//
// Only touches a request with exactly one Content-Length header and a
// clean header/body boundary - zero found means nothing to fix, more
// than one is a request smuggling test's own deliberately ambiguous
// framing, and guessing which one to rewrite there would be worse than
// leaving both alone.
func fixContentLength(raw []byte) []byte {
	sep := []byte("\r\n\r\n")
	idx := bytes.Index(raw, sep)
	if idx < 0 {
		return raw
	}
	headerBlock, body := raw[:idx], raw[idx+len(sep):]

	lines := bytes.Split(headerBlock, []byte("\r\n"))
	foundIdx, count := -1, 0
	for i, line := range lines {
		if i == 0 {
			continue // request line, not a header
		}
		colon := bytes.IndexByte(line, ':')
		if colon < 0 {
			continue
		}
		if bytes.EqualFold(bytes.TrimSpace(line[:colon]), []byte("Content-Length")) {
			count++
			foundIdx = i
		}
	}
	if count != 1 {
		return raw
	}

	lines[foundIdx] = []byte(fmt.Sprintf("Content-Length: %d", len(body)))
	var out bytes.Buffer
	out.Write(bytes.Join(lines, []byte("\r\n")))
	out.Write(sep)
	out.Write(body)
	return out.Bytes()
}