#!python

#
#  ktool | MAIN SCRIPT
#  ktool
#
#  This file is the main command-line script providing utilities for using ktool.
#
#  This file is part of ktool. ktool is free software that
#  is made available under the MIT license. Consult the
#  file "LICENSE" that is distributed together with this file
#  for the exact licensing terms.
#
#  Copyright (c) kat 2021.
#

import os
import pprint

from ktool import (
    Dyld,
    TBDGenerator,
    FatMachOGenerator,
    HeaderGenerator,
    MachOFileType,
    MachOFile,
    ObjCLibrary,
    TapiYAMLWriter,
    log,
    LogLevel
)

from kimg4.img4 import IM4P

from enum import Enum
from argparse import ArgumentParser, ArgumentError


class KToolError(Enum):
    ArgumentError = 1


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


def main():

    parser = 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_img4 = subparsers.add_parser('img4', help='img4/IM4P parsing utilities')

    parser_img4.add_argument('filename', nargs='?', default='')
    parser_img4.add_argument('--kbag', dest='get_kbag', action='store_true', help="Decode keybags in an im4p file")
    parser_img4.add_argument('--dec', dest='do_decrypt', action='store_true', help="Decrypt an im4p file with iv/key")
    parser_img4.add_argument('--iv', dest='aes_iv', type=str, help='IV for decryption')
    parser_img4.add_argument('--key', dest='aes_key', type=str, help='Key for decryption')
    parser_img4.add_argument('--out', dest='out', help="Output file destination for decryption")
    parser_img4.set_defaults(func=img4, get_kbag=False, do_decrypt=False, aes_iv=None, aes_key=None, out=None)

    parser_file = subparsers.add_parser('file', help='Print File Type (thin/fat MachO)')
    parser_file.add_argument('filename', nargs='?', default='')
    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='*', default='')
    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('filename', nargs='?', default='')
    parser_info.set_defaults(func=info, get_vm=False, get_lcs=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', nargs='?', default='')
    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('--binding', dest='get_binding', action='store_true', help="Print Binding Info Actions")
    parser_list.add_argument('filename', nargs='?', default='')
    parser_list.set_defaults(func=_list, get_syms=False, get_classes=False, get_protos=False, get_linked=False, get_binding=False)

    args = parser.parse_args()

    if not hasattr(args, 'filename'):
        # this is our default function
        args.func(args)
        exit()

    if not args.filename or args.filename == '':
        print(args.func.__doc__)
        exit()

    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):
    """ Usage: ktool [command] <flags> [filename]
    
    Commands for ktool:

    MachO ---
        dump - Tools to reconstruct certain files (headers, .tbds) from compiled MachOs
        file - Print very basic info about the MachO
        lipo - Utilities for combining/separating slices in fat MachO files.
        list - Dumps certain tables from MachOs
        info - Dump misc info about the target mach-o

    IMG4  ---
        img4 - IMG4 Utilities
    
    Run `ktool [command]`  for info/examples on using that command
        """
    print(help_prompt.__doc__)


def img4(args):
    """
    ktool img4
    IMG4 Utilities

    Getting keybags
        ktool img4 --kbag <filename>

    Decrypting an im4p
        ktool img4 --dec --iv AES_IV --key AES_KEY [--out <output-filename>] <filename>

    """
    if args.get_kbag:
        with open(args.filename, 'rb') as fp:
            im4p = IM4P(fp.read())
            for bag in im4p.kbag.keybags:
                print(f'{bag.iv.hex()}{bag.key.hex()}')
    if args.do_decrypt:
        if not args.aes_key or not args.aes_iv:
            exit_with_error(KToolError.ArgumentError, "--dec option requires --iv and --key")
        out = args.out
        if not out:
            out = args.filename + '.dec'
        with open(args.filename, 'rb') as fp:
            with open(out, 'wb') as out_fp:
                im4p = IM4P(fp.read())
                out_fp.write(im4p.decrypt_data(args.aes_iv, args.aes_key))

        print(f'Attempted decrypt of data with key/iv and saved to {out}')


def lipo(args):
    """
    ktool lipo
    Utilities for combining/separating slices in fat MachO files.

    Extract a slice from a fat binary

        ktool lipo --extract [slicename] [filename]

    Create a fat Macho Binary from multiple thin binaries

        ktool lipo --create [--out filename] [filenames]
    """
    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):
    """
    ktool list
    Outputs some lists of things

    To print the symbol table of a file

        ktool list --symbols [filename]

    To print the list of classes

        ktool list --classes [filename]

    To print the list of protocols

        ktool list --protocols [filename]

    To print a  list of linked libraries

        ktool list --linked [filename]

    To print binding actionms for a file

        ktool list [--slice n] --binding [filename]

    """
    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_binding:
            filt = True
            print('\nBinding Info Actions '.ljust(60, '-') + '\n')
            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
        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('(Weak) ' + exlib.install_name if exlib.weak else '' + exlib.install_name)


def _file(args):
    """
    ktool file
    ----------

    Print basic information about a file (e.g 'Thin MachO Binary')

        ktool file [filename]
    """
    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):
    """
    ktool info
    Some misc info about the target mach-o

    Print generic info about a MachO file

        ktool info [--slice n] [filename]

    Print VM -> Slice -> Filename address mapping for a slice
    of a MachO file

        ktool info [--slice n] --vm [filename]

    Print a list of load commands for a file

        ktool info [--slice n] --cmds [filename]
    """
    fd = open(args.filename, 'rb')
    machofile = MachOFile(fd)
    library = Dyld.load(machofile.slices[args.slice_index], load_symtab=False, load_binding=False)
    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 not filt:
        print(f'Name: {library.name}')
        print(f'Filetype: {library.macho_header.filetype.name}')
        print(f'Flags: {", ".join([i.name for i in library.macho_header.flags])}')
        print(f'UUID: {library.uuid.hex().upper()}')
        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):
    """
    ktool dump
    Tools to reconstruct certain files from compiled MachOs

    To dump a set of headers for a bin/framework

        ktool dump --headers --out <directory> [filename]

    To dump .tbd files for a framework

        ktool dump --tbd [filename]
    """
    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()
