package winpwn

import (
	"bytes"
	"fmt"
)

// FindXORKey finds a single byte k such that XORing data with the
// repeating key k produces no byte present in badChars -- the search step
// behind pwntools' encoders.xor, scoped to a single-byte key (the common
// case: avoiding \x00/\x0a/\x0d in a payload that itself gets typed/parsed
// as text before reaching the target). Returns an error if no byte in
// 0..255 works, which can happen if badChars is large enough that every
// possible key XORs at least one data byte into a forbidden value.
func FindXORKey(data []byte, badChars []byte) (byte, error) {
	var bad [256]bool
	for _, b := range badChars {
		bad[b] = true
	}

candidate:
	for k := 0; k < 256; k++ {
		key := byte(k)
		if bad[key] {
			continue // the key itself ends up nowhere in the output, but
			// disallowing it too keeps the result usable as a literal
			// byte elsewhere in the same payload (e.g. a decoder stub
			// that embeds the key as an immediate).
		}
		for _, b := range data {
			if bad[b^key] {
				continue candidate
			}
		}
		return key, nil
	}
	return 0, fmt.Errorf("no single-byte XOR key avoids all %d bad chars for this %d-byte input", len(badChars), len(data))
}

// EncodeXOR finds a single-byte XOR key avoiding badChars (via FindXORKey)
// and returns data encoded with it, plus the key itself. Decode by XORing
// again with the same key (see Xor in fiddling.go) -- this is the
// data-level half of pwntools' bad-character avoidance; it does not emit a
// self-decoding stub, so the receiving side needs to already know how to
// undo it (e.g. the target's own code does the XOR, or your script decodes
// a leaked buffer before parsing it).
func EncodeXOR(data []byte, badChars []byte) (encoded []byte, key byte, err error) {
	key, err = FindXORKey(data, badChars)
	if err != nil {
		return nil, 0, err
	}
	return Xor(data, []byte{key}), key, nil
}

// HasBadChars reports whether data contains any byte in badChars -- the
// quick check before bothering with an encoder at all.
func HasBadChars(data []byte, badChars []byte) bool {
	return bytes.ContainsAny(data, string(badChars))
}