package dibaguard

import (
	"bytes"
	"encoding/json"
	"math"
	"net/http"
	"net/url"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"unicode"
)

// The Python request parser (dibaguard/http/parse.py Request.params / model_text) in Go: query pairs, form, JSON
// (every occurrence of a duplicate key, Python's number spelling), multipart (every field; a file's name and the
// head of its content), cookies, the four standard headers and any other header whose value looks like a payload
// (every value of a repeated header), all under the same judging budget - so Canonical is the Python model text.

const (
	parseMaxPart      = 8192    // parse.MAX_PART
	parseChunkOverlap = 512     // parse.CHUNK_OVERLAP
	parseMaxText      = 32768   // parse.MAX_TEXT
	parseExtraText    = 32768   // parse.EXTRA_TEXT
	parseMaxFields    = 2000    // parse.MAX_FIELDS
	parseMaxDepth     = 512     // parse.MAX_DEPTH
	parseMaxFile      = 8192    // parse.MAX_FILE
	parseRestMax      = 1 << 20 // parse.REST_MAX
)

var (
	stdHeaders   = []string{"user-agent", "referer", "x-forwarded-for", "content-type"}
	neverHeaders = map[string]bool{"cookie": true, "host": true, "content-length": true, "connection": true, "accept-encoding": true,
		"cache-control": true, "pragma": true, "upgrade-insecure-requests": true}
	hdrSuspicious = regexp.MustCompile("['\"<>`]|\\$\\{|\\{\\{|--|/\\*|\\.\\./|(?:^|[^a-z0-9_])(?:union|select|sleep|benchmark|jndi|script|javascript|vbscript|document|onerror|eval|exec" +
		"|cmd|wget|curl|waitfor|extractvalue|updatexml|information_schema|passwd|whoami)(?:$|[^a-z0-9_])")
	keepWords = map[string]bool{}
	reBound   = regexp.MustCompile(`(?i)boundary=(?:"([^"]+)"|([^";,\s]+))`)
	reDispNm  = regexp.MustCompile(`(?i)(?:^|;)\s*name\s*=\s*(?:"([^"]*)"|([^;\s]*))`)
	reDispFn  = regexp.MustCompile(`(?i)(?:^|;)\s*filename\s*=\s*(?:"([^"]*)"|([^;\s]*))`)
)

func init() {
	for _, w := range strings.Fields("or and xor not union select sleep benchmark waitfor having exec execute declare drop insert update delete truncate " +
		"shutdown grant information_schema pg_sleep extractvalue updatexml dual concat char chr rlike regexp like null jndi script javascript passwd " +
		"whoami cmd eval system") {
		keepWords[w] = true
	}
}

// keepPart is parse._keep: can this canonical text carry a payload?
func keepPart(s string) bool {
	if strings.Contains(s, "--") || strings.Contains(s, "..") {
		return true
	}
	for _, r := range s {
		if !(unicode.IsLetter(r) || unicode.IsNumber(r) || r == '_' || r == ' ' || r == '.' || r == '-') {
			return true
		}
	}
	for i := 0; i < len(s); {
		c := s[i]
		if 'a' <= c && c <= 'z' || '0' <= c && c <= '9' || c == '_' {
			j := i
			for j < len(s) && ('a' <= s[j] && s[j] <= 'z' || '0' <= s[j] && s[j] <= '9' || s[j] == '_') {
				j++
			}
			if keepWords[s[i:j]] {
				return true
			}
			i = j
			continue
		}
		i++
	}
	return false
}

// chunks is parse._chunks (code points).
func chunks(v string) []string {
	n := runeLen(v)
	if n <= parseMaxPart {
		return []string{v}
	}
	rs := []rune(v)
	step := parseMaxPart - parseChunkOverlap
	var out []string
	for i := 0; i < n-parseChunkOverlap; i += step {
		e := i + parseMaxPart
		if e > n {
			e = n
		}
		out = append(out, string(rs[i:e]))
	}
	return out
}

// textBuilder is parse._Text.
type textBuilder struct {
	parts                      []string
	n, extra, over, rawLeft, d int
}

func (t *textBuilder) add(prefix, value string) {
	pl, vl := runeLen(prefix), runeLen(value)
	if t.n+pl+vl+3 <= parseMaxText {
		t.parts = append(t.parts, prefix+value)
		t.n += pl + vl + 3
		return
	}
	for _, c := range chunks(value) {
		k := pl + runeLen(c) + 3
		switch {
		case t.n+k <= parseMaxText:
			t.parts = append(t.parts, prefix+c)
			t.n += k
		case keepPart(c):
			if t.extra+k <= parseExtraText {
				t.parts = append(t.parts, prefix+c)
				t.extra += k
			} else {
				t.over += runeLen(c)
			}
		}
	}
}

func (t *textBuilder) addRaw(prefix, raw string) {
	n := runeLen(raw)
	if n > t.rawLeft {
		t.over += n - t.rawLeft
		raw = runeCut(raw, t.rawLeft)
		n = t.rawLeft
	}
	t.rawLeft -= n
	if n == 0 {
		return
	}
	rs := []rune(raw)
	step := canonMaxLen - parseChunkOverlap
	hi := n - parseChunkOverlap
	if hi < 1 {
		hi = 1
	}
	for i := 0; i < hi; i += step {
		e := i + canonMaxLen
		if e > n {
			e = n
		}
		a, p, ok, r := canonText(string(rs[i:e]), false)
		if r > t.d {
			t.d = r
		}
		t.add(prefix, a)
		if ok {
			t.add(prefix, p)
		}
	}
}

func (t *textBuilder) text() string {
	if t.over > 0 {
		t.parts = append(t.parts, "x:oversized="+strconv.Itoa(t.over))
	}
	return strings.Join(t.parts, " | ")
}

type param struct{ where, k, v string }

// pairs is parse._pairs: split on '&' only, a pair at its first '='; nothing decoded yet.
func pairs(s string, where string, out []param) []param {
	for _, p := range strings.Split(s, "&") {
		if p == "" {
			continue
		}
		k, v, _ := strings.Cut(p, "=")
		out = append(out, param{where, k, v})
	}
	return out
}

// pyFloat is Python's repr(float).
func pyFloat(f float64) string {
	switch {
	case math.IsInf(f, 1):
		return "inf"
	case math.IsInf(f, -1):
		return "-inf"
	}
	e := strconv.FormatFloat(f, 'e', -1, 64) // [-]d[.ddd]e±XX, the shortest digits that round-trip
	sign := ""
	if e[0] == '-' {
		sign, e = "-", e[1:]
	}
	mant, exps, _ := strings.Cut(e, "e")
	digits := strings.Replace(mant, ".", "", 1)
	x, _ := strconv.Atoi(exps)
	decpt := x + 1
	if decpt <= -4 || decpt > 16 {
		m := digits[:1]
		if len(digits) > 1 {
			m += "." + digits[1:]
		}
		es := "+"
		if x < 0 {
			es, x = "-", -x
		}
		xs := strconv.Itoa(x)
		if len(xs) < 2 {
			xs = "0" + xs
		}
		return sign + m + "e" + es + xs
	}
	switch {
	case decpt <= 0:
		return sign + "0." + strings.Repeat("0", -decpt) + digits
	case decpt >= len(digits):
		return sign + digits + strings.Repeat("0", decpt-len(digits)) + ".0"
	}
	return sign + digits[:decpt] + "." + digits[decpt:]
}

type jsonFail struct{}

// flattenJSON is json.loads(b, object_pairs_hook=..) + parse._flatten: every leaf as (path, Python's str of it), every
// duplicate key kept, in order. ok false: not JSON for Python (judged as one raw value).
func flattenJSON(b string) (out []param, ok bool) {
	defer func() {
		if r := recover(); r != nil {
			if _, is := r.(jsonFail); !is {
				panic(r)
			}
			out, ok = nil, false
		}
	}()
	dec := json.NewDecoder(strings.NewReader(b))
	dec.UseNumber()
	next := func() json.Token {
		t, err := dec.Token()
		if err != nil {
			panic(jsonFail{})
		}
		return t
	}
	var value func(t json.Token, prefix string, depth int)
	value = func(t json.Token, prefix string, depth int) {
		if depth > parseMaxDepth {
			panic(jsonFail{})
		}
		switch x := t.(type) {
		case json.Delim:
			switch x {
			case '{':
				for {
					kt := next()
					if d, is := kt.(json.Delim); is && d == '}' {
						return
					}
					k, is := kt.(string)
					if !is {
						panic(jsonFail{})
					}
					p := k
					if prefix != "" {
						p = prefix + "." + k
					}
					value(next(), p, depth+1)
				}
			case '[':
				for i := 0; ; i++ {
					vt := next()
					if d, is := vt.(json.Delim); is && d == ']' {
						return
					}
					value(vt, prefix+"["+strconv.Itoa(i)+"]", depth+1)
				}
			}
			panic(jsonFail{})
		case string:
			out = append(out, param{"b", prefix, x})
		case json.Number:
			s := string(x)
			if strings.ContainsAny(s, ".eE") {
				f, _ := strconv.ParseFloat(s, 64) // the decoder checked the syntax; out of range is ±Inf or 0, as float()
				out = append(out, param{"b", prefix, pyFloat(f)})
			} else {
				digits := strings.TrimPrefix(s, "-")
				if len(digits) > 4300 { // Python's own int limit (parse._int)
					panic(jsonFail{})
				}
				if strings.Trim(digits, "0") == "" {
					s = "0"
				}
				out = append(out, param{"b", prefix, s})
			}
		case bool:
			if x {
				out = append(out, param{"b", prefix, "True"})
			} else {
				out = append(out, param{"b", prefix, "False"})
			}
		case nil:
			out = append(out, param{"b", prefix, ""})
		default:
			panic(jsonFail{})
		}
	}
	value(next(), "", 0)
	// json.loads: nothing but JSON whitespace may follow the value
	if strings.Trim(b[int(dec.InputOffset()):], " \t\n\r") != "" {
		return nil, false
	}
	return out, true
}

// multipart is parse._multipart.
func multipart(body, ct string) ([]param, bool) {
	m := reBound.FindStringSubmatch(ct)
	if m == nil {
		return nil, false
	}
	bnd := m[1]
	if bnd == "" {
		bnd = m[2]
	}
	pieces := strings.Split(body, "--"+bnd)
	if len(pieces) < 2 {
		return nil, false
	}
	out := []param{}
	for _, piece := range pieces[1:] {
		if strings.HasPrefix(piece, "--") {
			break
		}
		if strings.HasPrefix(piece, "\r\n") {
			piece = piece[2:]
		} else if strings.HasPrefix(piece, "\n") {
			piece = piece[1:]
		}
		i, j := strings.Index(piece, "\r\n\r\n"), strings.Index(piece, "\n\n")
		if i < 0 && j < 0 {
			out = append(out, param{"b", "", piece})
			continue
		}
		cut, skip := j, 2
		if j < 0 || 0 <= i && i < j {
			cut, skip = i, 4
		}
		head, content := piece[:cut], piece[cut+skip:]
		if strings.HasSuffix(content, "\r\n") {
			content = content[:len(content)-2]
		} else if strings.HasSuffix(content, "\n") {
			content = content[:len(content)-1]
		}
		disp := ""
		for _, ln := range strings.Split(head, "\n") {
			hk, hv, sep := strings.Cut(ln, ":")
			if sep && strings.ToLower(pyStrip(hk)) == "content-disposition" {
				disp = pyStrip(hv)
			}
		}
		name := ""
		if nm := reDispNm.FindStringSubmatchIndex(disp); nm != nil {
			if nm[2] >= 0 {
				name = disp[nm[2]:nm[3]]
			} else if nm[4] >= 0 {
				name = disp[nm[4]:nm[5]]
			}
		}
		if fn := reDispFn.FindStringSubmatchIndex(disp); fn != nil {
			f := ""
			if fn[2] >= 0 {
				f = disp[fn[2]:fn[3]]
			} else if fn[4] >= 0 {
				f = disp[fn[4]:fn[5]]
			}
			out = append(out, param{"b", name, f}, param{"b", name, runeCut(content, parseMaxFile)})
		} else {
			out = append(out, param{"b", name, content})
		}
	}
	return out, true
}

// pyRequest is the dibaguard.http.parse.Request a request becomes: lower-cased header names, every value of a
// repeated header joined (cookies with "; ", others with ", "), the raw path and query.
type pyRequest struct {
	method, path, query, body string
	headers                   map[string]string
}

func newPyRequest(method string, u *url.URL, hdr http.Header, body []byte) pyRequest {
	// the path as it came on the wire: RawPath whenever Go's own escaping would differ (EscapedPath re-escapes a raw
	// "{" that the client sent unescaped, which would add a decoding layer the Python edge never saw)
	path := u.RawPath
	if path == "" {
		path = u.EscapedPath()
	}
	r := pyRequest{method: strings.ToUpper(method), path: validUTF8([]byte(path)), query: validUTF8([]byte(u.RawQuery)),
		body: validUTF8(body), headers: map[string]string{}}
	if r.path == "" {
		r.path = "/"
	}
	for k, vs := range hdr {
		name := strings.ToLower(k)
		sep := ", "
		if name == "cookie" {
			sep = "; "
		}
		v := validUTF8([]byte(strings.Join(vs, sep)))
		if old, ok := r.headers[name]; ok { // two spellings of one name: joined like a repeat
			v = old + sep + v
		}
		r.headers[name] = v
	}
	return r
}

func (r pyRequest) params() []param {
	out := pairs(r.query, "q", nil)
	ctRaw := r.headers["content-type"]
	ct := strings.ToLower(ctRaw)
	b := pyStrip(r.body)
	if b != "" {
		switch {
		case strings.Contains(ct, "json") || b[0] == '[' || b[0] == '{':
			if js, ok := flattenJSON(b); ok {
				out = append(out, js...)
			} else {
				out = append(out, param{"b", "", b})
			}
		case strings.Contains(ct, "multipart/"):
			if mp, ok := multipart(b, ctRaw); ok {
				out = append(out, mp...)
			} else {
				out = append(out, param{"b", "", b})
			}
		case strings.Contains(b, "="):
			out = pairs(b, "b", out)
		default:
			out = append(out, param{"b", "", b})
		}
	}
	if ck := r.headers["cookie"]; ck != "" {
		for _, p := range strings.Split(ck, ";") {
			if pyStrip(p) == "" {
				continue
			}
			k, v, _ := strings.Cut(p, "=")
			out = append(out, param{"c", pyStrip(k), v})
		}
	}
	return out
}

// modelText is parse.Request.model_text.
func (r pyRequest) modelText() (string, int) {
	t := &textBuilder{rawLeft: parseRestMax}
	p, pa, pok, d := canonText(r.path, false)
	t.d = d
	t.add(strings.ToLower(r.method)+" ", p)
	if pok {
		t.add("p:=", pa)
	}
	ps := r.params()
	for i, x := range ps {
		if i >= parseMaxFields {
			break
		}
		cv, va, vok, d := canonText(x.v, x.where != "c")
		ck, ka, kok, _ := canonText(x.k, false)
		if d > t.d {
			t.d = d
		}
		prefix := x.where + ":" + ck + "="
		t.add(prefix, cv)
		if vok {
			t.add(prefix, va)
		}
		if kok {
			t.add(x.where+":"+ka+"=", cv)
		}
		if runeLen(x.v) > canonMaxLen {
			t.addRaw(prefix, runeFrom(x.v, canonMaxLen-parseChunkOverlap))
		}
	}
	if len(ps) > parseMaxFields {
		var rest bytes.Buffer
		for i, x := range ps[parseMaxFields:] {
			if i > 0 {
				rest.WriteByte('&')
			}
			rest.WriteString(x.k + "=" + x.v)
		}
		t.addRaw("x:rest=", rest.String())
	}
	header := func(h, v string) {
		ch, ha, hok, d := canonText(v, false)
		if d > t.d {
			t.d = d
		}
		t.add("h:"+h+"=", ch)
		if hok {
			t.add("h:"+h+"=", ha)
		}
	}
	for _, h := range stdHeaders {
		if v := r.headers[h]; v != "" {
			header(h, v)
		}
	}
	names := make([]string, 0, len(r.headers))
	for h := range r.headers {
		names = append(names, h)
	}
	sort.Strings(names)
	for _, h := range names {
		v := r.headers[h]
		if v == "" || neverHeaders[h] || h == "user-agent" || h == "referer" || h == "x-forwarded-for" || h == "content-type" {
			continue
		}
		ch, ha, hok, _ := canonText(v, false)
		if hdrSuspicious.MatchString(ch) || hok && hdrSuspicious.MatchString(ha) {
			header(h, v)
		}
	}
	return t.text(), t.d
}