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(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