feat: 添加自动提取图标和创建快捷方式

eaio/function/check.py

@@ -9,7 +9,7 @@
 from eaio import __electron_repo_root__, __electron_repo__
 from eaio.util.utils import dir_tree, file_crc, get_all_drives, to_drive
 from eaio.util.status import LinkStatus, RepoStatus
-from eaio.util.error import ScanError, RepoError, TargetError
+from eaio.util.error import ScanError, RepoError, TargetError, PEError
 
 
 def is_electron_exe(path: Path) -> bool:
@@ -52,26 +52,23 @@ def parse_electron_exe(path: Path) -> tuple[str, str]:
             case _:
                 msg = f'{path.name} 的 CPU 架构未知:{pefile.MACHINE_TYPE[pe.FILE_HEADER.Machine]}，如确认为应用入口，则需要提交 issue'
                 logger.warning(msg)
-                raise ScanError(msg)
+                raise PEError(msg)
 
         section_rdata = list(filter(lambda section: section.Name.strip(b'\x00') == b'.rdata', pe.sections))
         if len(section_rdata) == 0:
             msg = f'{path.name} 无 .rdata 段，如确认为应用入口，则需要提交 issue'
             logger.warning(msg)
-            raise ScanError(msg)
+            raise PEError(msg)
         elif len(section_rdata) > 1:
             logger.warning(f'{path.name} 的 .rdata 段不唯一，默认使用第一个')
         section_rdata = section_rdata[0]
-
-    with open(path, 'rb') as f:
-        f.seek(section_rdata.PointerToRawData)
-        rdata = f.read(section_rdata.SizeOfRawData)
+        rdata = pe.get_data(section_rdata.VirtualAddress, section_rdata.SizeOfRawData)
 
     versions = [i.decode() for i in (set(re.findall(rb'Chrome/(?:[0-9.]+?|%s) Electron/(\S+?)\x00', rdata)))]
     if len(versions) == 0:
         msg = f'{path.name} 的 .rdata 段中找不到版本信息，如确认为应用入口，则需要提交 issue'
         logger.warning(msg)
-        raise ScanError(msg)
+        raise PEError(msg)
     elif len(versions) > 1:
         logger.warning(f'{path.name} 的 .rdata 段中版本信息不唯一:{versions}，默认使用第一个')
     electron_version = versions[0]

eaio/function/link.py

@@ -5,10 +5,11 @@
 from loguru import logger
 
 from eaio import __electron_repo_root__, __electron_repo__
-from eaio.util.utils import to_drive
+from eaio.util.error import PEError
+from eaio.util.utils import to_drive, extract_icon, create_win_lnk
 
 
-def create_link(repo: Path, repo_name: Path, target: Path, target_name: Path):
+def create_link(repo: Path, repo_name: Path | str, target: Path, target_name: Path):
     logger.debug(f"将 {repo_name} 链接到 {target_name}")
     repo_file = repo.joinpath(repo_name)
     target_file = target.joinpath(target_name)
@@ -22,7 +23,19 @@ def link(app_entry: Path, arch: str, version: str, files: Iterable[Path]):
     target = app_entry.parent
     for file in files:
         relative_name = file.relative_to(target)
-        create_link(repo, 'electron.exe' if file == app_entry else relative_name, target, relative_name)
+        if file == app_entry:
+            ico_data = b''
+            try:
+                ico_data = extract_icon(app_entry)
+            except PEError as e:
+                logger.error(f"提取图标失败\t{e}")
+            if ico_data:
+                with open(target.joinpath('eaio.ico'), 'wb') as f:
+                    f.write(ico_data)
+            create_win_lnk(app_entry.with_suffix('.lnk'), app_entry, target.joinpath('eaio.ico') if ico_data else None)
+            create_link(repo, 'electron.exe', target, relative_name)
+        else:
+            create_link(repo,  relative_name, target, relative_name)
 
 
 def delete_link(target: Path, target_name: Path | str):

eaio/util/error.py

@@ -24,3 +24,10 @@ class DownloadError(Exception):
     下载时错误
     """
     pass
+
+
+class PEError(Exception):
+    """
+    解析 PE 文件时错误
+    """
+    pass

eaio/util/utils.py

@@ -1,11 +1,15 @@
 import binascii
 import io
 import math
+import struct
 from pathlib import Path
 from typing import Generator
-from ctypes import create_string_buffer
+from ctypes import create_string_buffer, wintypes
 
 from loguru import logger
+import pefile
+
+from eaio.util.error import PEError
 
 
 def dir_tree(path: Path, depth: int = 0) -> Generator[tuple[Path, int], None, None]:
@@ -84,4 +88,93 @@ def str_size(size_bytes: int) -> str:
     return "%s %s" % (s, size_name[i])
 
 
+def parse_icon_group(icon_group: bytes, icon_datas: dict[int, bytes]):
+    result_header = b''
+    result_body = b''
+
+    ico_reserved, ico_type, ico_number = struct.unpack('<HHH', icon_group[:6])
+    if ico_reserved != 0:
+        msg = f'{ico_reserved} != 0'
+        logger.warning(msg)
+        raise PEError(msg)
+    if ico_type != 1:
+        msg = f'{ico_reserved} != 1'
+        logger.warning(msg)
+        raise PEError(msg)
+    result_header += struct.pack('<HHH', ico_reserved, ico_type, ico_number)
+    for i in range(ico_number):
+        ico_image_width, ico_image_height, ico_image_color_count, ico_image_reserved, ico_image_color_places, ico_image_bits, ico_image_size, ico_image_offset = struct.unpack('<BBBBHHIH', icon_group[6+i*14:6+(i+1)*14])
+        result_header += struct.pack('<BBBBHHII', ico_image_width, ico_image_height, ico_image_color_count, ico_image_reserved, ico_image_color_places, ico_image_bits, ico_image_size, 6 + 16 * ico_number + len(result_body))
+        result_body += icon_datas[ico_image_offset]
+
+    return result_header + result_body
+
+
+def extract_icon(target):
+    with pefile.PE(target, fast_load=True) as pe:
+        pe.parse_data_directories([
+            pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_RESOURCE'],
+        ])
+
+        if not hasattr(pe, 'DIRECTORY_ENTRY_RESOURCE'):
+            msg = f'{target.name} 没有 IMAGE_DIRECTORY_ENTRY_RESOURCE'
+            logger.warning(msg)
+            raise PEError(msg)
+
+        icon_group_entries = [resource for resource in pe.DIRECTORY_ENTRY_RESOURCE.entries if resource.id == 14]
+        if len(icon_group_entries) == 0:
+            msg = f'{target.name} 没有 RT_GROUP_ICON'
+            logger.warning(msg)
+            raise PEError(msg)
+        elif len(icon_group_entries) > 1:
+            logger.warning(f'{target.name} 的 RT_GROUP_ICON 不唯一，默认使用第一个')
+        icon_group_data = None
+        for entry in icon_group_entries[0].directory.entries:
+            if entry.struct.Id == 1:  # 1 represents the default icon group
+                data_entry = entry.directory.entries[0]
+                icon_group_data = pe.get_data(data_entry.data.struct.OffsetToData, data_entry.data.struct.Size)
+                break
+        if icon_group_data is None:
+            msg = f'{target.name} 的 RT_GROUP_ICON 中未找到默认图标 1'
+            logger.warning(msg)
+            raise PEError(msg)
+
+        icon_entries = [resource for resource in pe.DIRECTORY_ENTRY_RESOURCE.entries if resource.id == 3]
+        if len(icon_entries) == 0:
+            msg = f'{target.name} 没有 RT_ICON'
+            logger.warning(msg)
+            raise PEError(msg)
+        icon_datas = {}
+        for entry in icon_entries[0].directory.entries:
+            data_entry = entry.directory.entries[0]
+            icon_datas[entry.struct.Id] = pe.get_data(data_entry.data.struct.OffsetToData, data_entry.data.struct.Size)
+
+        return parse_icon_group(icon_group_data, icon_datas)
+
+
+def create_win_lnk(target: Path, source: Path, icon: Path | None = None):
+    import pylnk3
+    lnk = pylnk3.create(str(target))
+    lnk.link_flags.IsUnicode = True
+
+    levels = list(pylnk3.path_levels(source))
+    elements = [pylnk3.RootEntry(pylnk3.ROOT_MY_COMPUTER),
+                pylnk3.DriveEntry(levels[0])]
+    for level in levels[1:]:
+        segment = pylnk3.PathSegmentEntry.create_for_path(level)
+        elements.append(segment)
+    lnk.shell_item_id_list = pylnk3.LinkTargetIDList()
+    lnk.shell_item_id_list.items = elements
+
+    if icon:
+        lnk.link_flags.HasIconLocation = True
+        lnk.icon = str(icon)
+        lnk.icon_index = 0
+
+    lnk.link_flags.HasWorkingDir = True
+    lnk.work_dir = str(source.parent)
+
+    lnk.save()
+
+
 log = io.StringIO()

pyproject.toml

@@ -10,6 +10,7 @@ dependencies = [
     "loguru>=0.7.0",
     "requests>=2.30.0",
     "pysocks>=1.7.1",
+    "pylnk3>=0.4.2",
 ]
 requires-python = ">=3.10"
 readme = "README.md"