import struct
import win32api  # pip3 install pywin32
import win32con


"""
struct ICONDIR
{
    WORD idReserved;
    WORD idType;
    WORD idCount;
    ICONDIRENTRY idEntries[];
};
struct ICONDIRENTRY
{
    BYTE bWidth;
    BYTE bHeight;
    BYTE bColorCount;
    BYTE bReserved;
    WORD wPlanes;
    WORD wBitCount;
    DWORD dwBytesInRes;
    DWORD dwImageOffset;
};
struct GRPICONDIRENTRY
{
    BYTE bWidth;
    BYTE bHeight;
    BYTE bColorCount;
    BYTE bReserved;
    WORD wPlanes;
    WORD wBitCount;
    DWORD dwBytesInRes;
    WORD nID;
};
"""

def replaceIcon(pe_file, icon_file, icon_group_id, icon_id, lang_id, icon_width):

    class ICONDIRENTRY:
        def __init__(self, data):
            self.bWidth, self.bHeight, self.bColorCount, self.bReserved, self.wPlanes, self.wBitCount, self.dwBytesInRes, self.dwImageOffset = struct.unpack('<4B2H2I', data)
            self.data = b''
        
        def pack(self, nid):
            return struct.pack('<3H4B2HIH', 0, 1, 1, self.bWidth, self.bHeight, self.bColorCount, self.bReserved, self.wPlanes, self.wBitCount, self.dwBytesInRes, nid)


    icon_list = []

    with open(icon_file, 'rb') as f:
        data = f.read(6)
        idReserved, idType, idCount = struct.unpack('<3H', data)
        print(icon_file)
        print('idCount:', idCount)

        for i in range(idCount):
            data = f.read(16)
            icon = ICONDIRENTRY(data)
            icon_list.append(icon)
            print(i+1, 'Width: %d, Height: %d, BitCount: %d' % (icon.bWidth, icon.bHeight, icon.wBitCount))
        
        data = f.read()
        head_offset = 6 + idCount * 16
        for icon in icon_list:
            offset = icon.dwImageOffset - head_offset
            icon.data = data[offset : offset + icon.dwBytesInRes]

    for icon in icon_list:
        # 只写入一张指定宽度和位深的图片
        if icon.bWidth == icon_width:
            data_IconGroup = icon.pack(icon_id)
            data_Icon = icon.data
            try:
                hdst = win32api.BeginUpdateResource(pe_file, 0)
                win32api.UpdateResource(hdst, win32con.RT_GROUP_ICON, icon_group_id, data_IconGroup, lang_id)
                win32api.UpdateResource(hdst, win32con.RT_ICON, icon_id, data_Icon, lang_id)
                win32api.EndUpdateResource (hdst, 0)
                return True
            except Exception as e:
                print(e)
            break
    return False


if __name__ == '__main__':
    print(replaceIcon('WDUIServer.exe', 'win-install.ico', 4008, 1, 2052, 32))