1#!/usr/bin/env node 2 3/* 4 * Copyright (c) 2022-2023 Huawei Device Co., Ltd. 5 * Licensed under the Apache License, Version 2.0 (the "License"); 6 * you may not use this file except in compliance with the License. 7 * You may obtain a copy of the License at 8 * 9 * http://www.apache.org/licenses/LICENSE-2.0 10 * 11 * Unless required by applicable law or agreed to in writing, software 12 * distributed under the License is distributed on an "AS IS" BASIS, 13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 * See the License for the specific language governing permissions and 15 * limitations under the License. 16 */ 17 18import fs from "node:fs" 19import path from "node:path"; 20 21const { O_RDWR } = fs.constants 22 23const DOS_HEADER_LEN = 0x40; 24const DOS_HEADER_MAGIC = 0x5A4D; 25const PE_HEADER_POS_OFFSET = 0x3c; 26const PE_HEADER_LEN = 0x112; // for 64-bit 27const PE_HEADER_MAGIC = 0x00004550; 28 29const COFF_HEADER_SIZE = 24; 30const PE_NUMBER_OF_SECTIONS_OFFSET = 6; 31const PE_MAGIC_OFFSET = COFF_HEADER_SIZE + 0; 32const PE_SIZE_OF_HEADERS_OFFSET = COFF_HEADER_SIZE + 60; 33const PE_NUMBER_OF_RVA_AND_SIZES_OFFSET_32 = COFF_HEADER_SIZE + 92; 34const PE_NUMBER_OF_RVA_AND_SIZES_OFFSET_64 = COFF_HEADER_SIZE + 108; 35const PE_RESOURCE_SECTION_INFO_OFFSET_32 = COFF_HEADER_SIZE + 112; 36const PE_RESOURCE_SECTION_INFO_OFFSET_64 = COFF_HEADER_SIZE + 128; 37const PE_SECTION_HEADER_SIZE = 40 38 39const RT_ICON = 3 40const RT_GROUP_ICON = RT_ICON + 11 41 42export function setIcon(filePath, iconPath) { 43 44 process.on("uncaughtException", err => { 45 console.error("" + err); 46 process.exit(1); 47 }) 48 49 if (!filePath || !iconPath) printUsageAndExit() 50 51 let fd = fs.openSync(filePath, O_RDWR); 52 let dosHeader = Buffer.alloc(DOS_HEADER_LEN); 53 let bytesRead = fs.readSync(fd, dosHeader, 0, DOS_HEADER_LEN); 54 55 if (bytesRead !== DOS_HEADER_LEN || dosHeader.readUInt16LE(0) !== DOS_HEADER_MAGIC) { 56 throw new Error("Not a PE file") 57 } 58 59 let peOffset = dosHeader.readUInt32LE(PE_HEADER_POS_OFFSET); 60 let peHeader = Buffer.alloc(PE_HEADER_LEN); 61 bytesRead = fs.readSync(fd, peHeader, 0, PE_HEADER_LEN, peOffset); 62 if (bytesRead !== PE_HEADER_LEN || peHeader.readUInt32LE(0) !== PE_HEADER_MAGIC) { 63 throw new Error("Not a PE file") 64 } 65 66 let sections = readSections(fd, peHeader, peOffset); 67 68 let rsrc = sections.find(s => s.name === ".rsrc"); 69 if (!rsrc) { 70 console.error("no resource section") 71 process.exit(1); 72 } 73 74 let resources = readResourceSection(fd, rsrc.size, rsrc.offset); 75 76 const iconDirResource = resources.items.find(r => r.id === RT_ICON); 77 const iconResources = iconDirResource.items 78 .map(res => ({ name: "icon-" + res.id, id: res.id, size: res.items[0].size, offset: rsrc.offset + res.items[0].offset - rsrc.virtualAddr, tableOffset: res.items[0].fileOffset })) 79 const groupIconResources = resources.items.find(r => r.id === RT_GROUP_ICON)?.items 80 .map(res => ({ name: "groupicon-" + res.id, size: res.items[0].size, offset: rsrc.offset + res.items[0].offset - rsrc.virtualAddr, tableOffset: res.items[0].fileOffset })) 81 82 if (!checkResourcesAreSequential(iconResources)) { 83 console.error("Icons must be placed sequentially, WIP") 84 process.exit(1) 85 } 86 87 const newIcon = fs.readFileSync(iconPath) 88 const lastResource = iconResources[iconResources.length - 1] 89 const imageResource = iconResources[0] 90 const icoHeaderResource = groupIconResources[0] 91 writeImage(fd, newIcon, imageResource.offset, lastResource.offset + lastResource.size); 92 fixImageDirectoryResource(fd, icoHeaderResource.tableOffset,icoHeaderResource.offset, icoHeaderResource.size, imageResource.id, newIcon.length); 93 fixImageDataEntry(fd, imageResource.tableOffset, newIcon.length); 94 fixImageDirectoryEntry(fd, iconDirResource.fileOffset, imageResource.id, iconResources.length); 95 96 fs.closeSync(fd); 97} 98 99function readSections(fd, peHeader, peOffset) { 100 101 let is64 = peHeader.readUint16LE(PE_MAGIC_OFFSET) === 0x20b; 102 103 let numRVAAndSizesOffset = is64 ? PE_NUMBER_OF_RVA_AND_SIZES_OFFSET_64 : PE_NUMBER_OF_RVA_AND_SIZES_OFFSET_32; 104 let numRVAAndSizes = peHeader.readUInt32LE(numRVAAndSizesOffset); 105 106 let sectionsOffset = peOffset + numRVAAndSizesOffset + 4 + 8 * numRVAAndSizes; 107 let numSections = peHeader.readUInt16LE(PE_NUMBER_OF_SECTIONS_OFFSET); 108 let sectionsHeaders = Buffer.alloc(numSections * PE_SECTION_HEADER_SIZE); 109 let _bytesRead = fs.readSync(fd, sectionsHeaders, 0, numSections * PE_SECTION_HEADER_SIZE, sectionsOffset); 110 let sections = []; 111 for (let i = 0; i < numSections; i++) { 112 let sectionHeader = sectionsHeaders.subarray(i * PE_SECTION_HEADER_SIZE, (i + 1) * PE_SECTION_HEADER_SIZE); 113 let sectionName = sectionHeader.toString("utf8", 0, cstrlen(sectionHeader, 0, 8)); 114 let sectionOffset = sectionHeader.readUInt32LE(20); 115 let sectionSize = sectionHeader.readUInt32LE(16); 116 let sectionVirtAddr = sectionHeader.readUint32LE(12); 117 sections.push({ name: sectionName, size: sectionSize, offset: sectionOffset, virtualAddr: sectionVirtAddr }); 118 } 119 return sections; 120} 121 122function readResourceSection(fd, size, offset) { 123 let section = Buffer.alloc(size) 124 let _ = fs.readSync(fd, section, 0, size, offset) 125 const fileOffset = offset 126 127 const TABLE_HEADER_SIZE = 16 128 const DIRECTORY_ENTRY_SIZE = 8 129 130 function readDirectoryTable(offset, id) { 131 let entry = { type: "directory", id: id, items: [], fileOffset: fileOffset + offset } 132 let numNamed = section.readUint16LE(offset + 12) 133 let numIds = section.readUint16LE(offset + 14) 134 135 let off = offset + TABLE_HEADER_SIZE 136 for (let i = 0; i < numNamed; i++) { 137 entry.items.push(readDirectoryEntry(off, true)) 138 off += DIRECTORY_ENTRY_SIZE 139 } 140 for (let i = 0; i < numIds; i++) { 141 entry.items.push(readDirectoryEntry(off, false)) 142 off += DIRECTORY_ENTRY_SIZE 143 } 144 145 return entry; 146 } 147 148 function readDataEntry(offset, id) { 149 let entry = { type: "data", id: id, offset: 0, size: 0, fileOffset: fileOffset + offset } 150 entry.offset = section.readUInt32LE(offset) 151 entry.size = section.readUInt32LE(offset + 4) 152 return entry; 153 } 154 155 function readStringEntry(offset) { 156 let length = section.readUint16LE(offset); 157 let value = section.toString("utf16le", offset + 2, offset + 2 + length) 158 if (value.endsWith("\0")) value = value.slice(0, -1) 159 return value 160 } 161 162 function readDirectoryEntry(offset, named) { 163 let id = section.readUInt32LE(offset); 164 let itemOffset = section.readUInt32LE(offset + 4) 165 if (named) { 166 id = readStringEntry(id) 167 } 168 if (itemOffset & 0x80000000) { 169 itemOffset &= 0x0FFFFFFF 170 return readDirectoryTable(itemOffset, id) 171 } else { 172 return readDataEntry(itemOffset, id) 173 } 174 } 175 176 return readDirectoryTable(0, 0) 177} 178 179function printUsageAndExit() { 180 let exe = path.basename(process.argv[0]) 181 let script = path.basename(process.argv[1]) 182 console.log(`USAGE: ${exe} ${script} <file> <icon path>`) 183 process.exit(1); 184} 185 186function cstrlen(buf, start, end) { 187 for (let i = start; i < end; i++) { 188 if (buf[i] === 0) return i; 189 } 190 191 return end - start; 192} 193 194function checkResourcesAreSequential(items) { 195 if (!items.length) return true 196 let offset = items[0].offset 197 for (const item of items) { 198 if (offset !== item.offset) { 199 return false 200 } 201 offset += item.size 202 } 203 204 return true 205} 206 207function writeImage(fd, data, start, end) { 208 writeWithTrailingZeroes(fd, data, start, end) 209} 210 211function writeWithTrailingZeroes(fd, data, start, end) { 212 if (end - start < data.length) { 213 throw new Error("Image too large") 214 } 215 216 fs.writeSync(fd, data, 0, data.length, start); 217 let offset = start + data.length 218 let zeros = Buffer.alloc(4096); 219 while (offset < end) { 220 let delta = Math.min(zeros.length, end - offset); 221 offset += fs.writeSync(fd, zeros, 0, delta, offset); 222 } 223} 224 225function fixImageDirectoryResource(fd, tableOffset, dataOffset, dataSize, iconId, iconSize) { 226 let newIcoDirectory = Buffer.alloc(6 + 14); 227 newIcoDirectory.writeUInt16LE(0, 0); // Must be zero 228 newIcoDirectory.writeUInt16LE(1, 2); // 1 for ICON 229 newIcoDirectory.writeUInt16LE(1, 4); // num images, only one is supported by this script 230 newIcoDirectory.writeUInt8(0, 6); // width, icon must 256x256 231 newIcoDirectory.writeUInt8(0, 7); // height 232 newIcoDirectory.writeUInt8(0, 8); // no palette 233 newIcoDirectory.writeUInt8(0, 9); // reserved 234 newIcoDirectory.writeUInt16LE(1, 10); // 1 plane 235 newIcoDirectory.writeUInt16LE(32, 12); // bpp 236 newIcoDirectory.writeUInt32LE(iconSize, 14); // size 237 newIcoDirectory.writeUInt16LE(iconId, 18); // id 238 239 writeWithTrailingZeroes(fd, newIcoDirectory, dataOffset, dataOffset + dataSize); 240 241 let buffer = Buffer.alloc(4) 242 buffer.writeUInt32LE(newIcoDirectory.length) 243 fs.writeSync(fd, buffer, 0, 4, tableOffset + 4); 244} 245 246function fixImageDataEntry(fd, tableOffset, iconSize) { 247 let buffer = Buffer.alloc(4) 248 buffer.writeUInt32LE(iconSize) 249 fs.writeSync(fd, buffer, 0, 4, tableOffset + 4); 250} 251 252function fixImageDirectoryEntry(fd, tableOffset, iconId, numIcons) { 253 const DIRECTORY_ENTRY_SIZE = 8 254 const TABLE_HEADER_SIZE = 16 255 256 let buffer = Buffer.alloc(8) 257 buffer.writeUInt16LE(0, 0) // No entries with string names 258 buffer.writeUInt16LE(1, 2) // 1 entry with id 259 buffer.writeUInt32LE(iconId, 4) // point to current icon ID 260 fs.writeSync(fd, buffer, 0, 8, tableOffset + 12); 261 let zerosStart = tableOffset + TABLE_HEADER_SIZE + DIRECTORY_ENTRY_SIZE 262 let zerosEnd = zerosStart + (numIcons - 1) * DIRECTORY_ENTRY_SIZE 263 writeWithTrailingZeroes(fd, Buffer.alloc(0), zerosStart, zerosEnd) 264} 265