def test_rope_dump():
r = rope(b'abcd', b'efgh', b'ijkl')
r = rope(r[:5], r[9:], r[5:9])
fout = six.BytesIO()
grope.dump(r, fout)
assert fout.getvalue() == b'abcdejklfghi'
def to_blob(self, update_checksum=False):
self._opt_header.CheckSum = 0
self._opt_header.SizeOfImage = max(
self._mem_align(sec.hdr.VirtualAddress + sec.hdr.VirtualSize)
for sec in self._sections)
self._check_vm_overlaps()
header_end = (
len(self._dos_stub) + 4 + self._file_header.size + 2 +
self._opt_header.size +
len(self._data_directories) * _IMAGE_DATA_DIRECTORY.size +
len(self._sections) * _IMAGE_SECTION_HEADER.size)
section_offset = self._file_align(header_end)
header_pad = section_offset - header_end
for sec in self._sections:
if sec.hdr.PointerToRawData == 0:
continue
sec.hdr.PointerToRawData = section_offset
sec.hdr.SizeOfRawData = self._file_align(len(sec.data))
section_offset = section_offset + sec.hdr.SizeOfRawData
new_file = []
new_file.append(self._dos_stub)
new_file.append(b'PE\0\0')
new_file.append(self._file_header.pack())
new_file.append(struct.pack('<H', self._opt_header.sig))
new_file.append(self._opt_header.pack())
for dd in self._data_directories:
new_file.append(dd.pack())
for sec in self._sections:
new_file.append(sec.hdr.pack())
new_file.append(b'\0' * header_pad)
for sec in self._sections:
if sec.data is None:
continue
new_file.append(sec.data)
with_pad = self._file_align(len(sec.data))
pad = with_pad - len(sec.data)
if pad:
new_file.append(b'\0' * pad)
new_file.append(self._trailer)
out_blob = rope(*new_file)
if update_checksum:
new_checksum = pe_checksum(out_blob)
else:
new_checksum = 0
return rope(out_blob[:self._checksum_offs],
struct.pack('<I', new_checksum),
out_blob[self._checksum_offs + 4:])
def test_blob_io_swap_bytes():
fin = six.BytesIO(b'abcdefgh' * 2**18)
blob = wrap_io(fin)
blob = rope(blob[:4], b'x', blob[6:])
assert bytes(blob[:7]) == b'abcdxgh'
def save(self):
log_num_tables = int(math.floor(math.log2(len(self._tables))))
hdr = _OTF_OFFSET_TABLE(
sfntVersion=0x10000,
numTables=len(self._tables),
searchRange=2**log_num_tables * 16,
entrySelector=log_num_tables,
rangeShift=(len(self._tables) - 2**log_num_tables) * 16)
table_blobs = [tab.pack() for tab in self._tables]
table_hdrs = []
blobs_with_pad = []
head_offset = None
data_offset = _OTF_OFFSET_TABLE.size + _OTF_TABLE_RECORD.size * len(
self._tables)
for tab, blob in zip(self._tables, table_blobs):
if tab.name == b'head':
head_offset = data_offset
table_hdrs.append(
(tab.name,
_OTF_TABLE_RECORD(tag=tab.name,
checkSum=_otf_table_checksum(blob),
offset=data_offset,
length=len(blob)).pack()))
blobs_with_pad.append(blob)
blobs_with_pad.append(_pad4(len(blob)))
data_offset += _align4(len(blob))
table_hdrs.sort()
pre_full_checksum = rope(hdr.pack(),
*(blob for name, blob in table_hdrs),
*blobs_with_pad)
full_checksum = _otf_table_checksum(pre_full_checksum)
head_blob = self._head.pack(checksum=(0x1b1b0afba - full_checksum)
& 0xffffffff)
return rope(pre_full_checksum[:head_offset], head_blob,
pre_full_checksum[head_offset + len(head_blob):])
def get_vm(self, start, stop):
for sec in self._sections:
if sec.hdr.VirtualAddress <= start and sec.hdr.VirtualAddress + sec.hdr.VirtualSize >= stop:
sec_offs = start - sec.hdr.VirtualAddress
init_size = min(sec.hdr.SizeOfRawData - sec_offs, stop - start)
uninit_size = stop - start - init_size
if len(sec.data) < sec_offs + init_size:
raise RuntimeError('PE file corrupt: missing section content')
return rope(sec.data[sec_offs:sec_offs + init_size], b'\0'*uninit_size)
def get_vm(self, start, stop):
for sec in self._sections:
if sec.hdr.VirtualAddress <= start and sec.hdr.VirtualAddress + sec.hdr.VirtualSize >= stop:
sec_offs = start - sec.hdr.VirtualAddress
init_size = min(sec.hdr.SizeOfRawData - sec_offs, stop - start)
uninit_size = stop - start - init_size
return rope(sec.data[sec_offs:sec_offs + init_size],
b'\0' * uninit_size)
def pe_resources_prepack(rsrc):
entries = _prepack(rsrc)
strings = []
string_map = {}
def add_string(s):
r = string_map.get(s)
if r is None:
encoded = s.encode('utf-16le')
r = sum(len(ss) for ss in strings)
string_map[s] = r
strings.append(_STRING_HEADER(Length=len(encoded) // 2).pack())
strings.append(encoded)
return r
_entry_offsets = {}
offs = 0
for ent in entries:
_entry_offsets[ent] = offs
offs += ent.size
table_size = offs
for ent in entries:
if isinstance(ent, _RESOURCE_DIRECTORY_ENTRY):
if isinstance(ent.NameOrId, six.string_types):
ent.NameOrId = (1 << 31) | (table_size +
add_string(ent.NameOrId))
strings = b''.join(strings)
aligned_strings_len = align16(len(strings))
strings += b'\0' * (aligned_strings_len - len(strings))
data_offs = table_size + len(strings)
blobs = []
for ent in entries:
if isinstance(ent, _RESOURCE_DIRECTORY_ENTRY):
if isinstance(ent.Offset, _RESOURCE_DIRECTORY_TABLE):
ent.Offset = (1 << 31) | _entry_offsets[ent.Offset]
else:
ent.Offset = _entry_offsets[ent.Offset]
elif isinstance(ent, _RESOURCE_DATA_ENTRY):
blob = ent.DataRva
ent.DataRva = data_offs
blobs.append(blob)
aligned_size = align8(len(blob))
pad = aligned_size - len(blob)
if pad:
blobs.append(b'\0' * pad)
data_offs += aligned_size
return _PrepackedResources(entries, strings, rope(*blobs))
def _pack_node(node):
children = []
for child in node.children:
if children:
children.append(b'\0' *
(align4(len(children[-1])) - len(children[-1])))
children.append(_pack_node(child))
children = rope(*children)
name = node.name.encode('utf-16le') + b'\0\0'
children_offset = align4(_NODE_HEADER.size + len(name))
name_pad = b'\0' * (children_offset - _NODE_HEADER.size - len(name))
hdr = _NODE_HEADER()
if node.value is None:
value = b''
hdr.wValueLength = 0
hdr.wType = 1
elif isinstance(node.value, six.string_types):
value = node.value.encode('utf-16le') + b'\0\0'
hdr.wValueLength = len(value) // 2
hdr.wType = 1
else:
value = node.value
hdr.wValueLength = len(value)
hdr.wType = 0
if not children:
hdr.wLength = _NODE_HEADER.size + len(name) + len(name_pad) + len(
value)
value_pad = b''
else:
value_len_aligned = align4(len(value))
value_pad = b'\0' * (value_len_aligned - len(value))
hdr.wLength = _NODE_HEADER.size + len(name) + len(name_pad) + len(
value) + len(value_pad) + len(children)
return rope(hdr.pack(), name, name_pad, value, value_pad, children)
def test_indexing():
r = rope('abcd', 'efgh')
assert r[0] == 'a'
assert r[3] == 'd'
assert r[4] == 'e'
assert r[7] == 'h'
assert r[-1] == 'h'
assert r[-8] == 'a'
with pytest.raises(IndexError):
r[8]
with pytest.raises(IndexError):
r[-9]
def pack(self):
segments = []
seg_cid = None
seg_gid = None
for cid, gid in enumerate(self._map):
if seg_cid is not None and (gid == 0
or cid - seg_cid + seg_gid != gid):
segments.append((seg_cid, seg_gid, cid - seg_cid))
seg_cid = None
if seg_cid is None and gid != 0:
seg_cid = cid
seg_gid = gid
if seg_cid is not None:
segments.append((seg_cid, seg_gid, cid - seg_cid))
segments.append((0xffff, 0, 1))
seg_words = []
seg_words.extend(cid + length - 1 for cid, gid, length in segments)
seg_words.append(0)
seg_words.extend(cid for cid, gid, length in segments)
seg_words.extend((gid - cid) & 0xffff for cid, gid, length in segments)
seg_words.extend(0 for cid, gid, length in segments)
seg_spec = struct.pack('>{}H'.format(len(seg_words)), *seg_words)
hdr = _cmap_header(version=0, numTables=1).pack()
record = _cmap_encodingRecord(platformID=3,
encodingID=1,
offset=_cmap_header.size +
_cmap_encodingRecord.size).pack()
entry_selector = int(math.floor(math.log2(len(segments))))
search_range = 2 * (2**entry_selector)
subheader = _cmap_fmt4_header(
format=4,
length=len(seg_spec) + _cmap_fmt4_header.size,
language=0,
segCountX2=len(segments) * 2,
searchRange=search_range,
entrySelector=entry_selector,
rangeShift=len(segments) * 2 - search_range).pack()
return rope(hdr, record, subheader, seg_spec)
def test_empty_slice():
r = rope(_MockString(0, 1000))
for ch in r.chunks:
assert len(ch) != 0
rr = r[0:0]
for ch in rr.chunks:
assert len(ch) != 0
rr = r[999:999]
for ch in rr.chunks:
assert len(ch) != 0
rr = r[1000:1000]
for ch in rr.chunks:
assert len(ch) != 0
def pe_checksum(blob):
total_len = len(blob)
r = 0
while len(blob) >= 0x1000:
words = struct.unpack('<2048H', bytes(blob[:0x1000]))
r += sum(words)
blob = blob[0x1000:]
if len(blob) % 2 != 0:
blob = rope(blob, b'\0')
words = struct.unpack('<' + 'H' * (len(blob) // 2), bytes(blob))
r += sum(words)
while r > 0xffff:
c = r
r = 0
while c:
r += c & 0xffff
c >>= 16
return r + total_len
def test_idempotence():
r = rope('test')
assert rope(r)._tree is r._tree
def test_balancing():
r = rope(*string.ascii_lowercase)
assert str(r) == string.ascii_lowercase
def test_unicode_conversion():
r = rope(u'abcd', u'efgh')
assert unicode(r) == u'abcdefgh'
def test_empty_rope():
r = rope()
assert isinstance(r, rope)
assert len(r) == 0
def test_singleton_rope():
r = rope('a')
assert repr(r) == "rope('a')"
assert isinstance(r, rope)
assert len(r) == 1
def __init__(self, blob, verify_checksum=False):
pe_offs, = _read(blob[0x3c:], '<H')
fin = BlobIO(blob[pe_offs:])
sig = fin.read(4)
if sig != b'PE\0\0':
raise RuntimeError('Not a PE file: PE signature is missing.')
hdr = _IMAGE_FILE_HEADER.unpack_from_io(fin)
opt_sig, = struct.unpack('<H', fin.read(2))
if opt_sig == IMAGE_NT_OPTIONAL_HDR32_MAGIC:
opt = _IMAGE_OPTIONAL_HEADER32.unpack_from_io(fin)
opt.sig = opt_sig
elif opt_sig == IMAGE_NT_OPTIONAL_HDR64_MAGIC:
opt = _IMAGE_OPTIONAL_HEADER64.unpack_from_io(fin)
opt.sig = opt_sig
else:
raise RuntimeError('Unknown optional header type.')
self._checksum_offs = pe_offs + 4 + _IMAGE_FILE_HEADER.size + 4 * 16
if verify_checksum:
if opt.CheckSum == 0:
self.checksum_correct = False
else:
real_checksum = pe_checksum(
rope(blob[:self._checksum_offs], b'\0\0\0\0',
blob[self._checksum_offs + 4:]))
self.checksum_correct = real_checksum == opt.CheckSum
if opt.FileAlignment == 0:
raise RuntimeError(
'IMAGE_OPTIONAL_HEADER.FileAlignment must be nonzero')
dds = [
_IMAGE_DATA_DIRECTORY.unpack_from_io(fin)
for dd_idx in range(opt.NumberOfRvaAndSizes)
]
def make_pe_section(idx, hdr):
name = hdr.Name.rstrip(b'\0')
if hdr.PointerToRawData % opt.FileAlignment != 0:
raise RuntimeError('Section {}@{} is misaligned ({})'.format(
name, idx, hdr.PointerToRawData))
if hdr.SizeOfRawData % opt.FileAlignment != 0:
raise RuntimeError(
'Size of section {}@{} is misaligned ({})'.format(
name, idx, hdr.SizeOfRawData))
if hdr.PointerToRawData == 0:
data = None
else:
data = blob[hdr.PointerToRawData:hdr.PointerToRawData +
hdr.SizeOfRawData]
return _PeSection(hdr, data)
sections = [
make_pe_section(sec_idx, _IMAGE_SECTION_HEADER.unpack_from_io(fin))
for sec_idx in range(hdr.NumberOfSections)
]
present_secs = sorted(
(sec for sec in sections if sec.hdr.SizeOfRawData != 0),
key=lambda sec: sec.hdr.PointerToRawData)
if not present_secs:
raise RuntimeError('no present sections')
i = 1
while i < len(present_secs):
if present_secs[i - 1].hdr.PointerToRawData + present_secs[
i - 1].hdr.SizeOfRawData != present_secs[
i].hdr.PointerToRawData:
raise RuntimeError('there are holes between sections')
i += 1
last_sec = present_secs[-1]
end_of_image = last_sec.hdr.PointerToRawData + last_sec.hdr.SizeOfRawData
self._blob = blob
self._dos_stub = blob[:pe_offs]
self._file_header = hdr
self._opt_header = opt
self._data_directories = dds
self._sections = sections
self._trailer_offset = end_of_image
self._trailer = blob[end_of_image:]
self._check_vm_overlaps()
def test_concat_rope():
r = rope('a', 'b')
assert isinstance(r, rope)
assert len(r) == 2
def test_empty_leaves():
r = rope('', '')
assert len(r) == 0
assert str(r) == ''
def make_pdb_stream(blocks, size):
r = rope(*(blob[hdr.block_size * block:][:hdr.block_size]
for block in blocks))
return r[:size]
def test_rope_iteration():
r = rope('a', 'bc', '', 'd')
assert list(r) == ['a', 'b', 'c', 'd']
def test_heterogenous_rope():
r = rope('abcd', b'efgh')
assert r[0] == 'a'
assert r[4] == b'e'[0]
def test_slicing():
r = rope('abcd', 'efgh')
rr = r[2:6]
assert len(rr) == 4
assert str(rr) == 'cdef'
def test_rope_chunks():
r = rope('a', 'bc', '', 'd')
assert ''.join(r.chunks) == 'abcd'
