Extractor for Higurashi Switch ROM

HigurashiSwitchReader.swift

import Foundation

guard CommandLine.arguments.count > 2 else {
	print("Usage: \(CommandLine.arguments[0]) data.rom outputFolder")
	exit(EXIT_FAILURE)
}
let outputPath = CommandLine.arguments[2]

func fileInfo(path: String) -> (exists: Bool, isDirectory: Bool) {
	var isDir: ObjCBool = false
	let exists = FileManager.default.fileExists(atPath: path, isDirectory: &isDir)
	return (exists, exists ? isDir.boolValue : false)
}

guard fileInfo(path: outputPath).isDirectory else {
	print("\(outputPath) wasn't a folder!")
	exit(EXIT_FAILURE)
}

extension MutableCollection {
	mutating func mapInPlace(_ transform: (Element) throws -> Element) rethrows {
		var i = startIndex
		while i != endIndex {
			self[i] = try transform(self[i])
			i = index(after: i)
		}
	}
}

let fileStart: UInt64
let offsetMultiplier: UInt64

let file = fopen(CommandLine.arguments[1], "r")
let header = { () -> [UInt32] in
	var tmp = [UInt32](repeating: 0, count: 4);
	guard fread(&tmp, 4, 4, file) == 4 else {
		fatalError("Failed to read header")
	}
	tmp.mapInPlace { UInt32(littleEndian: $0) }
	return tmp
}()

if header[0] == 0x204d4f52 /*ROM */ {
	fileStart = 0x10
	offsetMultiplier = 0x800
}
else if header[0] == 0x324d4f52 /*ROM2*/ {
	fileStart = 0x20
	offsetMultiplier = 0x200
}
else {
	print("Bad file header, first 4 bytes were \(String(header[0].bigEndian, radix: 16))")
	exit(0)
}

let infoSize = header[2]

let infoData = { () -> UnsafeRawBufferPointer in
	let tmp = UnsafeMutableRawBufferPointer.allocate(byteCount: Int(infoSize), alignment: MemoryLayout<UInt32>.alignment)
	fseek(file, Int(fileStart), SEEK_SET)
	guard fread(tmp.baseAddress, 1, Int(infoSize), file) == Int(infoSize) else {
		fatalError("Failed to read folder listing")
	}
	return UnsafeRawBufferPointer(tmp)
}()

struct File {
	let name: String
	let offset: UInt32
	let size: UInt32
	let isFolder: Bool
	var globalOffset: UInt64 {
		if isFolder {
			return UInt64(offset) &+ fileStart
		}
		else {
			return UInt64(offset) &* offsetMultiplier
		}
	}
	init(ptr: UnsafeRawBufferPointer, nameOffset: UInt32, offset: UInt32, size: UInt32, end: inout UInt32) {
		self.offset = offset
		self.size = size
		self.isFolder = nameOffset & (1 << 31) != 0
		let nameOffset = Int(nameOffset & ((1 << 31) - 1))
		let nameBase = ptr.bindMemory(to: CChar.self).baseAddress!.advanced(by: nameOffset)
		let len = strlen(nameBase)
		end = max(end, UInt32(nameOffset + len + 1))
		name = String(decoding: ptr[nameOffset..<(nameOffset + len)], as: UTF8.self)
	}
}

struct Folder {
	let offset: UInt32
	let parentOffset: UInt32
	let files: [File]
	init(ptr: inout UnsafeRawBufferPointer) {
		let uint32list = ptr.bindMemory(to: UInt32.self).lazy.map(UInt32.init(littleEndian:))
		let count = uint32list[0]
		var end: UInt32 = 0
		let files = (0..<count).map { pos -> File in
			let index = Int(pos) * 3 + 1
			return File(
				ptr: ptr,
				nameOffset: uint32list[index],
				offset: uint32list[index + 1],
				size: uint32list[index + 2],
				end: &end
			)
		}
		end = (end + 0x10 - 1) & ~(0x10 - 1)
		self.files = files
		
		guard
			files.count >= 2,
			files[0].name == ".",
			files[1].name == "..",
			files[0].isFolder,
			files[1].isFolder
		else {
			fatalError("Read invalid folder!")
		}
		
		offset = UInt32(files[0].offset)
		parentOffset = UInt32(files[0].offset)
		
		ptr = UnsafeRawBufferPointer(rebasing: ptr[Int(end)...])
	}
}

var folders: [UInt32: Folder] = [:]
var folderReadPtr = infoData
while !folderReadPtr.isEmpty {
	let folder = Folder(ptr: &folderReadPtr)
	guard !folders.keys.contains(folder.offset) else {
		fatalError("Folder referenced something other than itself as itself")
	}
	folders[folder.offset] = folder
}
var unwrittenFolders = Set(folders.keys)
guard let firstFolder = folders[0] else {
	fatalError("Missing base folder!")
}

struct FileIOError: Error { let msg: String }

func writeFile(_ fileReference: File, url: URL) throws {
	guard !FileManager.default.fileExists(atPath: url.path) else {
		print("File already exists at \(url.path), skipping")
		return
	}
	let output = fopen(url.path, "w")
	defer { fclose(output) }
	let blockSize = 1024 * 1024
	let buffer = UnsafeMutableRawPointer.allocate(byteCount: min(blockSize, Int(fileReference.size)), alignment: 1)
	defer { buffer.deallocate() }
	fseek(file, Int(fileReference.globalOffset), SEEK_SET);
	var leftToCopy = fileReference.size
	func copyData(size: Int) throws {
		guard fread(buffer, 1, size, file) == size else {
			throw FileIOError(msg: "Failed to read contents of \(url.path)")
		}
		guard fwrite(buffer, 1, size, output) == size else {
			throw FileIOError(msg: "Failed to write contents of \(url.path)")
		}
	}
	while leftToCopy >= blockSize {
		try copyData(size: blockSize)
		leftToCopy -= UInt32(blockSize)
	}
	if leftToCopy > 0 {
		try copyData(size: Int(leftToCopy))
	}
}

extension Folder {
	func writeToPath(path: URL) {
		unwrittenFolders.remove(offset)
		for file in files[2...] {
			let newPath = path.appendingPathComponent(file.name)
			print("Writing \(newPath.path)")
			if file.isFolder {
				do {
					try FileManager.default.createDirectory(at: newPath, withIntermediateDirectories: false, attributes: nil)
				} catch {
					print("Failed to create \(newPath.path): \(error)")
				}
				let folder = folders[file.offset]
				folder?.writeToPath(path: newPath)
			}
			else {
				do {
					try writeFile(file, url: newPath)
				} catch let error as FileIOError {
					print(error.msg)
					// Delete any partial files
					try? FileManager.default.removeItem(at: newPath)
				} catch {
					fatalError("Only FileIOErrors should be thrown, this shouldn't happen!")
				}
			}
		}
	}
}

firstFolder.writeToPath(path: URL(fileURLWithPath: outputPath))
if !unwrittenFolders.isEmpty {
	print("Warning: These folders didn't get written:")
	for folder in unwrittenFolders.lazy.map({ folders[$0]! }) {
		print("\tFolder \(folder.offset):")
		for file in folder.files {
			print("\t\t\(file.name)\(file.isFolder ? "/" : "")")
		}
	}
}