package exploit

import (
	"strings"

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

// BOLASpec turns a confirmed BOLA/IDOR hit into an external-oracle re-proof
// recipe. A BOLA hit is inherently a differential test: the harvested victim id
// returned a 200 for a session that does not own it, so replaying that id (the
// attack) against a well-formed but non-existent id (the control) and comparing
// responses re-proves the broken authorization with a signal no model produced.
//
// template is an id-templated URL (the "{id}" slot DeriveBOLATargets produces);
// victimID is the id that yielded unauthorized access; authHeader is the
// session's own credential (e.g. "Bearer <jwt>") so the probe reproduces the
// exact authenticated context. Returns an error if no "{id}" slot is present.
func BOLASpec(template, victimID, authHeader string) (*verify.ProbeSpec, error) {
	raw, err := bolaRawRequest(template, authHeader)
	if err != nil {
		return nil, err
	}
	// The control is a well-formed id of the same shape that should not exist,
	// so a correctly-authorizing endpoint denies it — the differential collapses
	// to a 401/403/404 and the oracle can both prove and refute.
	return verify.NewDifferentialSpec(raw, "", victimID, controlID(victimID), "")
}

// bolaRawRequest builds a raw GET request from an id-templated URL. Host and
// path are split by hand (not via net/url) so the injection marker is placed
// into the path verbatim — url.Parse would percent-encode the § marker.
func bolaRawRequest(template, authHeader string) (string, error) {
	rest := template
	for _, p := range []string{"https://", "http://"} {
		rest = strings.TrimPrefix(rest, p)
	}
	host, path := rest, "/"
	if i := strings.IndexAny(rest, "/?#"); i >= 0 {
		host, path = rest[:i], rest[i:]
	}
	path = strings.ReplaceAll(path, "{id}", verify.DefaultMarker)

	var b strings.Builder
	b.WriteString("GET " + path + " HTTP/1.1\r\n")
	if host != "" {
		b.WriteString("Host: " + host + "\r\n")
	}
	if authHeader != "" {
		b.WriteString("Authorization: " + authHeader + "\r\n")
	}
	b.WriteString("\r\n")
	return b.String(), nil
}

// BOLATemplateAndID recovers the id-templated form and the concrete object id
// from a BOLA hit URL (which has the victim id already substituted in). It
// templatises the first UUID/numeric path segment — the same slot
// DeriveBOLATargets chooses — returning the "{id}" template and that id. ok is
// false when the URL has no id-shaped segment to pivot on.
func BOLATemplateAndID(concreteURL string) (template, id string, ok bool) {
	rest := concreteURL
	scheme := ""
	for _, p := range []string{"https://", "http://"} {
		if strings.HasPrefix(rest, p) {
			scheme, rest = p, strings.TrimPrefix(rest, p)
			break
		}
	}
	host, path := rest, ""
	if i := strings.IndexAny(rest, "/?#"); i >= 0 {
		host, path = rest[:i], rest[i:]
	}
	segs := strings.Split(path, "/")
	for i, s := range segs {
		if uuidRe.MatchString(s) || (len(s) >= 2 && isAllDigits(s)) {
			id = s
			segs[i] = "{id}"
			return scheme + host + strings.Join(segs, "/"), id, true
		}
	}
	return "", "", false
}

// controlID returns a well-formed but almost-certainly-nonexistent id matching
// the victim id's shape (UUID vs numeric), so the control request is a fair
// negative — same route, same type, an id the user also shouldn't reach.
func controlID(victimID string) string {
	if uuidRe.MatchString(victimID) {
		return "00000000-0000-0000-0000-000000000000"
	}
	if isAllDigits(victimID) {
		// A large id unlikely to exist; distinct from any small harvested id.
		return "999999999"
	}
	return "nonexistent-" + victimID
}