Dump macho segment names from swift

segment-dumper.swift

//
//  segment-dumper.swift
//
//  Created by Cero on 2021-10-03.
//

import Foundation
import ArgumentParser
import MachO.loader
import MachO.swap

typealias cString = UnsafePointer<CChar>
typealias mFile = UnsafeMutablePointer<FILE>

let cpuTypeArray: [CPUTypeName] = [
    CPUTypeName(cpu_type: CPU_TYPE_I386, name: "i386"),
    CPUTypeName(cpu_type: CPU_TYPE_X86_64, name: "x86_64"),
    CPUTypeName(cpu_type: CPU_TYPE_ARM, name: "arm"),
    CPUTypeName(cpu_type: CPU_TYPE_ARM64, name: "arm64")
]

struct CPUTypeName {
    let cpu_type: cpu_type_t
    let name: String
    
    static func getSelf(fromType type: cpu_type_t) -> CPUTypeName {
        for typeName in cpuTypeArray {
            if typeName.cpu_type == type {
                return typeName
            }
        }
        return CPUTypeName(cpu_type: CPU_TYPE_ANY, name: "unknown")
    }
}


func mainFunc(_ filePath: String) {
    print("Starting with contents of file in path: \(filePath)")
    let charPath: cString? = NSString.init(string: filePath).utf8String
    guard charPath != nil else {
        print("bruv")
        return
    }
    
    let openMode: cString? = NSString.init(string: "rb").utf8String
    guard openMode != nil else {
        print("I really can't with this rn")
        return
    }
    let file: mFile? = fopen(charPath!, openMode!)
    
    guard file != nil else {
        print("Was given invalid path: \(filePath), returning")
        return
    }
    
    dumpSegments(file!)
}

func loadBytes(_ file: mFile,_ offset: off_t,_ size: Int) -> UnsafeMutableRawPointer {
    let buffer: UnsafeMutableRawPointer = calloc(1, size)
    fseeko(file, offset, SEEK_SET)
    fread(buffer, size, 1, file)
    return buffer
}

func isMagic64(_ magic: UInt32) -> Bool {
    return magic == MH_MAGIC_64 || magic == MH_CIGAM_64
}

func shouldSwapBytes(_ magic: UInt32) -> Bool {
    return magic == MH_CIGAM || magic == MH_CIGAM_64
}


func readMagic(_ file: mFile,_ offset: off_t) -> UInt32 {
    var magic: UInt32 = 0
    fseeko(file, offset, SEEK_SET)
    fread(&magic, MemoryLayout<uint32>.size, 1, file)
    return magic
}

func dumpSegments(_ file: mFile) {
    
    func isFat(_ magic: UInt32) -> Bool {
        return magic == FAT_MAGIC || magic == FAT_CIGAM
    }
    
    let magic: UInt32 = readMagic(file, 0)
    let is64: Bool = isMagic64(magic)
    let shouldSwap: Bool = shouldSwapBytes(magic)
    let isFat: Bool = isFat(magic)
    if isFat && !is64 {
        dumpFatHeader(file, shouldSwap)
    } else if is64 && !isFat {
        dumpMachHeader(file, 0, is64, shouldSwap)
    } else {
        print("Bro...")
        return
    }
}

func dumpFatHeader(
    _ file: mFile,
    _ isSwap: Bool
) {
    let headerSize: size_t = MemoryLayout.size(ofValue: fat_header.self)
    let archSize: size_t = MemoryLayout.size(ofValue: fat_arch.self)
    
    var header: fat_header = loadBytes(file, 0, headerSize).load(as: fat_header.self)
    if isSwap {
        swap_fat_header(&header, NXByteOrder(0))
    }
    
    var archOffset: off_t = Int64(headerSize)
    for _ in 0..<header.nfat_arch {
        var arch: fat_arch = loadBytes(file, archOffset, archSize).load(as: fat_arch.self)
        
        if isSwap {
            swap_fat_arch(&arch, 1, NXByteOrder(0))
        }
        
        let machHeaderOffset: off_t = Int64(arch.offset)
        archOffset += Int64(archSize)
        
        let magic: UInt32 = readMagic(file, machHeaderOffset)
        let is64: Bool = isMagic64(magic)
        let isSwapMach = shouldSwapBytes(magic)
        dumpMachHeader(file, machHeaderOffset, is64, isSwapMach)
    }
}

func dumpMachHeader(
    _ file: mFile,
    _ offset: off_t,
    _ is64: Bool,
    _ isSwap: Bool
) {
    var loadCommandsOffset: off_t = offset
    var headerSize: Int
    var ncmds: UInt32
    if is64 {
        headerSize = MemoryLayout<mach_header_64>.size
        var header: mach_header_64 = loadBytes(file, offset, headerSize).load(as: mach_header_64.self)
        if isSwap {
            swap_mach_header_64(&header, NXByteOrder(0))
        }
        ncmds = header.ncmds
        loadCommandsOffset += Int64(headerSize)
        
        print("CPU Type: \(CPUTypeName.getSelf(fromType: header.cputype).name)")
    } else {
        headerSize = MemoryLayout<mach_header>.size
        var header: mach_header = loadBytes(file, offset, headerSize).load(as: mach_header.self)
        if isSwap {
            swap_mach_header(&header, NXByteOrder(0))
        }
        ncmds = header.ncmds
        loadCommandsOffset += Int64(headerSize)
        print("CPU Type: \(CPUTypeName.getSelf(fromType: header.cputype).name)")
    }
    dumpSegmentCommands(file, loadCommandsOffset, isSwap, ncmds)
}


func dumpSegmentCommands(
    _ file: mFile,
    _ offset: off_t,
    _ isSwap: Bool,
    _ ncmds: UInt32
) {
    var actualOffset: off_t = offset
    
    for _ in 0..<ncmds {
        var command: load_command = loadBytes(file, actualOffset, MemoryLayout<load_command>.size).load(as: load_command.self)
        if isSwap {
            swap_load_command(&command, NXByteOrder(0))
        }
        if command.cmd == LC_SEGMENT_64 {
            var segment: segment_command_64 = loadBytes(file, actualOffset, MemoryLayout<segment_command_64>.size).load(as: segment_command_64.self)
            if isSwap {
                swap_segment_command_64(&segment, NXByteOrder(0))
            }
            let name: String = convertCharTupleToString(&segment.segname)
            print("Segment Name: \(name)")
        } else if command.cmd == LC_SEGMENT {
            var segment: segment_command = loadBytes(file, actualOffset, MemoryLayout<segment_command>.size).load(as: segment_command.self)
            if isSwap {
                swap_segment_command(&segment, NXByteOrder(0))
            }
            let name: String = convertCharTupleToString(&segment.segname)
            print("Segment Name: \(name)")
        }
        actualOffset += Int64(command.cmdsize)
    }
}

typealias garbage = (CChar, CChar, CChar, CChar, CChar, CChar, CChar, CChar, CChar, CChar, CChar, CChar, CChar, CChar, CChar, CChar)

func convertCharTupleToString(_ tuple: UnsafePointer<garbage>) -> String {
    tuple.withMemoryRebound(to: UInt8.self, capacity: MemoryLayout.size(ofValue: tuple)) { cool in
        return String(cString: cool)
    }
}


//MARK: Command Crap
struct Main: ParsableCommand {
    static var configuration: CommandConfiguration = CommandConfiguration(
        commandName: "segdump",
        abstract: "dump segments from MachO files",
        discussion: """
        look I have no idea what I am doing...
        """,
        version: "no u",
        shouldDisplay: true
    )
    
    @Argument(
        help: ArgumentHelp(
            "The complete file path to dump segements",
            shouldDisplay: true
        )
    ) var filePath: String = ""
    
    func run() throws {
        guard filePath != "" else {
            throw ValidationError("Bruh give me a path\n")
        }
        mainFunc(filePath)
    }
}

Main.main()