#!python

import argparse
import os
import pprint

from ktool.dyld import Dyld
from ktool.generator import TBDGenerator, FatMachOGenerator
from ktool.headers import HeaderGenerator
from ktool.macho import MachOFileType, MachOFile
from ktool.objc import ObjCLibrary
from ktool.util import TapiYAMLWriter, log, LogLevel
from enum import Enum


class KToolError(Enum):
    ArgumentError = 1


def exit_with_error(error: KToolError, msg):
    print(f'Encountered an Error ({error.name})')
    print(f'{msg}')
    exit(error.value)


def main():

    parser = argparse.ArgumentParser(description="ktool")
    parser.add_argument('--bench', dest='bench', action='store_true')
    parser.add_argument('-v', dest='logging_level', type=int)
    parser.set_defaults(func=help_prompt, bench=False, logging_level=0)
    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_lipo = subparsers.add_parser('lipo', help='Extract/Combine slices')
    parser_lipo.add_argument('--extract', dest='extract', type=str, help='Extract a slice (--extract arm64)')
    parser_lipo.add_argument('--out', dest='out', help="Output File")
    parser_lipo.add_argument('--create', dest='combine', action='store_true', help="Combine files to create a fat mach-o library")
    parser_lipo.add_argument('filename', nargs='+')
    parser_lipo.set_defaults(func=lipo, out="", combine=False)

    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()
    log.LOG_LEVEL = LogLevel(args.logging_level)

    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 lipo:
ktool lipo --extract [slicename] [filename] - Extract a slice from a fat binary
ktool lipo --create [--out filename] [filenames] - Create a fat MachO Binary from multiple thin binaries

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 lipo(args):
    if args.combine:
        output = args.out
        if output == "":
            output = args.filename[0] + '.fat'
        slices = []
        for filename in args.filename:
            # Slice() might hold a ref preventing it from being closed? but i'm just going to let it close on exit()
            fd = open(filename, 'rb')
            macho_file = MachOFile(fd)
            if macho_file.type != MachOFileType.THIN:
                exit_with_error(KToolError.ArgumentError, "Fat mach-o passed to --create")
            slices.append(macho_file.slices[0])

        fat_generator = FatMachOGenerator(slices)

        with open(output, 'wb') as fd:
            fd.write(fat_generator.fat_head)
            for arch in fat_generator.fat_archs:
                fd.seek(arch.offset)
                fd.write(arch.slice.full_bytes_for_slice())

    elif args.extract != "":
        with open(args.filename[0], 'rb') as fd:
            macho_file = MachOFile(fd)
            output = args.out
            if output == "":
                output = args.filename[0] + '.' + args.extract.lower()
            for slice in macho_file.slices:
                if slice.type.name.lower() == args.extract:
                    with open(output, 'wb') as out:
                        out.write(slice.full_bytes_for_slice())
                    return
            exit_with_error(KToolError.ArgumentError,
                            f'Architecture {args.extract} wasn\'t found (found: {[slice.type.name.lower() for slice in macho_file.slices]})')


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:
                pass # 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:
            exit_with_error(KToolError.ArgumentError, "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()
