#!/usr/bin/env python3

import argparse
import os
import pprint

from ktool.dyld import Dyld
from ktool.generator import HeaderGenerator, TBDGenerator
from ktool.macho import MachOFileType, MachOFile
from ktool.objc import ObjCLibrary
from ktool.util import TapiYAMLWriter


def main():
    parser = argparse.ArgumentParser(description="ktool")
    subparsers = parser.add_subparsers(help='sub-command help')

    parser_file = subparsers.add_parser('file', help='Print File Type (thin/fat MachO)')
    parser_file.add_argument('filename')
    parser_file.set_defaults(func=file)

    parser_info = subparsers.add_parser('info', help='Print Info about a MachO Library')
    parser_info.add_argument('--slice', dest='slice_index', type=int, help="Specify Index of Slice (in FAT MachO) to examine")
    parser_info.add_argument('--vm', dest='get_vm', action='store_true', help="Print VM Mapping for MachO Library")
    parser_info.add_argument('--cmds', dest='get_lcs', action='store_true', help="Print Load Commands")
    parser_info.add_argument('--binding', dest='get_binding', action='store_true', help="Print Binding Info Actions")
    parser_info.add_argument('filename')
    parser_info.set_defaults(func=info, get_vm=False, get_lcs=False, get_binding=False, slice_index=0)

    parser_dump = subparsers.add_parser('dump', help='Dump items (headers) from binary')
    parser_dump.add_argument('--slice', dest='slice_index', type=int, help="Specify Index of Slice (in FAT MachO) to examine")
    parser_dump.add_argument('--headers', dest='do_headers', action='store_true')
    parser_dump.add_argument('--tbd', dest='do_tbd', action='store_true')
    parser_dump.add_argument('--out', dest='outdir', help="Directory to dump headers into")
    parser_dump.add_argument('filename')
    parser_dump.set_defaults(func=dump, do_headers=False, do_tbd=False, slice_index=0)

    args = parser.parse_args()

    args.func(args)


def file(args):
    fd = open(args.filename, 'rb')
    machofile = MachOFile(fd)
    print(f'\n{args.filename} '.ljust(60, '-') + '\n')

    if machofile.type == MachOFileType.FAT:
        print('Fat MachO Binary')
        print(f'{len(machofile.slices)} Slices:')

        print(f'{"Offset".ljust(15, " ")} | {"CPU Type".ljust(15, " ")} | {"CPU Subtype".ljust(15, " ")}')
        for slice in machofile.slices:
            print(f'{hex(slice.offset).ljust(15, " ")} | {slice.type.name.ljust(15, " ")} | {slice.subtype.name.ljust(15, " ")}')
    else:
        print('Thin MachO Binary')


def info(args):
    fd = open(args.filename, 'rb')
    machofile = MachOFile(fd)
    library = Dyld.load(machofile.slices[args.slice_index])
    filt = False
    if args.get_vm:
        print(library.vm)
        filt = True
    if args.get_lcs:
        pprint.pprint(library.macho_header.load_commands)
        filt = True
    if args.get_binding:
        filt = True
        print('\nBinding Info Actions '.ljust(60, '-') + '\n')
        print(library.linked)
        for sym in library.binding_table.symbol_table:
            print(f'{hex(sym.addr).ljust(15, " ")} | {library.linked[int(sym.ordinal)-1].install_name} | {sym.name.ljust(20, " ")} | {sym.type}')

    if not filt:
        print(f'Name: {library.name}')
        print(f'UUID: {hex(library.uuid)}')
        print(f'Platform: {library.platform.name}')
        print(f'Minimum OS: {library.minos.x}.{library.minos.y}.{library.minos.z}')
        print(f'SDK Version: {library.sdk_version.x}.{library.sdk_version.y}.{library.sdk_version.z}')


    fd.close()


def dump(args):
    if args.do_headers:
        fd = open(args.filename, 'rb')
        machofile = MachOFile(fd)
        library = Dyld.load(machofile.slices[args.slice_index])
        if library.name == "":
            library.name = args.filename
        objc_lib = ObjCLibrary(library)

        if args.outdir is None:
            raise AssertionError("Missing --out flag (--out <directory>), specifies directory to place headers")
        generator = HeaderGenerator(objc_lib)
        for header_name, header in generator.headers.items():
            if args.outdir == "kdbg":  # something i can put into IDE args that wont accidentally get used by a user
                print('\n\n')
                print(header_name)
                print()
                print(header)
            else:
                os.makedirs(args.outdir, exist_ok=True)
                with open(args.outdir + '/' + header_name, 'w') as out:
                    out.write(str(header))

        fd.close()

    if args.do_tbd:
        fd = open(args.filename, 'rb')
        machofile = MachOFile(fd)
        library = Dyld.load(machofile.slices[args.slice_index])
        tbdgen = TBDGenerator(library, True)
        with open(library.name + '.tbd', 'w') as filen:
            filen.write(TapiYAMLWriter.write_out(tbdgen.dict))
        fd.close()


if __name__ == "__main__":
    main()
