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_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))
