extractor.py

import os
import cv2
import sys
import win32ui
import win32gui
import win32con
import win32api
import numpy as np
from PIL import Image

def extract_ico(target_file):
    ico_x = win32api.GetSystemMetrics(win32con.SM_CXICON)
    try:
        large, small = win32gui.ExtractIconEx(target_file,0)
        win32gui.DestroyIcon(small[0])
    except Exception as e:
        print(e)
        return
    hdc = win32ui.CreateDCFromHandle(win32gui.GetDC(0))
    hbmp = win32ui.CreateBitmap()
    hbmp.CreateCompatibleBitmap(hdc, ico_x, ico_x)
    hdc = hdc.CreateCompatibleDC()

    hdc.SelectObject(hbmp)
    hdc.DrawIcon((0,0), large[0])

    bmpstr = hbmp.GetBitmapBits(True)
    img = Image.frombuffer('RGBA', (32,32), bmpstr, 'raw', 'BGRA', 0, 1)
    return img

def get_img_file(dir_path):
    file_list = []
    for root, dir_names, file_names in os.walk(dir_path):
        for file_name in file_names:
            file_list.append(os.path.join(root, file_name))
        return file_list
    
if __name__ == "__main__":
    raw_file_dir = "C:/Users/g0mx/Desktop/Down/Down"
    file_list = get_img_file(raw_file_dir)
    for file_name in file_list:
        p_img = extract_ico(file_name)
        try:
            file_tuple = os.path.split(file_name)
            p_img.save("C:/Users/g0mx/Desktop/malware_icon/" + file_tuple[1] + ".png")
        except Exception as e:
            print(e)
            continue