Add pickle package for generating exploit payloads (#587)

* feat: add pickle package for generating exploit payloads

Author: vlobstein-vc

pickle/disasm.go

@@ -0,0 +1,237 @@
+package pickle
+
+import (
+	"bytes"
+	"errors"
+	"fmt"
+	"io"
+)
+
+// Op is a decoded pickle opcode. Arg is nil for no-arg opcodes; a
+// string for BINUNICODE / SHORT_BINUNICODE / GLOBAL; a []byte for
+// BINBYTES / SHORT_BINBYTES; an int for BININTx / PROTO / FRAME.
+// GLOBAL packs (module, attr) into a [2]string.
+type Op struct {
+	Name string
+	Code byte
+	Arg  any
+}
+
+// String formats an Op the way pickletools.dis does.
+func (o Op) String() string {
+	if o.Arg == nil {
+		return o.Name
+	}
+	switch v := o.Arg.(type) {
+	case string:
+		return fmt.Sprintf("%s %q", o.Name, v)
+	case [2]string:
+		return fmt.Sprintf("%s %q %q", o.Name, v[0], v[1])
+	case []byte:
+		return fmt.Sprintf("%s %q", o.Name, v)
+	case int:
+		return fmt.Sprintf("%s %d", o.Name, v)
+	default:
+		return fmt.Sprintf("%s %v", o.Name, v)
+	}
+}
+
+// Sentinel errors for errors.Is matching.
+var (
+	ErrUnknownOpcode = errors.New("pickle: unknown opcode")
+	ErrTruncated     = errors.New("pickle: truncated stream")
+)
+
+// noArgOps maps opcodes that take no argument to their symbolic name.
+var noArgOps = map[byte]string{
+	OpStop:       "STOP",
+	OpMark:       "MARK",
+	OpReduce:     "REDUCE",
+	OpBuild:      "BUILD",
+	OpEmptyTuple: "EMPTY_TUPLE",
+	OpTuple:      "TUPLE",
+	OpTuple1:     "TUPLE1",
+	OpTuple2:     "TUPLE2",
+	OpTuple3:     "TUPLE3",
+	OpEmptyList:  "EMPTY_LIST",
+	OpEmptyDict:  "EMPTY_DICT",
+	OpAppend:     "APPEND",
+	OpAppends:    "APPENDS",
+	OpSetItem:    "SETITEM",
+	OpSetItems:   "SETITEMS",
+	OpNewTrue:    "NEWTRUE",
+	OpNewFalse:   "NEWFALSE",
+	OpNone:       "NONE",
+	OpMemoize:    "MEMOIZE",
+	OpDup:        "DUP",
+	OpPop:        "POP",
+	OpPopMark:    "POP_MARK",
+}
+
+// Disassemble walks a pickle stream into a list of named opcodes.
+// Returns an error on truncated input or unknown opcodes rather than
+// silently skipping.
+func Disassemble(data []byte) ([]Op, error) {
+	r := bytes.NewReader(data)
+	ops := []Op{}
+
+	for r.Len() > 0 {
+		code, err := r.ReadByte()
+		if err != nil {
+			return ops, fmt.Errorf("read opcode: %w", err)
+		}
+		op, err := decodeOp(code, r)
+		if err != nil {
+			return ops, fmt.Errorf("at byte %d (opcode 0x%02x): %w",
+				len(data)-r.Len()-1, code, err)
+		}
+		ops = append(ops, op)
+		if code == OpStop {
+			break
+		}
+	}
+
+	return ops, nil
+}
+
+// argSpec describes how to decode an opcode's argument. read pulls the
+// arg value from r and returns it ready to drop into Op.Arg. Centralising
+// the (name, reader) pair as data lets decodeOp stay a one-liner lookup.
+type argSpec struct {
+	name string
+	read func(*bytes.Reader) (any, error)
+}
+
+// argOps maps every opcode that carries an argument to its spec. The
+// no-arg opcodes live in noArgOps above.
+var argOps = map[byte]argSpec{
+	OpProto:           {"PROTO", readUintArg(1)},
+	OpBinPut:          {"BINPUT", readUintArg(1)},
+	OpBinGet:          {"BINGET", readUintArg(1)},
+	OpBinInt1:         {"BININT1", readUintArg(1)},
+	OpBinInt2:         {"BININT2", readUintArg(2)},
+	OpBinInt:          {"BININT", readSignedInt32Arg},
+	OpFrame:           {"FRAME", readUintArg(8)},
+	OpGlobal:          {"GLOBAL", readGlobalArg},
+	OpBinUnicode:      {"BINUNICODE", readStringArg(4)},
+	OpShortBinUnicode: {"SHORT_BINUNICODE", readStringArg(1)},
+	OpBinUnicode8:     {"BINUNICODE8", readStringArg(8)},
+	OpBinBytes:        {"BINBYTES", readBytesArg(4)},
+	OpShortBinBytes:   {"SHORT_BINBYTES", readBytesArg(1)},
+}
+
+// decodeOp dispatches an opcode to its argument reader.
+func decodeOp(code byte, r *bytes.Reader) (Op, error) {
+	if name, ok := noArgOps[code]; ok {
+		return Op{Name: name, Code: code}, nil
+	}
+	spec, ok := argOps[code]
+	if !ok {
+		return Op{}, fmt.Errorf("%w: 0x%02x", ErrUnknownOpcode, code)
+	}
+	arg, err := spec.read(r)
+	if err != nil {
+		return Op{}, fmt.Errorf("%w reading %s arg: %w", ErrTruncated, spec.name, err)
+	}
+
+	return Op{Name: spec.name, Code: code, Arg: arg}, nil
+}
+
+// readUintArg returns a reader that pulls byteCount little-endian bytes
+// as an unsigned int.
+func readUintArg(byteCount int) func(*bytes.Reader) (any, error) {
+	return func(r *bytes.Reader) (any, error) {
+		u, err := readUint(r, byteCount)
+
+		return int(u), err
+	}
+}
+
+// readSignedInt32Arg pulls 4 LE bytes as a signed int (BININT).
+func readSignedInt32Arg(r *bytes.Reader) (any, error) {
+	u, err := readUint(r, 4)
+
+	return int(int32(u)), err
+}
+
+// readStringArg / readBytesArg share readPrefixedBytes; they differ only
+// in how the payload is typed.
+func readStringArg(lenBytes int) func(*bytes.Reader) (any, error) {
+	return func(r *bytes.Reader) (any, error) {
+		body, err := readPrefixedBytes(r, lenBytes)
+
+		return string(body), err
+	}
+}
+
+func readBytesArg(lenBytes int) func(*bytes.Reader) (any, error) {
+	return func(r *bytes.Reader) (any, error) {
+		body, err := readPrefixedBytes(r, lenBytes)
+
+		return body, err
+	}
+}
+
+// readGlobalArg pulls (module, attr) as two newline-terminated strings.
+func readGlobalArg(r *bytes.Reader) (any, error) {
+	mod, err := readLine(r)
+	if err != nil {
+		return nil, err
+	}
+	attr, err := readLine(r)
+	if err != nil {
+		return nil, err
+	}
+
+	return [2]string{mod, attr}, nil
+}
+
+// readUint reads byteCount little-endian bytes as a uint64. byteCount
+// must be 1, 2, 4 or 8 (the widths used by every pickle integer /
+// length opcode).
+func readUint(r *bytes.Reader, byteCount int) (uint64, error) {
+	if byteCount < 1 || byteCount > 8 {
+		return 0, fmt.Errorf("%w: invalid byte count %d", ErrTruncated, byteCount)
+	}
+	buf := make([]byte, byteCount)
+	if _, err := io.ReadFull(r, buf); err != nil {
+		return 0, fmt.Errorf("read %d bytes: %w", byteCount, err)
+	}
+	var v uint64
+	for i := byteCount - 1; i >= 0; i-- {
+		v = v<<8 | uint64(buf[i])
+	}
+
+	return v, nil
+}
+
+func readLine(r *bytes.Reader) (string, error) {
+	var buf bytes.Buffer
+	for {
+		b, err := r.ReadByte()
+		if err != nil {
+			return "", fmt.Errorf("read line byte: %w", err)
+		}
+		if b == '\n' {
+			return buf.String(), nil
+		}
+		buf.WriteByte(b)
+	}
+}
+
+// readPrefixedBytes reads a 1/4/8-byte LE length then that many bytes.
+func readPrefixedBytes(r *bytes.Reader, lenBytes int) ([]byte, error) {
+	n, err := readUint(r, lenBytes)
+	if err != nil {
+		return nil, fmt.Errorf("read length: %w", err)
+	}
+	if n > uint64(r.Len()) {
+		return nil, fmt.Errorf("%w: length %d exceeds remaining %d", ErrTruncated, n, r.Len())
+	}
+	buf := make([]byte, n)
+	if _, err := io.ReadFull(r, buf); err != nil {
+		return nil, fmt.Errorf("read body: %w", err)
+	}
+
+	return buf, nil
+}

pickle/fragments.go

@@ -0,0 +1,160 @@
+package pickle
+
+// Fragment is a pickle byte slice that pushes exactly one value onto
+// the unpickler stack. Primitives, composites and calls all return
+// Fragments, so they nest. Dump wraps a top-level Fragment with PROTO
+// + STOP into a complete stream.
+type Fragment []byte
+
+// Dump wraps top with PROTO + STOP into a complete pickle stream.
+func Dump(top Fragment) []byte {
+	out := make([]byte, 0, 2+len(top)+1)
+	out = append(out, OpProto, Protocol)
+	out = append(out, top...)
+	out = append(out, OpStop)
+
+	return out
+}
+
+// concatFragments concatenates parts into a single pre-sized Fragment.
+func concatFragments(parts ...Fragment) Fragment {
+	total := 0
+	for _, p := range parts {
+		total += len(p)
+	}
+	out := make(Fragment, 0, total)
+	for _, p := range parts {
+		out = append(out, p...)
+	}
+
+	return out
+}
+
+// Str pushes the Python string s onto the stack. Returns ok=false if s
+// exceeds the 4-byte length cap baked into BINUNICODE.
+func Str(s string) (Fragment, bool) {
+	out, ok := BinUnicode(s)
+	if !ok {
+		return nil, false
+	}
+
+	return Fragment(out), true
+}
+
+// Int pushes a Python int via the smallest BININTx opcode that fits.
+// Wider values (LONG1 / LONG4) are not implemented; use the int32
+// signature so callers convert at the call site rather than at runtime.
+func Int(n int32) Fragment {
+	switch {
+	case n >= 0 && n <= 0xff:
+		return Fragment(BinInt1(uint8(n)))
+	case n >= 0 && n <= 0xffff:
+		return Fragment(BinInt2(uint16(n)))
+	default:
+		return Fragment(BinInt(n))
+	}
+}
+
+// Bool pushes True or False onto the stack.
+func Bool(b bool) Fragment {
+	if b {
+		return Fragment{OpNewTrue}
+	}
+
+	return Fragment{OpNewFalse}
+}
+
+// None pushes Python None.
+func None() Fragment { return Fragment{OpNone} }
+
+// tupleSmallClosers indexes the protocol-2 fast-path opcodes by tuple
+// arity: EMPTY_TUPLE for 0, TUPLE1..TUPLE3 for 1..3.
+var tupleSmallClosers = [...]byte{OpEmptyTuple, OpTuple1, OpTuple2, OpTuple3}
+
+// TupleOf pushes a Python tuple. Arities 0..3 use the dedicated
+// EMPTY_TUPLE / TUPLE1..3 fast-path opcodes; larger tuples fall back
+// to MARK + items + TUPLE.
+func TupleOf(elems ...Fragment) Fragment {
+	if n := len(elems); n < len(tupleSmallClosers) {
+		parts := make([]Fragment, 0, n+1)
+		parts = append(parts, elems...)
+		parts = append(parts, Fragment{tupleSmallClosers[n]})
+
+		return concatFragments(parts...)
+	}
+	parts := make([]Fragment, 0, 2+len(elems))
+	parts = append(parts, Fragment{OpMark})
+	parts = append(parts, elems...)
+	parts = append(parts, Fragment{OpTuple})
+
+	return concatFragments(parts...)
+}
+
+// collection emits an EMPTY_<X> opener and, if items is non-empty,
+// wraps the items in MARK + items + closeOp. Shared body of ListOf
+// and DictOf, since the only thing that differs between them is the
+// pair of opcodes and how items are flattened from the input shape.
+func collection(emptyOp, closeOp byte, items []Fragment) Fragment {
+	if len(items) == 0 {
+		return Fragment{emptyOp}
+	}
+	parts := make([]Fragment, 0, 3+len(items))
+	parts = append(parts, Fragment{emptyOp}, Fragment{OpMark})
+	parts = append(parts, items...)
+	parts = append(parts, Fragment{closeOp})
+
+	return concatFragments(parts...)
+}
+
+// ListOf pushes a Python list. Wire form: EMPTY_LIST [+ MARK + items
+// + APPENDS].
+func ListOf(elems ...Fragment) Fragment {
+	return collection(OpEmptyList, OpAppends, elems)
+}
+
+// DictEntry is a key/value pair for DictOf. Both sides are Fragments
+// so they can themselves be calls or other composites.
+type DictEntry struct {
+	Key   Fragment
+	Value Fragment
+}
+
+// DictOf pushes a Python dict. Wire form: EMPTY_DICT [+ MARK + pairs
+// + SETITEMS].
+func DictOf(entries ...DictEntry) Fragment {
+	items := make([]Fragment, 0, 2*len(entries))
+	for _, e := range entries {
+		items = append(items, e.Key, e.Value)
+	}
+
+	return collection(OpEmptyDict, OpSetItems, items)
+}
+
+// Call pushes module.attr(*args) using the legacy GLOBAL opcode.
+func Call(module, attr string, args ...Fragment) Fragment {
+	return CallFragment(Fragment(Global(module, attr)), args...)
+}
+
+// CallFragment pushes callable(*args) for a callable already on the
+// stack (from a previous Call, Method, attribute access, etc.).
+func CallFragment(callable Fragment, args ...Fragment) Fragment {
+	return concatFragments(callable, TupleOf(args...), Fragment{OpReduce})
+}
+
+// Method pushes obj.name(*args) via getattr(obj, name)(*args); pickle
+// has no dedicated method-call opcode. Returns ok=false if name cannot
+// be encoded as a BINUNICODE string (delegated to Str).
+func Method(obj Fragment, name string, args ...Fragment) (Fragment, bool) {
+	nameFrag, ok := Str(name)
+	if !ok {
+		return nil, false
+	}
+
+	return CallFragment(Call("builtins", "getattr", obj, nameFrag), args...), true
+}
+
+// Build pushes obj with state applied via __setstate__. Wire form:
+// obj + state + BUILD.
+func Build(obj Fragment, state Fragment) Fragment {
+	return concatFragments(obj, state, Fragment{OpBuild})
+}