Skip to content

Commit

Permalink
Use a FileHandle instead of loading whole PE
Browse files Browse the repository at this point in the history
  • Loading branch information
@IsaacMarovitz
IsaacMarovitz committed Oct 27, 2023
1 parent 193b47a commit 3c8eec0
Show file tree
Hide file tree
Showing 8 changed files with 182 additions and 175 deletions.
4 changes: 3 additions & 1 deletion Whisky/Models/Bottle.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,9 @@ extension Bottle {

for link in linkURLs {
do {
if let program = ShellLinkHeader.getProgram(url: link, data: try Data(contentsOf: link), bottle: self) {
if let program = ShellLinkHeader.getProgram(url: link,
handle: try FileHandle(forReadingFrom: link),
bottle: self) {
if !startMenuPrograms.contains(where: { $0.url == program.url }) {
startMenuPrograms.append(program)
try FileManager.default.removeItem(at: link)
Expand Down
33 changes: 8 additions & 25 deletions WhiskyKit/Extensions/DataExtensions.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,34 +18,17 @@

import Foundation

extension Data {
extension FileHandle {
public func extract<T>(_ type: T.Type, offset: Int = 0) -> T? {
if offset + MemoryLayout<T>.size < self.count {
let data = self[offset..<offset + MemoryLayout<T>.size]
return data.withUnsafeBytes { $0.loadUnaligned(as: T.self) }
} else {
do {
try self.seek(toOffset: UInt64(offset))
if let data = try self.read(upToCount: MemoryLayout<T>.size) {
return data.withUnsafeBytes { $0.loadUnaligned(as: T.self)}
}
} catch {
return nil
}
}

// Thanks ChatGPT
public func nullTerminatedStrings(using encoding: String.Encoding = .utf8) -> [String] {
var strings = [String]()
self.withUnsafeBytes { (ptr: UnsafeRawBufferPointer) in
if let baseAddress = ptr.baseAddress {
var strStart = baseAddress
let strEnd = baseAddress + self.count
while strStart < strEnd {
let strPtr = strStart.assumingMemoryBound(to: CChar.self)
let strLen = strnlen(strPtr, self.count)
let strData = Data(bytes: strPtr, count: strLen)
if let str = String(data: strData, encoding: encoding) {
strings.append(str)
}
strStart = strStart.advanced(by: strLen + 1)
}
}
}
return strings
return nil
}
}
54 changes: 27 additions & 27 deletions WhiskyKit/PE/BitmapInfo.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,39 +35,39 @@ public struct BitmapInfoHeader: Hashable {
public let originDirection: BitmapOriginDirection
public let colorFormat: ColorFormat

init(data: Data, offset: Int) {
init(handle: FileHandle, offset: Int) {
var offset = offset
self.size = data.extract(UInt32.self, offset: offset) ?? 0
self.size = handle.extract(UInt32.self, offset: offset) ?? 0
offset += 4
self.width = data.extract(Int32.self, offset: offset) ?? 0
self.width = handle.extract(Int32.self, offset: offset) ?? 0
offset += 4
self.height = data.extract(Int32.self, offset: offset) ?? 0
self.height = handle.extract(Int32.self, offset: offset) ?? 0
offset += 4
self.planes = data.extract(UInt16.self, offset: offset) ?? 0
self.planes = handle.extract(UInt16.self, offset: offset) ?? 0
offset += 2
self.bitCount = data.extract(UInt16.self, offset: offset) ?? 0
self.bitCount = handle.extract(UInt16.self, offset: offset) ?? 0
offset += 2
self.compression = BitmapCompression(rawValue: data.extract(UInt32.self, offset: offset) ?? 0) ?? .rgb
self.compression = BitmapCompression(rawValue: handle.extract(UInt32.self, offset: offset) ?? 0) ?? .rgb
offset += 4
self.sizeImage = data.extract(UInt32.self, offset: offset) ?? 0
self.sizeImage = handle.extract(UInt32.self, offset: offset) ?? 0
offset += 4
self.xPelsPerMeter = data.extract(Int32.self, offset: offset) ?? 0
self.xPelsPerMeter = handle.extract(Int32.self, offset: offset) ?? 0
offset += 4
self.yPelsPerMeter = data.extract(Int32.self, offset: offset) ?? 0
self.yPelsPerMeter = handle.extract(Int32.self, offset: offset) ?? 0
offset += 4
self.clrUsed = data.extract(UInt32.self, offset: offset) ?? 0
self.clrUsed = handle.extract(UInt32.self, offset: offset) ?? 0
offset += 4
self.clrImportant = data.extract(UInt32.self, offset: offset) ?? 0
self.clrImportant = handle.extract(UInt32.self, offset: offset) ?? 0
offset += 4

self.originDirection = self.height < 0 ? .upperLeft : .bottomLeft
self.colorFormat = ColorFormat(rawValue: bitCount) ?? .unknown
}

// swiftlint:disable:next cyclomatic_complexity function_body_length
func renderBitmap(data: Data, offset: Int) -> NSImage {
func renderBitmap(handle: FileHandle, offset: Int) -> NSImage {
var offset = offset
let colorTable = buildColorTable(offset: &offset, data: data)
let colorTable = buildColorTable(offset: &offset, handle: handle)

var pixels: [ColorQuad] = []

Expand All @@ -91,15 +91,15 @@ public struct BitmapInfoHeader: Hashable {
// Ditto .indexed1
break
case .indexed8:
let index = Int(data.extract(UInt8.self, offset: offset) ?? 0)
let index = Int(handle.extract(UInt8.self, offset: offset) ?? 0)
if index >= colorTable.count {
pixelRow.append(ColorQuad(red: 0, green: 0, blue: 0, alpha: 0))
} else {
pixelRow.append(colorTable[Int(index)])
}
offset += 1
case .sampled16:
let sample = data.extract(UInt16.self, offset: offset) ?? 0
let sample = handle.extract(UInt16.self, offset: offset) ?? 0
let red = sample & 0x001F
let green = (sample & 0x03E0) >> 5
let blue = (sample & 0x7C00) >> 10
Expand All @@ -109,24 +109,24 @@ public struct BitmapInfoHeader: Hashable {
alpha: 1))
offset += 2
case .sampled24:
let blue = data.extract(UInt8.self, offset: offset) ?? 0
let blue = handle.extract(UInt8.self, offset: offset) ?? 0
offset += 1
let green = data.extract(UInt8.self, offset: offset) ?? 0
let green = handle.extract(UInt8.self, offset: offset) ?? 0
offset += 1
let red = data.extract(UInt8.self, offset: offset) ?? 0
let red = handle.extract(UInt8.self, offset: offset) ?? 0
offset += 1
pixelRow.append(ColorQuad(red: red,
green: green,
blue: blue,
alpha: 1))
case .sampled32:
let blue = data.extract(UInt8.self, offset: offset) ?? 0
let blue = handle.extract(UInt8.self, offset: offset) ?? 0
offset += 1
let green = data.extract(UInt8.self, offset: offset) ?? 0
let green = handle.extract(UInt8.self, offset: offset) ?? 0
offset += 1
let red = data.extract(UInt8.self, offset: offset) ?? 0
let red = handle.extract(UInt8.self, offset: offset) ?? 0
offset += 1
let alpha = data.extract(UInt8.self, offset: offset) ?? 0
let alpha = handle.extract(UInt8.self, offset: offset) ?? 0
offset += 1
pixelRow.append(ColorQuad(red: red,
green: green,
Expand All @@ -147,15 +147,15 @@ public struct BitmapInfoHeader: Hashable {
return constructImage(pixels: pixels)
}

func buildColorTable(offset: inout Int, data: Data) -> [ColorQuad] {
func buildColorTable(offset: inout Int, handle: FileHandle) -> [ColorQuad] {
var colorTable: [ColorQuad] = []

for _ in 0..<clrUsed {
let blue = data.extract(UInt8.self, offset: offset) ?? 0
let blue = handle.extract(UInt8.self, offset: offset) ?? 0
offset += 1
let green = data.extract(UInt8.self, offset: offset) ?? 0
let green = handle.extract(UInt8.self, offset: offset) ?? 0
offset += 1
let red = data.extract(UInt8.self, offset: offset) ?? 0
let red = handle.extract(UInt8.self, offset: offset) ?? 0
offset += 2

colorTable.append(ColorQuad(red: red,
Expand Down
Loading

0 comments on commit 3c8eec0

Please sign in to comment.