//
//  machoparse.swift
//  Chimera13
//
//  Created by CoolStar on 3/31/20.
//  Copyright © 2020 coolstar. All rights reserved.
//

import Foundation
import MachO.dyld
import CommonCrypto

func parseMacho(path: String, symbol: String) -> UInt32 {
    guard let data = try? Data(contentsOf: URL(fileURLWithPath: path)) else {
        return 0
    }
    
    let cpu = UInt32(littleEndian: data.withUnsafeBytes { $0.load(fromByteOffset: 4, as: UInt32.self) })
    let ncmds = UInt32(littleEndian: data.withUnsafeBytes { $0.load(fromByteOffset: 16, as: UInt32.self) })
    
    guard cpu == 0x100000c else {
        return 0
    }
    
    var symOff = UInt32(0)
    var strOff = UInt32(0)
    var indirectSymOff = UInt32(0)
    var nIndirectSyms = UInt32(0)
    
    var segmentOffset = 32
    for _ in 0..<ncmds {
        let cmd = UInt32(littleEndian: data.withUnsafeBytes({ $0.load(fromByteOffset: segmentOffset, as: UInt32.self) }))
        let cmdSize = UInt32(littleEndian: data.withUnsafeBytes({ $0.load(fromByteOffset: segmentOffset + 4, as: UInt32.self) }))
        defer { segmentOffset += Int(cmdSize) }
        
        if cmd == LC_SYMTAB {
            let parsedCmd = data.withUnsafeBytes({ $0.load( fromByteOffset: segmentOffset, as: symtab_command.self) })
            
            symOff = parsedCmd.symoff
            strOff = parsedCmd.stroff
        }
        
        if cmd == LC_DYSYMTAB {
            let parsedCmd = data.withUnsafeBytes({ $0.load( fromByteOffset: segmentOffset, as: dysymtab_command.self) })
            
            indirectSymOff = parsedCmd.indirectsymoff
            nIndirectSyms = parsedCmd.nindirectsyms
        }
    }
    
    guard symOff != 0,
        strOff != 0,
        indirectSymOff != 0,
        nIndirectSyms != 0 else {
            return 0
    }
    
    segmentOffset = 32
    for _ in 0..<ncmds {
        let cmd = UInt32(littleEndian: data.withUnsafeBytes({ $0.load(fromByteOffset: segmentOffset, as: UInt32.self) }))
        let cmdSize = UInt32(littleEndian: data.withUnsafeBytes({ $0.load(fromByteOffset: segmentOffset + 4, as: UInt32.self) }))
        defer { segmentOffset += Int(cmdSize) }
        
        if cmd == LC_SEGMENT_64 {
            var segment = data.withUnsafeBytes({ $0.load(fromByteOffset: segmentOffset, as: segment_command_64.self) })
            let segNameSz = MemoryLayout.size(ofValue: segment.segname)
            let segNameArr = withUnsafePointer(to: &segment.segname.0) {
                [Int8](UnsafeBufferPointer(start: $0, count: segNameSz))
            }
            guard let segNameRaw = String(bytes: segNameArr.map { UInt8($0) }, encoding: .utf8) else {
                continue
            }
            let segName = segNameRaw.components(separatedBy: "\0")[0]
            if segName == "__DATA" || segName == "__DATA_CONST" {
                var sectionOff = segmentOffset + MemoryLayout.size(ofValue: segment)
                for _ in 0..<segment.nsects {
                    let section = data.withUnsafeBytes({ $0.load(fromByteOffset: sectionOff, as: section_64.self) })
                    if Int32(section.flags) & SECTION_TYPE == S_LAZY_SYMBOL_POINTERS ||
                        Int32(section.flags) & SECTION_TYPE == S_NON_LAZY_SYMBOL_POINTERS {
                        
                        let startIndex = section.reserved1
                        
                        var index = UInt32(0)
                        while index < UInt32(section.size / 8) {
                            defer { index += 1 }
                            
                            let symIndex = startIndex + index
                            let symtabIndex = UInt32(littleEndian: data.withUnsafeBytes({ $0.load(fromByteOffset: Int(indirectSymOff + (symIndex * 4)), as: UInt32.self) }))
                            
                            guard symtabIndex != INDIRECT_SYMBOL_ABS &&
                                symtabIndex != INDIRECT_SYMBOL_LOCAL &&
                                symtabIndex != (INDIRECT_SYMBOL_LOCAL | UInt32(INDIRECT_SYMBOL_ABS)) else {
                                continue
                            }
                            
                            let symbolOffset = Int(symOff) + Int(symtabIndex) * MemoryLayout<nlist_64>.size
                            let symtab = data.withUnsafeBytes({ $0.load(fromByteOffset: symbolOffset, as: nlist_64.self) })
                            guard symtab.n_un.n_strx != 0 else {
                                continue
                            }
                            
                            let stringOff = strOff + symtab.n_un.n_strx
                            let dataSlid = data.advanced(by: Int(stringOff))
                            guard let ptr = dataSlid.withUnsafeBytes({
                                    $0.bindMemory(to: UInt8.self)
                            }).baseAddress else {
                                continue
                            }
                            let str = String(decodingCString: ptr, as: UTF8.self)
                            if str == symbol {
                                let patchOffset = section.offset + (UInt32(index) * 8)
                                return patchOffset
                            }
                        }
                    }
                    sectionOff += MemoryLayout.size(ofValue: section)
                }
            }
        }
    }
    return 0
}

let CSSLOT_CODEDIRECTORY = 0
let CSSLOT_ALTERNATE_CODEDIRECTORIES = 0x1000
let CSSLOT_ALTERNATURE_CODEDIRECORY_MAX = 5
let CSSLOT_ALTERNATE_CODEDIRECTORY_LIMIT = CSSLOT_ALTERNATE_CODEDIRECTORIES + CSSLOT_ALTERNATURE_CODEDIRECORY_MAX
let CSMAGIC_CODEDIRECTORY = 0xfade0c02

let CS_HASHTYPE_SHA1 = 1
let CS_HASHTYPE_SHA256 = 2
let CS_HASHTYPE_SHA256_TRUNCATED = 3
let CS_HASHTYPE_SHA384 = 4

let CODEDIR_SIZE = 0x58

let CS_CDHASH_LEN = 20
let CS_HASH_MAX_SIZE = 48

func parseSuperblob(codedir: Data) -> [UInt8] {
    let count = UInt32(bigEndian: codedir.withUnsafeBytes { $0.load(fromByteOffset: 8, as: UInt32.self) })
   
    var highestHash = 0
    var outputDigest = [UInt8]()
    
    var idxOff = 12
    for _ in 0..<count {
        defer { idxOff += 8 }
        
        guard codedir.count >= idxOff + 8 else {
            break
        }
        let type = UInt32(bigEndian: codedir.withUnsafeBytes { $0.load(fromByteOffset: idxOff, as: UInt32.self) })
        let offset = UInt32(bigEndian: codedir.withUnsafeBytes { $0.load(fromByteOffset: idxOff + 4, as: UInt32.self) })
        
        if type == CSSLOT_CODEDIRECTORY || (type >= CSSLOT_ALTERNATE_CODEDIRECTORIES && type < CSSLOT_ALTERNATE_CODEDIRECTORY_LIMIT) {
            let subblob = Data(codedir[offset..<offset + UInt32(CODEDIR_SIZE)])
            
            let magic = UInt32(bigEndian: subblob.withUnsafeBytes { $0.load(fromByteOffset: 0, as: UInt32.self) })
            let realsize = UInt32(bigEndian: subblob.withUnsafeBytes { $0.load(fromByteOffset: 4, as: UInt32.self) })
            
            if magic == CSMAGIC_CODEDIRECTORY {
                guard codedir.count >= offset + realsize else {
                    continue
                }
                
                let realsubblob = Data(codedir[offset..<offset + realsize])
                
                let hashType = Int(realsubblob.withUnsafeBytes { $0.load(fromByteOffset: 0x25, as: UInt8.self) })
                
                var digest = [UInt8](repeating: 0, count: CS_HASH_MAX_SIZE)
                
                switch hashType {
                case CS_HASHTYPE_SHA1:
                    realsubblob.withUnsafeBytes {
                        _ = CC_SHA1($0.baseAddress, CC_LONG(realsize), &digest)
                    }
                case CS_HASHTYPE_SHA256:
                    fallthrough
                case CS_HASHTYPE_SHA256_TRUNCATED:
                    realsubblob.withUnsafeBytes {
                        _ = CC_SHA256($0.baseAddress, CC_LONG(realsize), &digest)
                    }
                case CS_HASHTYPE_SHA384:
                    realsubblob.withUnsafeBytes {
                        _ = CC_SHA384($0.baseAddress, CC_LONG(realsize), &digest)
                    }
                default:
                    continue
                }
                
                if hashType > highestHash {
                    highestHash = hashType
                    let finalDigest = digest[0..<CS_CDHASH_LEN]
                    outputDigest = [UInt8](finalDigest)
                }
            }
        }
    }
    
    //let hexBytes = outputDigest.map { String(format: "%02hhx", $0) }.joined()
    return outputDigest
}

func getCodeSignature(path: String) -> [UInt8] {
    guard let data = try? Data(contentsOf: URL(fileURLWithPath: path)) else {
        return []
    }
    
    guard data.count >= 8 else {
        return []
    }
    
    var machOffset = 0
    let preferredCPUSubtypes = isArm64e() ? [CPU_SUBTYPE_ARM64E, CPU_SUBTYPE_ARM_V8] : [CPU_SUBTYPE_ARM_V8]
    
    let fat_hdr = UInt32(bigEndian: data.withUnsafeBytes { $0.load(fromByteOffset: 0, as: UInt32.self) })
    for preferredCPUSubtype in preferredCPUSubtypes {
        if fat_hdr == FAT_MAGIC || fat_hdr == FAT_MAGIC_64 {
            let nfatarch = UInt32(bigEndian: data.withUnsafeBytes { $0.load(fromByteOffset: 4, as: UInt32.self) })
            
            var fatOffset = 8
            
            let fatArchSz = MemoryLayout<fat_arch>.size
            for _ in 0..<nfatarch {
                defer { fatOffset += fatArchSz }
                
                guard data.count >= fatOffset + fatArchSz else {
                    continue
                }
                
                let arch = data.withUnsafeBytes { $0.load(fromByteOffset: fatOffset, as: fat_arch.self) }
                
                let cputype = Int32(bigEndian: arch.cputype)
                let subtype = Int32(bigEndian: arch.cpusubtype)
                let offset = UInt32(bigEndian: arch.offset)
                
                guard cputype == CPU_TYPE_ARM64 else {
                    continue
                }
                if subtype == preferredCPUSubtype || (preferredCPUSubtype == CPU_SUBTYPE_ARM_V8 && subtype == CPU_SUBTYPE_ARM64_ALL) {
                    machOffset = Int(offset)
                }
            }
        }
        if machOffset != 0 {
            break
        }
    }
    
    guard machOffset % 4 == 0 else {
        return []
    }
    guard data.count >= machOffset + 32 else {
        return []
    }
    let cpu = UInt32(littleEndian: data.withUnsafeBytes { $0.load(fromByteOffset: machOffset + 4, as: UInt32.self) })
    let ncmds = UInt32(littleEndian: data.withUnsafeBytes { $0.load(fromByteOffset: machOffset + 16, as: UInt32.self) })
    
    guard cpu == 0x100000c else {
        return []
    }
    
    var segmentOffset = machOffset + 32
    for _ in 0..<ncmds {
        guard data.count >= segmentOffset + 8 else {
            continue
        }
        let cmd = UInt32(littleEndian: data.withUnsafeBytes({ $0.load(fromByteOffset: segmentOffset, as: UInt32.self) }))
        let cmdSize = UInt32(littleEndian: data.withUnsafeBytes({ $0.load(fromByteOffset: segmentOffset + 4, as: UInt32.self) }))
        defer { segmentOffset += Int(cmdSize) }
        
        if cmd == LC_CODE_SIGNATURE {
            let loadCmdSize = MemoryLayout<load_command>.size
            guard data.count >= segmentOffset + loadCmdSize + 8 else {
                continue
            }
            
            let off_cd = UInt32(machOffset) + UInt32(littleEndian: data.withUnsafeBytes({
                $0.load(fromByteOffset: segmentOffset + loadCmdSize, as: UInt32.self) }))
            let size_cs = UInt32(littleEndian: data.withUnsafeBytes({ $0.load(fromByteOffset: segmentOffset + loadCmdSize + 4, as: UInt32.self) }))
            
            guard data.count >= off_cd + size_cs else {
                continue
            }
            let subdata = Data(data[off_cd..<off_cd+size_cs])
            return parseSuperblob(codedir: subdata)
        }
    }
    
    return []
}
