package exploit

import (
	"encoding/json"
	"fmt"
	"sort"
	"strings"

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

// MassAssignmentSpec turns a proven mass-assignment httpreq step into an
// external-oracle differential re-proof recipe. A mass-assignment hit is
// inherently a differential test: the attack body smuggled a privileged field
// (e.g. "role":"admin", "is_admin":true) and the server honored it, so replaying
// that exact body (the attack) against the same body with the privileged field
// removed (the control) and comparing responses re-proves the flaw with a signal
// no model produced — the privileged field took effect only when it was present.
//
// httpreqCommand is the chain's proof step (an `httpreq --method ... --url ...
// --body '{...}'` line); vars resolves any {{capture}} placeholders (e.g. the
// session token) the chain substituted at run time. Returns an error when the
// step is not a body-mutating request, the body is not JSON, or it carries no
// recognizable privileged field — so a producer can't attach a recipe that
// would silently never fire.
func MassAssignmentSpec(httpreqCommand string, vars map[string]string) (*verify.ProbeSpec, error) {
	cmd := substituteVars(httpreqCommand, vars)
	parts, err := parseCommand(cmd)
	if err != nil {
		return nil, fmt.Errorf("mass-assignment: parse step command: %w", err)
	}
	if len(parts) == 0 || !strings.EqualFold(parts[0], "httpreq") {
		return nil, fmt.Errorf("mass-assignment: not an httpreq step")
	}
	opts, err := parseHTTPReqArgs(parts[1:])
	if err != nil {
		return nil, fmt.Errorf("mass-assignment: %w", err)
	}
	switch opts.method {
	case "POST", "PUT", "PATCH":
	default:
		return nil, fmt.Errorf("mass-assignment: %s is not a body-mutating method", opts.method)
	}
	if strings.TrimSpace(opts.body) == "" {
		return nil, fmt.Errorf("mass-assignment: step has no --body to diff")
	}

	var obj map[string]any
	if err := json.Unmarshal([]byte(opts.body), &obj); err != nil {
		return nil, fmt.Errorf("mass-assignment: body is not a JSON object: %w", err)
	}
	field, ok := privilegedField(obj)
	if !ok {
		return nil, fmt.Errorf("mass-assignment: no privileged field found in body")
	}

	// Control = the same body with the privileged field removed. A correctly-
	// authorizing server returns the same result as the attack (the field was
	// ignored anyway); a vulnerable one diverges only when the field is present.
	control := make(map[string]any, len(obj))
	for k, v := range obj {
		if k == field {
			continue
		}
		control[k] = v
	}
	controlBody, err := json.Marshal(control)
	if err != nil {
		return nil, fmt.Errorf("mass-assignment: marshal control body: %w", err)
	}

	raw, err := jsonBodyRawRequest(opts)
	if err != nil {
		return nil, err
	}
	// Indicator: a distinctive privileged VALUE (e.g. "admin", "superuser")
	// reflected in the attack response but not the control is the strongest
	// signal the field was mass-assigned. Boolean/numeric values make poor
	// indicators (too common / easily coincidental), so we leave the indicator
	// empty there and let the oracle fall back to response divergence.
	indicator := distinctiveValue(obj[field])
	return verify.NewDifferentialSpec(raw, "", opts.body, string(controlBody), indicator)
}

// jsonBodyRawRequest builds a raw HTTP/1.1 request whose body is the injection
// marker (the differential oracle substitutes the attack and control bodies
// there). Host and path are split by hand so the marker is placed verbatim.
// Headers from the step are preserved (they carry the session credential); a
// JSON content type is ensured.
func jsonBodyRawRequest(opts httpReqOptions) (string, error) {
	rest := opts.url
	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:]
	}
	if path == "" {
		path = "/"
	}

	var b strings.Builder
	b.WriteString(opts.method + " " + path + " HTTP/1.1\r\n")
	if host != "" {
		b.WriteString("Host: " + host + "\r\n")
	}
	hasContentType := false
	for _, h := range opts.headers {
		if k, _, ok := splitHeader(h); ok && strings.EqualFold(k, "Content-Type") {
			hasContentType = true
		}
		b.WriteString(h + "\r\n")
	}
	if !hasContentType {
		b.WriteString("Content-Type: application/json\r\n")
	}
	b.WriteString("\r\n")
	b.WriteString(verify.DefaultMarker)
	return b.String(), nil
}

// privilegedFieldKeys are the classic mass-assignment escalation fields an API
// should never accept from a client-supplied body.
var privilegedFieldKeys = map[string]struct{}{
	"is_admin": {}, "isadmin": {}, "admin": {},
	"role": {}, "roles": {}, "user_role": {}, "userrole": {},
	"is_staff": {}, "isstaff": {}, "staff": {},
	"is_superuser": {}, "superuser": {}, "is_super": {},
	"privilege": {}, "privileges": {}, "perm": {}, "perms": {},
	"permission": {}, "permissions": {}, "scope": {}, "scopes": {},
	"account_type": {}, "accounttype": {}, "type": {}, "tier": {}, "plan": {}, "grade": {},
	"is_verified": {}, "isverified": {}, "verified": {}, "email_verified": {},
	"is_premium": {}, "premium": {}, "balance": {}, "credit": {}, "credits": {}, "wallet": {},
	"owner": {}, "owner_id": {}, "ownerid": {}, "user_id": {}, "userid": {}, "org_id": {},
}

// privilegedField returns the first top-level body key that names a classic
// mass-assignment escalation field. Keys are checked case-insensitively, in
// sorted order so the choice is deterministic.
func privilegedField(obj map[string]any) (string, bool) {
	keys := make([]string, 0, len(obj))
	for k := range obj {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for _, k := range keys {
		if _, ok := privilegedFieldKeys[strings.ToLower(k)]; ok {
			return k, true
		}
	}
	return "", false
}

// distinctiveValue returns a privileged field's value as a success indicator
// when it is a distinctive string (length ≥ 3, so "admin"/"superuser" qualify
// but a bare "1" or a boolean does not). Non-string or short values return ""
// so the oracle falls back to response divergence rather than a weak indicator.
func distinctiveValue(v any) string {
	s, ok := v.(string)
	if !ok {
		return ""
	}
	s = strings.TrimSpace(s)
	if len(s) < 3 {
		return ""
	}
	return s
}