generate-mangled-pickles.py

#!/usr/bin/env python3

import pickle, io, collections, sys

x = collections.deque()
x.append(10)
x.append(3.14)
x.append(10000000000000)
x.append((1, 2, 3))
x.append([1, 2, 3])
x.append({1, 2, 3})
x.append({'1': 1, '2': 2, '3': 3})
x.append('ПАСЬЯНС')

for protocol in range(pickle.HIGHEST_PROTOCOL + 1):
    b = io.BytesIO()
    pickle.dump(x, b, protocol)
    b = b.getvalue()

    with open('{}.pickle'.format(protocol), 'wb') as f:
        f.write(b)

    if '--mangle' in sys.argv:
        with open('{}-garbage.pickle'.format(protocol), 'wb') as f:
            f.write(b)
            f.write(b'GARBAGE')
        if protocol >= 3:
            with open('{}-noproto.pickle'.format(protocol), 'wb') as f:
                f.write(b[2:])
            with open('{}-noproto-garbage.pickle'.format(protocol), 'wb') as f:
                f.write(b[2:])
                f.write(b'GARBAGE')
        if protocol >= 4:
            if b[2] != 0x95: # FRAME opcode
                raise AssertionError('expected protocol {} to use framing'.format(protocol))
            with open('{}-noframe.pickle'.format(protocol), 'wb') as f:
                f.write(b[:2])
                f.write(b[11:])
            with open('{}-noframe-garbage.pickle'.format(protocol), 'wb') as f:
                f.write(b[:2])
                f.write(b[11:])
                f.write(b'GARBAGE')
            with open('{}-noproto-noframe.pickle'.format(protocol), 'wb') as f:
                f.write(b[11:])
            with open('{}-noproto-noframe-garbage.pickle'.format(protocol), 'wb') as f:
                f.write(b[11:])
                f.write(b'GARBAGE')