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'