def test_get_icon_file(self, sample):
return_obj = PEGroupIconDir(data=sample(
"438117c7bd53653b3113903bcdb8bd369904a152b524b4676b18a626c2b60e82")
)
assert (
return_obj.get_icon_file(idx=0, data=b"FOOBAR") ==
b"\x00\x00\x01\x00\x01\x00\x00\x00\x04\x00\x00\x00\xff\xff\x00\x00\xb8\x00\x16\x00\x00\x00FOOBAR"
)
def _get_icon_info(self):
"""Get icon in PNG format and information for searching for similar icons
@return: tuple of (image data in PNG format encoded as base64, md5 hash of image data, md5 hash of "simplified" image for fuzzy matching)
"""
if not self.pe:
return None, None, None
try:
rt_group_icon_idx = [entry.id for entry in self.pe.DIRECTORY_ENTRY_RESOURCE.entries].index(pefile.RESOURCE_TYPE['RT_GROUP_ICON'])
rt_group_icon_dir = self.pe.DIRECTORY_ENTRY_RESOURCE.entries[rt_group_icon_idx]
entry = rt_group_icon_dir.directory.entries[0]
offset = entry.directory.entries[0].data.struct.OffsetToData
size = entry.directory.entries[0].data.struct.Size
peicon = PEGroupIconDir(self.pe.get_memory_mapped_image()[offset:offset+size])
bigwidth = 0
bigheight = 0
bigbpp = 0
bigidx = -1
iconidx = 0
for idx,icon in enumerate(peicon.icons):
if icon.bWidth >= bigwidth and icon.bHeight >= bigheight and icon.wBitCount >= bigbpp:
bigwidth = icon.bWidth
bigheight = icon.bHeight
bigbpp = icon.wBitCount
bigidx = icon.nID
iconidx = idx
rt_icon_idx = [entry.id for entry in self.pe.DIRECTORY_ENTRY_RESOURCE.entries].index(pefile.RESOURCE_TYPE['RT_ICON'])
rt_icon_dir = self.pe.DIRECTORY_ENTRY_RESOURCE.entries[rt_icon_idx]
for entry in rt_icon_dir.directory.entries:
if entry.id == bigidx:
offset = entry.directory.entries[0].data.struct.OffsetToData
size = entry.directory.entries[0].data.struct.Size
icon = peicon.get_icon_file(iconidx, self.pe.get_memory_mapped_image()[offset:offset+size])
strio = StringIO()
output = StringIO()
simplified = StringIO()
strio.write(icon)
strio.seek(0)
img = Image.open(strio)
img.save(output, format="PNG")
img = img.resize((16,16), Image.ANTIALIAS)
img = img.convert("RGB").convert("P", palette=Image.ADAPTIVE, colors=4)
img.save(simplified, format="PNG", colors=4)
m = hashlib.md5()
m.update(output.getvalue())
fullhash = m.hexdigest()
m = hashlib.md5()
m.update(simplified.getvalue())
simphash = m.hexdigest()
return base64.b64encode(output.getvalue()), fullhash, simphash
except:
return None, None, None
def _get_icon(self):
"""Get icon in PNG format.
@return: image data in PNG format, encoded as base64
"""
if not self.pe:
return None
try:
rt_group_icon_idx = [
entry.id for entry in self.pe.DIRECTORY_ENTRY_RESOURCE.entries
].index(pefile.RESOURCE_TYPE['RT_GROUP_ICON'])
rt_group_icon_dir = self.pe.DIRECTORY_ENTRY_RESOURCE.entries[
rt_group_icon_idx]
entry = rt_group_icon_dir.directory.entries[0]
offset = entry.directory.entries[0].data.struct.OffsetToData
size = entry.directory.entries[0].data.struct.Size
peicon = PEGroupIconDir(
self.pe.get_memory_mapped_image()[offset:offset + size])
bigwidth = 0
bigheight = 0
bigbpp = 0
bigidx = -1
iconidx = 0
for idx, icon in enumerate(peicon.icons):
if icon.bWidth >= bigwidth and icon.bHeight >= bigheight and icon.wBitCount >= bigbpp:
bigwidth = icon.bWidth
bigheight = icon.bHeight
bigbpp = icon.wBitCount
bigidx = icon.nID
iconidx = idx
rt_icon_idx = [
entry.id for entry in self.pe.DIRECTORY_ENTRY_RESOURCE.entries
].index(pefile.RESOURCE_TYPE['RT_ICON'])
rt_icon_dir = self.pe.DIRECTORY_ENTRY_RESOURCE.entries[rt_icon_idx]
for entry in rt_icon_dir.directory.entries:
if entry.id == bigidx:
offset = entry.directory.entries[
0].data.struct.OffsetToData
size = entry.directory.entries[0].data.struct.Size
icon = peicon.get_icon_file(
iconidx,
self.pe.get_memory_mapped_image()[offset:offset +
size])
strio = StringIO()
output = StringIO()
strio.write(icon)
strio.seek(0)
img = Image.open(strio)
img.save(output, format="PNG")
return base64.b64encode(output.getvalue())
except:
return None
def _get_icon(self):
"""Get icon in PNG format.
@return: image data in PNG format, encoded as base64
"""
if not self.pe:
return None
try:
rt_group_icon_idx = [entry.id for entry in self.pe.DIRECTORY_ENTRY_RESOURCE.entries].index(pefile.RESOURCE_TYPE['RT_GROUP_ICON'])
rt_group_icon_dir = self.pe.DIRECTORY_ENTRY_RESOURCE.entries[rt_group_icon_idx]
entry = rt_group_icon_dir.directory.entries[0]
offset = entry.directory.entries[0].data.struct.OffsetToData
size = entry.directory.entries[0].data.struct.Size
peicon = PEGroupIconDir(self.pe.get_memory_mapped_image()[offset:offset+size])
bigwidth = 0
bigheight = 0
bigbpp = 0
bigidx = -1
iconidx = 0
for idx,icon in enumerate(peicon.icons):
if icon.bWidth >= bigwidth and icon.bHeight >= bigheight and icon.wBitCount >= bigbpp:
bigwidth = icon.bWidth
bigheight = icon.bHeight
bigbpp = icon.wBitCount
bigidx = icon.nID
iconidx = idx
rt_icon_idx = [entry.id for entry in self.pe.DIRECTORY_ENTRY_RESOURCE.entries].index(pefile.RESOURCE_TYPE['RT_ICON'])
rt_icon_dir = self.pe.DIRECTORY_ENTRY_RESOURCE.entries[rt_icon_idx]
for entry in rt_icon_dir.directory.entries:
if entry.id == bigidx:
offset = entry.directory.entries[0].data.struct.OffsetToData
size = entry.directory.entries[0].data.struct.Size
icon = peicon.get_icon_file(iconidx, self.pe.get_memory_mapped_image()[offset:offset+size])
strio = StringIO()
output = StringIO()
strio.write(icon)
strio.seek(0)
img = Image.open(strio)
img.save(output, format="PNG")
return base64.b64encode(output.getvalue())
except:
return None
def test_init(self, sample):
return_obj = PEGroupIconDir(data=sample(
"438117c7bd53653b3113903bcdb8bd369904a152b524b4676b18a626c2b60e82")
)
assert len(return_obj.icons) == 3
def _get_icon_info(self):
"""Get icon in PNG format and information for searching for similar icons
@return: tuple of (image data in PNG format encoded as base64, md5 hash of image data, md5 hash of "simplified" image for fuzzy matching)
"""
if not self.pe:
return None, None, None
try:
rt_group_icon_idx = [
entry.id for entry in self.pe.DIRECTORY_ENTRY_RESOURCE.entries
].index(pefile.RESOURCE_TYPE['RT_GROUP_ICON'])
rt_group_icon_dir = self.pe.DIRECTORY_ENTRY_RESOURCE.entries[
rt_group_icon_idx]
entry = rt_group_icon_dir.directory.entries[0]
offset = entry.directory.entries[0].data.struct.OffsetToData
size = entry.directory.entries[0].data.struct.Size
peicon = PEGroupIconDir(
self.pe.get_memory_mapped_image()[offset:offset + size])
bigwidth = 0
bigheight = 0
bigbpp = 0
bigidx = -1
iconidx = 0
for idx, icon in enumerate(peicon.icons):
if icon.bWidth >= bigwidth and icon.bHeight >= bigheight and icon.wBitCount >= bigbpp:
bigwidth = icon.bWidth
bigheight = icon.bHeight
bigbpp = icon.wBitCount
bigidx = icon.nID
iconidx = idx
rt_icon_idx = [
entry.id for entry in self.pe.DIRECTORY_ENTRY_RESOURCE.entries
].index(pefile.RESOURCE_TYPE['RT_ICON'])
rt_icon_dir = self.pe.DIRECTORY_ENTRY_RESOURCE.entries[rt_icon_idx]
for entry in rt_icon_dir.directory.entries:
if entry.id == bigidx:
offset = entry.directory.entries[
0].data.struct.OffsetToData
size = entry.directory.entries[0].data.struct.Size
icon = peicon.get_icon_file(
iconidx,
self.pe.get_memory_mapped_image()[offset:offset +
size])
strio = StringIO()
output = StringIO()
strio.write(icon)
strio.seek(0)
img = Image.open(strio)
img.save(output, format="PNG")
img = img.resize((8, 8), Image.BILINEAR)
img = img.convert("RGB").convert("P",
palette=Image.ADAPTIVE,
colors=2).convert("L")
lowval = img.getextrema()[0]
img = img.point(lambda i: 255 if i > lowval else 0)
img = img.convert("1")
simplified = bytearray(img.getdata())
m = hashlib.md5()
m.update(output.getvalue())
fullhash = m.hexdigest()
m = hashlib.md5()
m.update(simplified)
simphash = m.hexdigest()
return base64.b64encode(
output.getvalue()), fullhash, simphash
except:
pass
return None, None, None
def get_icon_info(self, pe: pefile.PE) -> Tuple[str, str, str, str]:
"""Get icon in PNG format and information for searching for similar icons
@return: tuple of (image data in PNG format encoded as base64, md5 hash of image data, md5 hash of "simplified"
image for fuzzy matching)
"""
if not pe or not hasattr(pe, "DIRECTORY_ENTRY_RESOURCE"):
return None, None, None, None
try:
idx = [entry.id for entry in pe.DIRECTORY_ENTRY_RESOURCE.entries]
if pefile.RESOURCE_TYPE["RT_GROUP_ICON"] not in idx:
return None, None, None, None
rt_group_icon_idx = idx.index(
pefile.RESOURCE_TYPE["RT_GROUP_ICON"])
rt_group_icon_dir = pe.DIRECTORY_ENTRY_RESOURCE.entries[
rt_group_icon_idx]
entry = rt_group_icon_dir.directory.entries[0]
offset = entry.directory.entries[0].data.struct.OffsetToData
size = entry.directory.entries[0].data.struct.Size
peicon = PEGroupIconDir(
pe.get_memory_mapped_image()[offset:offset + size])
bigwidth = 0
bigheight = 0
bigbpp = 0
bigidx = -1
iconidx = 0
if hasattr(peicon, "icons") and peicon.icons:
for idx, icon in enumerate(peicon.icons):
if icon.bWidth >= bigwidth and icon.bHeight >= bigheight and icon.wBitCount >= bigbpp:
bigwidth = icon.bWidth
bigheight = icon.bHeight
bigbpp = icon.wBitCount
bigidx = icon.nID
iconidx = idx
rt_icon_idx = False
rt_icon_idx_tmp = [
entry.id for entry in pe.DIRECTORY_ENTRY_RESOURCE.entries
]
if pefile.RESOURCE_TYPE["RT_ICON"] in rt_icon_idx_tmp:
rt_icon_idx = rt_icon_idx_tmp.index(
pefile.RESOURCE_TYPE["RT_ICON"])
if not isinstance(rt_icon_idx, int):
return None, None, None, None
rt_icon_dir = pe.DIRECTORY_ENTRY_RESOURCE.entries[rt_icon_idx]
for entry in rt_icon_dir.directory.entries:
if entry.id == bigidx:
offset = entry.directory.entries[
0].data.struct.OffsetToData
size = entry.directory.entries[0].data.struct.Size
icon = peicon.get_icon_file(
iconidx,
pe.get_memory_mapped_image()[offset:offset + size])
with BytesIO() as byteio:
byteio.write(icon)
byteio.seek(0)
try:
with Image.open(
byteio) as img, BytesIO() as output:
img.save(output, format="PNG")
dhash = self.generate_icon_dhash(img)
img = (img.resize(
(8, 8),
Image.BILINEAR).convert("RGB").convert(
"P", palette=Image.ADAPTIVE,
colors=2).convert("L"))
lowval = img.getextrema()[0]
img = img.point(lambda i: 255 if i > lowval
else 0).convert("1")
simplified = bytearray(img.getdata())
m = hashlib.md5()
m.update(output.getvalue())
fullhash = m.hexdigest()
m = hashlib.md5()
m.update(simplified)
simphash = m.hexdigest()
icon = base64.b64encode(
output.getvalue()).decode()
except ValueError:
log.error(
"parse_pe.py -> get_incon_info -> buffer is not large enough"
)
return None, None, None, None
except OSError as e:
log.error(e)
return None, None, None, None
return icon, fullhash, simphash, dhash
except Exception as e:
log.error(e, exc_info=True)
return None, None, None, None
