107 lines
3.9 KiB
Python
107 lines
3.9 KiB
Python
"""Extract the largest embedded icon from a Windows PE executable as PNG."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import io
|
|
import struct
|
|
from pathlib import Path
|
|
|
|
from PIL import Image
|
|
|
|
try:
|
|
import pefile
|
|
except ImportError: # pragma: no cover - pefile is a dev/runtime dependency
|
|
pefile = None # type: ignore[assignment]
|
|
|
|
|
|
def _raw_icon_to_image(raw: bytes) -> Image.Image | None:
|
|
"""Decode a single RT_ICON blob (PNG-embedded or BMP) to RGBA."""
|
|
png_start = raw.find(b"\x89PNG")
|
|
if png_start >= 0:
|
|
return Image.open(io.BytesIO(raw[png_start:])).convert("RGBA")
|
|
|
|
if len(raw) < 40:
|
|
return None
|
|
|
|
bi_size = struct.unpack("<I", raw[0:4])[0]
|
|
bi_width = struct.unpack("<i", raw[4:8])[0]
|
|
bi_height = struct.unpack("<i", raw[8:12])[0]
|
|
bi_bpp = struct.unpack("<H", raw[14:16])[0]
|
|
actual_h = abs(bi_height) // 2
|
|
if bi_width <= 0 or actual_h <= 0:
|
|
return None
|
|
|
|
pixel_offset = bi_size
|
|
if bi_bpp <= 8:
|
|
num_colors = struct.unpack("<I", raw[32:36])[0] or (2**bi_bpp)
|
|
pixel_offset += num_colors * 4
|
|
|
|
row_bytes = ((bi_width * bi_bpp + 31) // 32) * 4
|
|
xor_data = raw[pixel_offset : pixel_offset + row_bytes * actual_h]
|
|
|
|
if bi_bpp == 32:
|
|
img = Image.frombytes("RGBA", (bi_width, actual_h), xor_data[: bi_width * actual_h * 4])
|
|
elif bi_bpp == 24:
|
|
img = Image.frombytes("RGB", (bi_width, actual_h), xor_data[: bi_width * actual_h * 3]).convert("RGBA")
|
|
elif bi_bpp == 8:
|
|
num_colors = struct.unpack("<I", raw[32:36])[0] or 256
|
|
palette: list[int] = []
|
|
for c in range(num_colors):
|
|
b, g, r, _ = struct.unpack("BBBB", raw[40 + c * 4 : 44 + c * 4])
|
|
palette.extend([r, g, b])
|
|
img = Image.frombytes("P", (bi_width, actual_h), xor_data[: bi_width * actual_h])
|
|
img.putpalette(palette)
|
|
img = img.convert("RGBA")
|
|
else:
|
|
return None
|
|
|
|
return img.transpose(Image.FLIP_TOP_BOTTOM)
|
|
|
|
|
|
def extract_icon_blobs(exe_path: str | Path) -> dict[int, bytes]:
|
|
"""Return ``{resource_id: raw_icon_bytes}`` from *exe_path*."""
|
|
if pefile is None:
|
|
raise RuntimeError("pefile is required to extract icons from executables")
|
|
|
|
pe = pefile.PE(str(exe_path), fast_load=True)
|
|
pe.parse_data_directories(directories=[pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_RESOURCE"]])
|
|
|
|
blobs: dict[int, bytes] = {}
|
|
if hasattr(pe, "DIRECTORY_ENTRY_RESOURCE"):
|
|
for entry in pe.DIRECTORY_ENTRY_RESOURCE.entries:
|
|
if entry.id != pefile.RESOURCE_TYPE["RT_ICON"]:
|
|
continue
|
|
for icon_entry in entry.directory.entries:
|
|
icon_id = icon_entry.struct.Id
|
|
for data_entry in icon_entry.directory.entries:
|
|
offset = data_entry.data.struct.OffsetToData
|
|
size = data_entry.data.struct.Size
|
|
blobs[icon_id] = pe.get_data(offset, size)
|
|
pe.close()
|
|
return blobs
|
|
|
|
|
|
def extract_best_icon_image(exe_path: str | Path) -> Image.Image | None:
|
|
"""Return the largest decodable icon from *exe_path*, or ``None``."""
|
|
blobs = extract_icon_blobs(exe_path)
|
|
if not blobs:
|
|
return None
|
|
|
|
best_raw = max(blobs.values(), key=len)
|
|
return _raw_icon_to_image(best_raw)
|
|
|
|
|
|
def extract_best_icon_png(exe_path: str | Path, out_path: str | Path, *, min_size: int = 0) -> tuple[int, int] | None:
|
|
"""Write the largest icon from *exe_path* to *out_path*; return size or ``None``."""
|
|
img = extract_best_icon_image(exe_path)
|
|
if img is None:
|
|
return None
|
|
if min_size and max(img.size) < min_size:
|
|
scale = min_size / max(img.size)
|
|
new_w = max(1, int(round(img.width * scale)))
|
|
new_h = max(1, int(round(img.height * scale)))
|
|
img = img.resize((new_w, new_h), Image.Resampling.NEAREST)
|
|
out_path = Path(out_path)
|
|
out_path.parent.mkdir(parents=True, exist_ok=True)
|
|
img.save(out_path)
|
|
return img.size
|