defmodule LibPE do
@moduledoc """
Implementation of the Windows PE executable format for reading and writing PE binaries.
Most struct member names are taken directly from the windows documentation:
https://docs.microsoft.com/en-us/windows/win32/debug/pe-format
This library has been created specifically to archieve the following:
* Update the PE checksum in `erl.exe` after making changes
* Insert a Microsoft manifest file after compilation: https://docs.microsoft.com/de-de/windows/win32/sbscs/application-manifests
* Insert an Executable Icon after compilation
"""
@exe_header_size 4
defstruct [
:meta,
:meta2,
:msdos_stub,
:format,
:machine,
:timestamp,
:object_offset,
:object_entry_count,
:coff_header,
:coff_flags,
:coff_sections,
:rest
]
@spec parse_file(binary()) :: {:ok, %LibPE{}}
def parse_file(filename) do
parse_string(File.read!(filename))
end
@spec parse_string(binary()) :: {:ok, %LibPE{}}
def parse_string(
<<"MZ", meta::binary-size(6), @exe_header_size::little-size(16), meta2::binary-size(50),
offset::little-size(32), rest::binary>> = full_image
) do
# exe header is always 64 bytes
stub_size = offset - @exe_header_size * 16
<<msdos_stub::binary-size(stub_size), coff::binary>> = rest
pe =
%LibPE{meta: meta, meta2: meta2, msdos_stub: msdos_stub, rest: ""}
|> parse_coff(coff, full_image)
end_offset = byte_size(encode(pe))
pe =
if end_offset < byte_size(full_image) do
%LibPE{pe | rest: binary_part(full_image, end_offset, byte_size(full_image) - end_offset)}
else
pe
end
{:ok, pe}
end
@spec encode(%LibPE{}) :: binary()
def encode(%LibPE{rest: rest, coff_sections: sections, coff_header: header} = pe) do
image = encode_header(pe)
# Compare with Section.parse_section()
# We assume the `raw_data` is the actual "virtual" payload
# So we fit it here (padding or cutting) into the raw size because
# often the real payload is smaller than the file space, but sometimes it can
# be larger as well (when the tail is only made of zeros)
image =
Enum.reduce(sections, image, fn sec, image ->
binary_pad_trailing(image, sec.pointer_to_raw_data) <> sec.raw_data
end)
image =
if header.certificate_data != nil do
{start, _size} = header.certificate_table
binary_pad_trailing(image, start) <> header.certificate_data
else
image
end
image <> rest
end
defp encode_header(%LibPE{meta: meta, meta2: meta2, msdos_stub: msdos_stub} = pe) do
stub_size = byte_size(msdos_stub)
offset = stub_size + @exe_header_size * 16
coff = encode_coff(pe)
<<"MZ", meta::binary-size(6), @exe_header_size::little-size(16), meta2::binary-size(50),
offset::little-size(32), msdos_stub::binary-size(stub_size), coff::binary>>
end
defp parse_coff(pe, <<"PE\0\0", rest::binary>>, full_image) do
parse_coff(%LibPE{pe | format: :pe}, rest, full_image)
end
defp parse_coff(
%LibPE{} = pe,
<<machine::little-size(16), number_of_sections::little-size(16),
timestamp::little-size(32), object_offset::little-size(32),
object_entry_count::little-size(32), coff_header_size::little-size(16),
coff_flags::little-size(16), rest::binary>>,
full_image
) do
<<header::binary-size(coff_header_size), rest::binary>> = rest
header = LibPE.OptionalHeader.parse(header, full_image)
{sections, _rest} = LibPE.Section.parse(rest, number_of_sections, full_image)
%LibPE{
pe
| machine: LibPE.MachineType.decode(machine),
timestamp: timestamp,
object_offset: object_offset,
object_entry_count: object_entry_count,
coff_header: header,
coff_flags: LibPE.Characteristics.decode(coff_flags),
coff_sections: sections
}
end
defp encode_coff(%LibPE{format: :pe} = pe) do
"PE\0\0" <> encode_coff(%LibPE{pe | format: nil})
end
defp encode_coff(%LibPE{
machine: machine,
timestamp: timestamp,
object_offset: object_offset,
object_entry_count: object_entry_count,
coff_header: header,
coff_flags: coff_flags,
coff_sections: sections
}) do
machine = LibPE.MachineType.encode(machine)
coff_flags = LibPE.Characteristics.encode(coff_flags)
header = LibPE.OptionalHeader.encode(header)
coff_header_size = byte_size(header)
number_of_sections = length(sections)
sections = Enum.map_join(sections, &LibPE.Section.encode/1)
<<machine::little-size(16), number_of_sections::little-size(16), timestamp::little-size(32),
object_offset::little-size(32), object_entry_count::little-size(32),
coff_header_size::little-size(16), coff_flags::little-size(16), header::binary,
sections::binary>>
end
@doc """
Update the PE image checksum of a PE file.
"""
def update_checksum(%LibPE{} = pe) do
tmp_pe =
%LibPE{pe | coff_header: %LibPE.OptionalHeader{pe.coff_header | checksum: 0}}
|> encode()
# size = byte_size(tmp_pe) + byte_size(LibPE.OptionalHeader.encode_checksum(pe.coff_header))
size = byte_size(tmp_pe)
# correcting size for the missing checksum field
new_checksum = LibPE.Checksum.checksum(tmp_pe, size)
%LibPE{pe | coff_header: %LibPE.OptionalHeader{pe.coff_header | checksum: new_checksum}}
end
@doc """
Update the section & certificate layout after a section size has
been changed
"""
@quad_word_size 64
def update_layout(%LibPE{coff_sections: sections, coff_header: header} = pe) do
offset = byte_size(encode_header(pe))
%LibPE.OptionalHeader{file_alignment: file_alignment, section_alignment: virtual_alignment} =
header
{sections, offsets} =
Enum.map_reduce(sections, {offset, offset}, fn %LibPE.Section{
virtual_data: virtual_data
} = sec,
{virtual, raw} ->
virtual = ceil(virtual / virtual_alignment) * virtual_alignment
raw = ceil(raw / file_alignment) * file_alignment
virtual_size = byte_size(virtual_data)
raw_size = byte_size(String.trim_trailing(virtual_data, "\0"))
raw_size = ceil(raw_size / file_alignment) * file_alignment
raw_data =
if virtual_size < raw_size do
binary_pad_trailing(virtual_data, raw_size, sec.padding)
else
binary_part(virtual_data, 0, raw_size)
end
sec = %LibPE.Section{
sec
| pointer_to_raw_data: raw,
size_of_raw_data: raw_size,
virtual_address: virtual,
virtual_size: virtual_size,
raw_data: raw_data
}
{sec, {virtual + virtual_size, raw + raw_size}}
end)
header =
if header.certificate_data != nil do
{_virtual, raw} = offsets
start = ceil(raw / @quad_word_size) * @quad_word_size
%LibPE.OptionalHeader{
header
| certificate_table: {start, byte_size(header.certificate_data)}
}
else
%LibPE.OptionalHeader{header | certificate_table: {0, 0}}
end
# Only updating the physical tables `.rsrc` and `.reloc`
header =
[
# export_table: ".edata",
# import_table: ".idata",
resource_table: ".rsrc",
# exception_table: ".pdata",
base_relocation_table: ".reloc"
# debug: ".debug",
# tls_table: ".tls"
]
|> Enum.reduce(header, fn {field_name, section_name}, header ->
case Enum.find(sections, fn %LibPE.Section{name: name} -> name == section_name end) do
nil ->
Map.put(header, field_name, {0, 0})
%LibPE.Section{virtual_address: virtual_address, virtual_size: virtual_size} ->
Map.put(header, field_name, {virtual_address, virtual_size})
end
end)
header =
case List.last(sections) do
%LibPE.Section{virtual_address: addr, virtual_size: size} ->
%LibPE.OptionalHeader{
header
| size_of_image: ceil((addr + size) / virtual_alignment) * virtual_alignment
}
nil ->
header
end
%LibPE{pe | coff_sections: sections, coff_header: header}
end
def get_resources(%LibPE{coff_sections: sections}) do
case Enum.find(sections, fn %LibPE.Section{name: name} -> name == ".rsrc" end) do
%LibPE.Section{virtual_data: virtual_data, virtual_address: virtual_address} ->
LibPE.ResourceTable.parse(virtual_data, virtual_address)
nil ->
%LibPE.ResourceTable{}
end
end
defp ensure_resource_section(%LibPE{coff_sections: sections} = pe) do
idx = Enum.find_index(sections, fn %LibPE.Section{name: name} -> name == ".rsrc" end)
if idx == nil do
%LibPE{pe | coff_sections: sections ++ [%LibPE.Section{name: ".rsrc"}]}
else
pe
end
end
def set_resources(pe, resources = %LibPE.ResourceTable{}) do
# need to ensure that the virtual_address is up-to-date
%LibPE{coff_sections: sections} = pe = update_layout(ensure_resource_section(pe))
# now fetching and setting the resource
idx = Enum.find_index(sections, fn %LibPE.Section{name: name} -> name == ".rsrc" end)
section = %LibPE.Section{virtual_address: virtual_address} = Enum.at(sections, idx)
data = LibPE.ResourceTable.encode(resources, virtual_address)
section = %LibPE.Section{section | virtual_data: data}
sections = List.update_at(sections, idx, fn _ -> section end)
# updating offsets coming after the ".rsrc" section
%LibPE{pe | coff_sections: sections}
|> update_layout()
end
def set_resource(
pe,
resource_type,
data,
codepage \\ 0,
language \\ 1033
) do
resources =
get_resources(pe)
|> LibPE.ResourceTable.set_resource(resource_type, data, codepage, language)
set_resources(pe, resources)
end
@doc false
def binary_pad_trailing(binary, size, padding \\ <<0>>)
def binary_pad_trailing(binary, size, padding) when byte_size(binary) < size do
bytesize = size - byte_size(binary)
padding = String.duplicate(padding, ceil(bytesize / byte_size(padding)) * byte_size(padding))
binary <> binary_part(padding, 0, bytesize)
end
def binary_pad_trailing(binary, _size, _padding) do
binary
end
@doc false
def binary_extract(binary, start, size) do
content = binary_part(binary, start, size)
binary =
binary_part(binary, 0, start) <>
binary_pad_trailing("", size) <>
binary_part(binary, start + size, byte_size(binary) - (start + size))
{binary, content}
end
end
