"""Checks the validity of MachO binary signatures

MachO binaries sometimes include a LC_CODE_SIGNATURE load command
and corresponding section in the __LINKEDIT segment that together
work to "sign" the binary. This script is used to check the validity
of this signature.

Usage:
    ./code-signature-check.py my_binary 800 300 0 800

Arguments:
   binary - The MachO binary to be tested
   offset - The offset from the start of the binary to where the code signature section begins
   size - The size of the code signature section in the binary
   code_offset - The point in the binary to begin hashing
   code_size - The length starting from code_offset to hash
"""

import argparse
import collections
import hashlib
import itertools
import struct
import sys
import typing

class CodeDirectoryVersion:
    SUPPORTSSCATTER = 0x20100
    SUPPORTSTEAMID = 0x20200
    SUPPORTSCODELIMIT64 = 0x20300
    SUPPORTSEXECSEG = 0x20400

class CodeDirectory:
    @staticmethod
    def make(buf: memoryview) -> typing.Union['CodeDirectoryBase', 'CodeDirectoryV20100', 'CodeDirectoryV20200', 'CodeDirectoryV20300', 'CodeDirectoryV20400']:
        _magic, _length, version = struct.unpack_from(">III", buf, 0)
        subtype = {
            CodeDirectoryVersion.SUPPORTSSCATTER: CodeDirectoryV20100,
            CodeDirectoryVersion.SUPPORTSTEAMID: CodeDirectoryV20200,
            CodeDirectoryVersion.SUPPORTSCODELIMIT64: CodeDirectoryV20300,
            CodeDirectoryVersion.SUPPORTSEXECSEG: CodeDirectoryV20400,
        }.get(version, CodeDirectoryBase)

        return subtype._make(struct.unpack_from(subtype._format(), buf, 0))

class CodeDirectoryBase(typing.NamedTuple):
    magic: int
    length: int
    version: int
    flags: int
    hashOffset: int
    identOffset: int
    nSpecialSlots: int
    nCodeSlots: int
    codeLimit: int
    hashSize: int
    hashType: int
    platform: int
    pageSize: int
    spare2: int

    @staticmethod
    def _format() -> str:
        return ">IIIIIIIIIBBBBI"

class CodeDirectoryV20100(typing.NamedTuple):
    magic: int
    length: int
    version: int
    flags: int
    hashOffset: int
    identOffset: int
    nSpecialSlots: int
    nCodeSlots: int
    codeLimit: int
    hashSize: int
    hashType: int
    platform: int
    pageSize: int
    spare2: int

    scatterOffset: int

    @staticmethod
    def _format() -> str:
        return CodeDirectoryBase._format() + "I"

class CodeDirectoryV20200(typing.NamedTuple):
    magic: int
    length: int
    version: int
    flags: int
    hashOffset: int
    identOffset: int
    nSpecialSlots: int
    nCodeSlots: int
    codeLimit: int
    hashSize: int
    hashType: int
    platform: int
    pageSize: int
    spare2: int

    scatterOffset: int

    teamOffset: int

    @staticmethod
    def _format() -> str:
        return CodeDirectoryV20100._format() + "I"

class CodeDirectoryV20300(typing.NamedTuple):
    magic: int
    length: int
    version: int
    flags: int
    hashOffset: int
    identOffset: int
    nSpecialSlots: int
    nCodeSlots: int
    codeLimit: int
    hashSize: int
    hashType: int
    platform: int
    pageSize: int
    spare2: int

    scatterOffset: int

    teamOffset: int

    spare3: int
    codeLimit64: int

    @staticmethod
    def _format() -> str:
        return CodeDirectoryV20200._format() + "IQ"

class CodeDirectoryV20400(typing.NamedTuple):
    magic: int
    length: int
    version: int
    flags: int
    hashOffset: int
    identOffset: int
    nSpecialSlots: int
    nCodeSlots: int
    codeLimit: int
    hashSize: int
    hashType: int
    platform: int
    pageSize: int
    spare2: int

    scatterOffset: int

    teamOffset: int

    spare3: int
    codeLimit64: int

    execSegBase: int
    execSegLimit: int
    execSegFlags: int

    @staticmethod
    def _format() -> str:
        return CodeDirectoryV20300._format() + "QQQ"

class CodeDirectoryBlobIndex(typing.NamedTuple):
    type_: int
    offset: int

    @staticmethod
    def make(buf: memoryview) -> 'CodeDirectoryBlobIndex':
        return CodeDirectoryBlobIndex._make(struct.unpack_from(CodeDirectoryBlobIndex.__format(), buf, 0))

    @staticmethod
    def bytesize() -> int:
        return struct.calcsize(CodeDirectoryBlobIndex.__format())

    @staticmethod
    def __format() -> str:
        return ">II"

class CodeDirectorySuperBlob(typing.NamedTuple):
    magic: int
    length: int
    count: int
    blob_indices: typing.List[CodeDirectoryBlobIndex]

    @staticmethod
    def make(buf: memoryview) -> 'CodeDirectorySuperBlob':
        super_blob_layout = ">III"
        super_blob = struct.unpack_from(super_blob_layout, buf, 0)

        offset = struct.calcsize(super_blob_layout)
        blob_indices = []
        for idx in range(super_blob[2]):
            blob_indices.append(CodeDirectoryBlobIndex.make(buf[offset:]))
            offset += CodeDirectoryBlobIndex.bytesize()

        return CodeDirectorySuperBlob(*super_blob, blob_indices)

def unpack_null_terminated_string(buf: memoryview) -> str:
    b = bytes(itertools.takewhile(lambda b: b != 0, buf))
    return b.decode()

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('binary', type=argparse.FileType('rb'), help='The file to analyze')
    parser.add_argument('offset', type=int, help='Offset to start of Code Directory data')
    parser.add_argument('size', type=int, help='Size of Code Directory data')
    parser.add_argument('code_offset', type=int, help='Offset to start of code pages to hash')
    parser.add_argument('code_size', type=int, help='Size of the code pages to hash')

    args = parser.parse_args()

    args.binary.seek(args.offset)
    super_blob_bytes = args.binary.read(args.size)
    super_blob_mem = memoryview(super_blob_bytes)

    super_blob = CodeDirectorySuperBlob.make(super_blob_mem)
    print(super_blob)

    for blob_index in super_blob.blob_indices:
        code_directory_offset = blob_index.offset
        code_directory = CodeDirectory.make(super_blob_mem[code_directory_offset:])
        print(code_directory)

        ident_offset = code_directory_offset + code_directory.identOffset
        print("Code Directory ID: " + unpack_null_terminated_string(super_blob_mem[ident_offset:]))

        code_offset = args.code_offset
        code_end = code_offset + args.code_size
        page_size = 1 << code_directory.pageSize
        args.binary.seek(code_offset)

        hashes_offset = code_directory_offset + code_directory.hashOffset
        for idx in range(code_directory.nCodeSlots):
            hash_bytes = bytes(super_blob_mem[hashes_offset:hashes_offset+code_directory.hashSize])
            hashes_offset += code_directory.hashSize

            hasher = hashlib.sha256()
            read_size = min(page_size, code_end - code_offset)
            hasher.update(args.binary.read(read_size))
            calculated_hash_bytes = hasher.digest()
            code_offset += read_size

            print("%s <> %s" % (hash_bytes.hex(), calculated_hash_bytes.hex()))

            if hash_bytes != calculated_hash_bytes:
                sys.exit(-1)


if __name__ == '__main__':
    main()