#!/usr/bin/env python3

import argparse
import os
import pprint

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


def main():
    parser = argparse.ArgumentParser(description="ktool")
    parser.add_argument('--bench', dest='bench', action='store_true')
    parser.set_defaults(func=help_prompt, bench=False)
    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)

    parser_list = subparsers.add_parser('list', help='Print various lists')
    parser_list.add_argument('--symbols', dest='get_syms', action='store_true', help='Print symbol list')
    parser_list.add_argument('--classes', dest='get_classes', action='store_true', help='Print class list')
    parser_list.add_argument('--protocols', dest='get_protos', action='store_true', help='Print Protocol list')
    parser_list.add_argument('--linked', dest='get_linked', action='store_true', help='Print list of linked libraries')
    parser_list.add_argument('filename')
    parser_list.set_defaults(func=list, get_syms=False, get_classes=False, get_protos=False, get_linked=False)

    args = parser.parse_args()

    if args.bench:
        import cProfile
        import pstats

        profile = cProfile.Profile()
        profile.runcall(args.func, args)
        ps = pstats.Stats(profile)
        ps.sort_stats('time', 'cumtime')
        ps.print_stats(10)
    else:
        args.func(args)


def help_prompt(args):
    string = """usage: ktool [command] <flags> [filename]

ktool dump:
ktool dump --headers --out <directory> [filename] - Dump set of headers for a bin/framework
ktool dump --tbd [filename] - Dump .tbd for a framework

ktool file:
ktool file [filename] - Prints (very) basic info about a file (e.g. "Thin MachO Binary")

ktool list:
ktool list --symbols [filename] - Print the symbol table for the file
ktool list --classes [filename] - Print the list of classes
ktool list --protocols [filename] - Print the list of protocols
ktool list --linked [filename] - Print a list of linked libraries

ktool info:
usage: ktool info [-h] [--slice SLICE_INDEX] [--vm] [--cmds] [--binding] filename
ktool info [--slice n] [filename] - Print generic info about a MachO File
ktool info [--slice n] --vm [filename] - Print VM -> Slice -> File address mapping for a slice of a MachO File
ktool info [--slice n] --cmds [filename] - Print list of load commands for a file 
ktool info [--slice n] --binding [filename] - Print binding actions for a file"""
    print(string)


def list(args):
    with open(args.filename, 'rb') as fd:
        machofile = MachOFile(fd)
        library = Dyld.load(machofile.slices[0])
        print(f'\n{args.filename} '.ljust(60, '-') + '\n')
        if args.get_syms:
            for sym in library.symbol_table.table:
                print(f'Address: {sym.addr} | Name: {sym.fullname}')
        if args.get_classes:
            objc_lib = ObjCLibrary(library)
            for obj_class in objc_lib.classlist:
                print(f'{obj_class.name}')
        if args.get_protos:
            objc_lib = ObjCLibrary(library)
            for objc_proto in objc_lib.protolist:
                print(f'{objc_proto.name}')
        if args.get_linked:
            for exlib in library.linked:
                print(exlib.install_name)




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')
    fd.close()


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:
            try:
                print(
                    f'{hex(sym.addr).ljust(15, " ")} | {library.linked[int(sym.ordinal) - 1].install_name} | {sym.name.ljust(20, " ")} | {sym.type}')
            except:
                print(f'{int(sym.ordinal)} symbol ordinal broken')

    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()
