def test_util_push_nested(self) -> None:
obj = {(None, True): [{-1: 3.14}, 'boo']}
for proto in range(DEFAULT_TEST_PROTO + 1):
with self.subTest(obj=obj, proto=proto):
pa = PickleAssembler(proto=proto)
pa.util_push(obj)
self.assertEqual(pickle.loads(pa.assemble()), obj)
obj = {(None, True): [{-1: 3.14}, (b'baz', 'boo')]}
for proto in range(3, DEFAULT_TEST_PROTO + 1):
with self.subTest(obj=obj, proto=proto):
pa = PickleAssembler(proto=proto)
pa.util_push(obj)
self.assertEqual(pickle.loads(pa.assemble()), obj)
def test_append_raw(self) -> None:
pa = PickleAssembler(proto=0)
pa.append_raw(b'foo')
self.assertEqual(pa.assemble(), b'foo.')
with self.assertRaisesRegex(TypeError,
re.escape('raw data must be bytes')):
pa.append_raw('string') # type: ignore[arg-type]
def test_exploit(self) -> None:
pa = PickleAssembler(proto=0)
pa.push_mark()
pa.util_push(
'__import__("subprocess").check_output("echo hacked", shell=True, universal_newlines=True)'
)
pa.build_inst('builtins', 'eval')
payload = pa.assemble()
self.assertNotIn(b'R', payload)
self.assertEqual(pickle.loads(payload), 'hacked\n')
def test_util_memo(self) -> None:
obj = (42, )
indices = [255, 256] # type: list[int]
for proto in range(DEFAULT_TEST_PROTO + 1):
for index in indices:
with self.subTest(index=index, proto=proto):
pa = PickleAssembler(proto=proto)
pa.util_push(obj)
pa.util_memo_put(index)
pa.pop()
pa.util_memo_get(index)
self.assertEqual(pickle.loads(pa.assemble()), obj)
def test_pop(self) -> None:
with self.subTest(test_case='pop'):
pa = PickleAssembler(proto=DEFAULT_TEST_PROTO)
pa.push_mark()
pa.push_binint(1)
pa.push_binint(2)
pa.push_binint(3)
pa.pop()
pa.build_tuple()
self.assertEqual(pickle.loads(pa.assemble()), (1, 2))
with self.subTest(test_case='pop_mark'):
pa = PickleAssembler(proto=DEFAULT_TEST_PROTO)
pa.push_mark()
pa.push_binint(1)
pa.push_mark()
pa.push_binint(2)
pa.push_binint(3)
pa.pop_mark()
pa.build_tuple()
self.assertEqual(pickle.loads(pa.assemble()), (1, ))
def test_string_encoding(self) -> None:
test_cases = [
('push_string', '\xcc', 'latin-1'),
('push_binstring', '\xcc', 'latin-1'),
('push_binstring', '\u4e2d\u6587', 'utf-8'),
('push_short_binstring', '\xcc', 'latin-1'),
('push_short_binstring', '\u4e2d\u6587', 'utf-8'),
] # type: list[tuple[str, str, str]]
for test_case in test_cases:
function, string, encoding = test_case
with self.subTest(test_case=test_case):
pa = PickleAssembler(proto=DEFAULT_TEST_PROTO)
getattr(pa, function)(string, encoding)
self.assertEqual(
pickle.loads(pa.assemble(), encoding=encoding), string)
def test_push_0(self) -> None:
test_cases = [
('push_none', None),
('push_false', False),
('push_true', True),
('push_empty_tuple', ()),
('push_empty_list', []),
('push_empty_dict', {}),
('push_empty_set', set()),
] # type: list[tuple[str, object]]
for test_case in test_cases:
function, expected_result = test_case
with self.subTest(test_case=test_case):
pa = PickleAssembler(proto=DEFAULT_TEST_PROTO)
getattr(pa, function)()
self.assertEqual(pickle.loads(pa.assemble()), expected_result)
def test_push_1(self) -> None:
test_cases = [
('push_int', 66),
('push_int', False),
('push_int', True),
('push_binint', 2**31 - 1),
('push_binint', -2**31),
('push_binint1', 0),
('push_binint1', 2**8 - 1),
('push_binint2', 0),
('push_binint2', 2**16 - 1),
('push_long', 0xdeadbeefcafebabe),
('push_long1', 2**2039 - 1),
('push_long1', -2**2039),
('push_long4', -77),
('push_long4', 0xdeadbeefcafebabe),
('push_float', 3.14),
('push_binfloat', -1e1000),
('push_string', 'hello "world"\n'),
('push_binstring', 'hello "world"\n'),
('push_short_binstring', 'A' * 255),
('push_binbytes', b'\xcc\xdd'),
('push_binbytes8', b'\xcc\xdd'),
('push_short_binbytes', b'\xcc' * 255),
('push_unicode', 'hello "world"\n'),
('push_unicode', '\u4e2d\u6587'),
('push_binunicode', 'hello "world"\n'),
('push_binunicode', '\u4e2d\u6587'),
('push_binunicode8', 'hello "world"\n'),
('push_binunicode8', '\u4e2d\u6587'),
('push_short_binunicode', 'A' * 255),
('push_short_binunicode', '\u4e2d\u6587'),
] # type: list[tuple[str, object]]
if DEFAULT_TEST_PROTO >= 5: # pragma: no cover
test_cases.append(('push_bytearray8', bytearray(b'\xcc\xdd')))
for test_case in test_cases:
function, arg = test_case
with self.subTest(test_case=test_case):
pa = PickleAssembler(proto=DEFAULT_TEST_PROTO)
getattr(pa, function)(arg)
self.assertEqual(pickle.loads(pa.assemble()), arg)
def test_memo(self) -> None:
pa = PickleAssembler(proto=DEFAULT_TEST_PROTO)
pa.push_mark()
pa.push_binint(1)
pa.memo_put(66)
pa.push_binint(2)
pa.memo_binput(255)
pa.memo_binput(0)
pa.push_binint(3)
pa.memo_long_binput(0)
pa.memo_long_binput(777)
pa.push_binint(4)
pa.memo_memoize()
pa.memo_get(777)
pa.memo_binget(0)
pa.memo_binget(255)
pa.memo_long_binget(4)
pa.memo_binget(66)
pa.build_tuple()
self.assertEqual(pickle.loads(pa.assemble()),
(1, 2, 3, 4, 3, 3, 2, 4, 1))
def test_util_push(self) -> None:
test_cases = [
(None, 0, NONE + b'.'),
(True, 1, TRUE + b'.'),
(False, 2, PROTO + p8(2) + NEWFALSE + b'.'),
(1, 0, INT + b'1\n.'),
(255, 1, BININT1 + p8(255) + b'.'),
(256, 1, BININT2 + p16(256) + b'.'),
(65535, 1, BININT2 + p16(65535) + b'.'),
(65536, 1, BININT + p32(65536, signed=True) + b'.'),
(-1, 1, BININT + p32(-1, signed=True) + b'.'),
(2**31 - 1, 1, BININT + p32(2**31 - 1, signed=True) + b'.'),
(2**31, 1, INT + b'2147483648\n.'),
(2**31, 2, PROTO + p8(2) + LONG1 + p8(5) +
pack(2**31, endian='<', signed=True) + b'.'),
(2**2039 - 1, 2, PROTO + p8(2) + LONG1 + p8(255) +
pack(2**2039 - 1, endian='<', signed=True) + b'.'),
(2**2039, 2, PROTO + p8(2) + LONG4 + p32(256) +
pack(2**2039, endian='<', signed=True) + b'.'),
(2.1, 0, FLOAT + b'2.1\n.'),
(2.1, 1, BINFLOAT + struct.pack('>d', 2.1) + b'.'),
(b'foo', 3, PROTO + p8(3) + SHORT_BINBYTES + b'\x03foo.'),
(b'x' * 255, 3,
PROTO + p8(3) + SHORT_BINBYTES + p8(255) + b'x' * 255 + b'.'),
(b'x' * 256, 3,
PROTO + p8(3) + BINBYTES + p32(256) + b'x' * 256 + b'.'),
('bar\n\u4e2d\u6587', 0,
UNICODE + br'bar\u000a\u4e2d\u6587' + b'\n.'),
('bar\n\u4e2d\u6587', 3, PROTO + p8(3) + BINUNICODE + p32(10) +
b'bar\n\xe4\xb8\xad\xe6\x96\x87.'),
('bar\n\u4e2d\u6587', 4, PROTO + p8(4) + SHORT_BINUNICODE +
p8(10) + b'bar\n\xe4\xb8\xad\xe6\x96\x87.'),
('x' * 256, 4,
PROTO + p8(4) + BINUNICODE + p32(256) + b'x' * 256 + b'.'),
((), 0, MARK + TUPLE + b'.'),
((), 1, EMPTY_TUPLE + b'.'),
((1, ), 1, MARK + BININT1 + p8(1) + TUPLE + b'.'),
((1, ), 2, PROTO + p8(2) + BININT1 + p8(1) + TUPLE1 + b'.'),
((1, 2), 2, PROTO + p8(2) + BININT1 + p8(1) + BININT1 + p8(2) +
TUPLE2 + b'.'),
((1, 2, 3), 2, PROTO + p8(2) + BININT1 + p8(1) + BININT1 + p8(2) +
BININT1 + p8(3) + TUPLE3 + b'.'),
((1, 2, 3,
4), 2, PROTO + p8(2) + MARK + BININT1 + p8(1) + BININT1 + p8(2) +
BININT1 + p8(3) + BININT1 + p8(4) + TUPLE + b'.'), # noqa: E501 # pylint: disable=line-too-long
([], 0, MARK + LIST + b'.'),
([], 1, EMPTY_LIST + b'.'),
([1,
2], 1, MARK + BININT1 + p8(1) + BININT1 + p8(2) + LIST + b'.'),
({}, 0, MARK + DICT + b'.'),
({}, 1, EMPTY_DICT + b'.'),
({
1: 2
}, 1, MARK + BININT1 + p8(1) + BININT1 + p8(2) + DICT + b'.'),
] # type: list[tuple[object, int, bytes]]
for test_case in test_cases:
arg, proto, expected_result = test_case
with self.subTest(test_case=test_case):
pa = PickleAssembler(proto=proto)
pa.util_push(arg)
result = pa.assemble()
self.assertEqual(result, expected_result)
self.assertEqual(pickle.loads(result), arg)
def test_build(self) -> None:
with self.subTest(test_case='build_tuple'):
pa = PickleAssembler(proto=DEFAULT_TEST_PROTO)
pa.push_mark()
pa.push_binint(1)
pa.push_binint(2)
pa.push_binint(3)
pa.build_tuple()
self.assertEqual(pickle.loads(pa.assemble()), (1, 2, 3))
with self.subTest(test_case='build_tuple1'):
pa = PickleAssembler(proto=DEFAULT_TEST_PROTO)
pa.push_binint(1)
pa.build_tuple1()
self.assertEqual(pickle.loads(pa.assemble()), (1, ))
with self.subTest(test_case='build_tuple2'):
pa = PickleAssembler(proto=DEFAULT_TEST_PROTO)
pa.push_binint(1)
pa.push_binint(2)
pa.build_tuple2()
self.assertEqual(pickle.loads(pa.assemble()), (1, 2))
with self.subTest(test_case='build_tuple3'):
pa = PickleAssembler(proto=DEFAULT_TEST_PROTO)
pa.push_binint(1)
pa.push_binint(2)
pa.push_binint(3)
pa.build_tuple3()
self.assertEqual(pickle.loads(pa.assemble()), (1, 2, 3))
with self.subTest(test_case='build_list'):
pa = PickleAssembler(proto=DEFAULT_TEST_PROTO)
pa.push_mark()
pa.push_binint(1)
pa.push_binint(2)
pa.push_binint(3)
pa.build_list()
self.assertEqual(pickle.loads(pa.assemble()), [1, 2, 3])
with self.subTest(test_case='build_dict'):
pa = PickleAssembler(proto=DEFAULT_TEST_PROTO)
pa.push_mark()
pa.push_binint(1)
pa.push_binint(2)
pa.push_binint(3)
pa.push_binint(4)
pa.build_dict()
self.assertEqual(pickle.loads(pa.assemble()), {1: 2, 3: 4})
with self.subTest(test_case='build_frozenset'):
pa = PickleAssembler(proto=DEFAULT_TEST_PROTO)
pa.push_mark()
pa.push_binint(1)
pa.push_binint(2)
pa.push_binint(3)
pa.build_frozenset()
result = pickle.loads(pa.assemble())
self.assertEqual(result, frozenset({1, 2, 3}))
self.assertIsInstance(result, frozenset)
with self.subTest(test_case='build_append'):
pa = PickleAssembler(proto=DEFAULT_TEST_PROTO)
pa.push_mark()
pa.push_binint(1)
pa.push_binint(2)
pa.build_list()
pa.push_binint(3)
pa.build_append()
self.assertEqual(pickle.loads(pa.assemble()), [1, 2, 3])
with self.subTest(test_case='build_appends'):
pa = PickleAssembler(proto=DEFAULT_TEST_PROTO)
pa.push_mark()
pa.push_binint(1)
pa.build_list()
pa.push_mark()
pa.push_binint(2)
pa.push_binint(3)
pa.build_appends()
self.assertEqual(pickle.loads(pa.assemble()), [1, 2, 3])
with self.subTest(test_case='build_setitem'):
pa = PickleAssembler(proto=DEFAULT_TEST_PROTO)
pa.push_mark()
pa.push_binint(1)
pa.push_binint(2)
pa.build_dict()
pa.push_binint(3)
pa.push_binint(4)
pa.build_setitem()
self.assertEqual(pickle.loads(pa.assemble()), {1: 2, 3: 4})
with self.subTest(test_case='build_setitems'):
pa = PickleAssembler(proto=DEFAULT_TEST_PROTO)
pa.push_mark()
pa.push_binint(1)
pa.push_binint(2)
pa.build_dict()
pa.push_mark()
pa.push_binint(3)
pa.push_binint(4)
pa.push_binint(1)
pa.push_binint(5)
pa.build_setitems()
self.assertEqual(pickle.loads(pa.assemble()), {1: 5, 3: 4})
with self.subTest(test_case='build_additems'):
pa = PickleAssembler(proto=DEFAULT_TEST_PROTO)
pa.push_empty_set()
pa.push_mark()
pa.push_binint(1)
pa.push_binint(2)
pa.push_binint(3)
pa.build_additems()
result = pickle.loads(pa.assemble())
self.assertEqual(result, {1, 2, 3})
self.assertIsInstance(result, set)
with self.subTest(test_case='build_inst'):
pa = PickleAssembler(proto=DEFAULT_TEST_PROTO)
pa.push_mark()
pa.push_binunicode('foo')
pa.build_inst('__main__', 'SampleClass')
self.assertEqual(pickle.loads(pa.assemble()), SampleClass('foo'))
with self.subTest(test_case='build_obj'):
pa = PickleAssembler(proto=DEFAULT_TEST_PROTO)
pa.push_mark()
pa.push_global('__main__', 'SampleClass')
pa.push_binunicode('foo')
pa.build_obj()
self.assertEqual(pickle.loads(pa.assemble()), SampleClass('foo'))
with self.subTest(test_case='build_newobj'):
pa = PickleAssembler(proto=DEFAULT_TEST_PROTO)
pa.push_global('__main__', 'SampleClass')
pa.push_binunicode('foo')
pa.build_tuple1()
pa.build_newobj()
self.assertEqual(pickle.loads(pa.assemble()), SampleClass('foo'))
with self.subTest(test_case='build_newobj_ex'):
pa = PickleAssembler(proto=DEFAULT_TEST_PROTO)
pa.push_global('__main__', 'SampleClass')
pa.push_binunicode('foo')
pa.build_tuple1()
pa.push_mark()
pa.push_binunicode('attr2')
pa.push_binunicode('bar')
pa.build_dict()
pa.build_newobj_ex()
self.assertEqual(pickle.loads(pa.assemble()),
SampleClass('foo', attr2='bar'))
with self.subTest(test_case='build_stack_global'):
pa = PickleAssembler(proto=DEFAULT_TEST_PROTO)
pa.push_binunicode('__main__')
pa.push_binunicode('SampleClass')
pa.build_stack_global()
self.assertIs(pickle.loads(pa.assemble()), SampleClass)
with self.subTest(test_case='build_reduce'):
pa = PickleAssembler(proto=DEFAULT_TEST_PROTO)
pa.push_global('__main__', 'SampleClass')
pa.push_binunicode('foo')
pa.build_tuple1()
pa.build_reduce()
self.assertEqual(pickle.loads(pa.assemble()), SampleClass('foo'))
with self.subTest(test_case='build_build'):
pa = PickleAssembler(proto=DEFAULT_TEST_PROTO)
pa.push_global('__main__', 'SampleClass')
pa.push_empty_tuple()
pa.build_reduce()
pa.push_mark()
pa.push_binunicode('attr1')
pa.push_binunicode('foo')
pa.build_dict()
pa.build_build()
self.assertEqual(pickle.loads(pa.assemble()), SampleClass('foo'))
with self.subTest(test_case='build_dup'):
pa = PickleAssembler(proto=DEFAULT_TEST_PROTO)
pa.push_mark()
pa.push_binint(1)
pa.build_dup()
pa.build_tuple()
self.assertEqual(pickle.loads(pa.assemble()), (1, 1))
def test_push_global(self) -> None:
pa = PickleAssembler(proto=DEFAULT_TEST_PROTO)
pa.push_global('__main__', 'SampleClass')
self.assertIs(pickle.loads(pa.assemble()), SampleClass)
def test_assemble(self) -> None:
pa = PickleAssembler(proto=0)
pa.push_none()
self.assertEqual(pa.assemble(), NONE + b'.')