def test_html():
"""Tests the html construct."""
# Test doctests
results = doctest.testmod(construct.construct_html)
assert not results.failed
# Test with an example
EMBED_SPEC = construct.Struct(
'a' / construct.IP4Address,
'b' / construct.IP4Address,
'c' / construct.IP4Address,
'd' / construct.IP4Address
)
address_struct = construct.Struct(
'first' / construct.Struct('a' / construct.Byte, 'b' / construct.Byte),
'second' / construct.Struct('inner2' / construct.Bytes(2))
# 'internal' / IP4Address
)
PACKET = construct.Struct(
construct.Padding(0x9),
'Hardcoded Value 1' / construct.HexString(construct.Int32ul),
'Hardcoded Value 2' / construct.HexString(construct.Int32ul),
'Hardcoded Value 3' / construct.HexString(construct.Int32ul),
construct.Padding(0x17),
'Compromised Host IP' / construct.IP4Address, # Use IP adapter
# 'Unknown IP Addresses' / construct.Switch(
# this['Hardcoded Value 1'],
# {
# '0x1f4' : EMBED_SPEC
# },
# ),
'Unknown IP Addresses' / address_struct[4],
# 'Unknown IP Addresses' / IP4Address[4],
construct.Padding(8),
'Unknown Indicator' / construct.String(0xF),
construct.Padding(2),
'Number of CPUs' / construct.Int32ul,
'CPU Mhz' / construct.Int32ul,
'Total Memory (MB)' / construct.Int32ul,
'Compromised System Kernel' / construct.CString(),
'Possible Trojan Version' / construct.CString()
)
data = (b'\x01\x00\x00\x00}\x00\x00\x00\x00\xf4\x01\x00\x002\x00\x00\x00\xe8'
b'\x03\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01'
b'\x01\x00\x00\x00\x00\x01\x00\x00\x00\xc0\xa8\x01\r\xc0\xa8\x01\r\xc0'
b'\xa8\x01\r\xc0\xa8\x01\r\xc0\xa8\x01\r\xff\xff\x01\x00\x00\x00\x00\x00'
b'-== Love AV ==-:\x00\x01\x00\x00\x00d\n\x00\x00\xc4\x07\x00\x00'
b'Linux 3.13.0-93-generic\x001:G2.40\x00')
html_data = construct.html_hex(PACKET, data, depth=1)
with open(os.path.join(os.path.dirname(__file__), 'construct_html.html'), 'r') as fo:
expected_html_data = fo.read()
assert html_data == expected_html_data
def build(self, write_section_data=True):
"""
Generate PE file.
:param write_section_data: Whether to include section data (otherwise only the headers are written)
:returns bytes: PE file data.
:raises ValueError: If set attributes contains contradicting data.
"""
pe = copy.deepcopy(self)
# Pad dos stub to match e_lfanew. (Warn if dos stub is too large.)
dos_stub_size = pe.DosHeader.e_lfanew - construct.IMAGE_DOS_HEADER.sizeof(
)
if len(pe.DosStub) > dos_stub_size:
raise ValueError(
'Provided DOS stub is too large for provided DosHeader.e_lfanew: {}'
.format(pe.DosHeader.e_lfanew))
pe.DosStub = pe.DosStub.ljust(dos_stub_size, b'\x00')
# Fix file header.
file_header = pe.NTHeaders.FileHeader
if file_header.NumberOfSections and file_header.NumberOfSections != len(
pe.SectionTable):
logger.debug(
'NTHeaders.FileHeader.NumberOfSections does not equal the number of sections provided. Auto-adjusting.'
)
file_header.NumberOfSections = len(pe.SectionTable)
if isinstance(file_header.Characteristics, list):
file_header.Characteristics = {
flag: True
for flag in file_header.Characteristics
}
# Fix sections.
pe.SectionTable = list(map(self._fix_section, pe.SectionTable))
# Fix optional header.
optional_header = pe.NTHeaders.OptionalHeader
number_of_rva_and_sizes = len(optional_header.DataDirectory)
optional_header.NumberOfRvaAndSizes = number_of_rva_and_sizes
file_header.SizeOfOptionalHeader = construct.IMAGE_OPTIONAL_HEADER.sizeof(
NumberOfRvaAndSizes=number_of_rva_and_sizes)
if isinstance(optional_header.DllCharacteristics, list):
optional_header.DllCharacteristics = {
flag: True
for flag in optional_header.DllCharacteristics
}
# SizeOfHeaders is the sum of the headers rounded by FileAlignment.
headers_size = construct.PEFILE_HEADER.sizeof(**pe)
file_alignment = optional_header.FileAlignment - 1
headers_size = (headers_size +
file_alignment) & 0xffffffff - file_alignment
optional_header.SizeOfHeaders = headers_size
pe_data = construct.PEFILE_HEADER.build(pe)
# Add section data.
if write_section_data:
stream = io.BytesIO(pe_data)
spec = construct.Pointer(this.PointerToRawData,
construct.Bytes(this.SizeOfRawData))
for section in pe.SectionTable:
spec.build_stream(section.data, stream, **section)
pe_data = stream.getvalue()
return pe_data