#!/usr/bin/env python3

#
#  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
import threading

from kmacho import LOAD_COMMAND
from ktool import (
    Dyld,
    TBDGenerator,
    FatMachOGenerator,
    HeaderGenerator,
    MachOFileType,
    MachOFile,
    ObjCLibrary,
    TapiYAMLWriter,
    log,
    LogLevel, dylib_command_t, dylib, dylib_command
)

from ktool.util import KTOOL_VERSION

from kimg4.img4 import IM4P

from ktool.window import KToolScreen

from enum import Enum
from argparse import ArgumentParser
import json
import urllib.request
from packaging.version import Version

UPDATE_AVAILABLE = False


def check_for_update():
    endpoint = "https://pypi.org/pypi/k2l/json"
    try:
        with urllib.request.urlopen(endpoint, timeout=1) as url:
            data = json.loads(url.read().decode(), strict=False)
        if Version(KTOOL_VERSION) < Version(data['info']['version']):
            global UPDATE_AVAILABLE
            UPDATE_AVAILABLE = True
    except Exception:
        pass


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 require_args(args, always=None, one_of=None):
    if always:
        for i in always:
            if not hasattr(args, i):
                print(args.func.__doc__)
                exit_with_error(KToolError.ArgumentError, f'Missing required argument {i}')
            elif not getattr(args, i):
                print(args.func.__doc__)
                exit_with_error(KToolError.ArgumentError, f'Missing required argument {i}')

    if one_of:
        found_one = False
        for i in one_of:
            if hasattr(args, i):
                if getattr(args, i):
                    found_one = True
                    break
        if not found_one:
            print(args.func.__doc__)
            exit_with_error(KToolError.ArgumentError, f'Missing one of [{",".join(one_of)}]')


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_open = subparsers.add_parser('open', help='open ktool GUI and browse file')
    parser_open.add_argument('filename', nargs='?', default='')
    parser_open.set_defaults(func=_open)

    parser_insert = subparsers.add_parser('insert', help='Insert data into MachO Binary')
    parser_insert.add_argument('filename', nargs='?', default='')
    parser_insert.add_argument('--lc', dest='lc', help="Type of Load Command to insert")
    parser_insert.add_argument('--payload', dest='payload', help="Payload (if required) for insertion")
    parser_insert.add_argument('--out', dest='out', help="Output file destination for patches")
    parser_insert.set_defaults(func=insert, out=None, lc=None, payload=None)

    parser_edit = subparsers.add_parser('edit', help='Edit attributes of the MachO')
    parser_edit.add_argument('filename', nargs='?', default='')
    parser_edit.add_argument('--iname', dest='iname', help='Modify the Install Name of a library')
    parser_edit.add_argument('--apad', dest='apad', help='Add MachO Header Padding (You probably dont want to do this manually)')
    parser_edit.add_argument('--out', dest='out', help="Output file destination for patches")
    parser_edit.set_defaults(func=edit, out=None, iname=None, apad=None)

    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('--sorted', dest='sort_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, sort_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()

    download_thread = threading.Thread(target=check_for_update, name="UpdateChecker")
    download_thread.start()

    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)

    if UPDATE_AVAILABLE:
        print(f'\n\nUpdate Available ---')
        print(f'run `pip3 install --upgrade k2l` to fetch the latest update')

    exit(0)


def help_prompt(args):
    """Usage: ktool [command] <flags> [filename]

Commands:

GUI (Still in active development) ---
    ktool open [filename] - Open the ktool command line GUI and browse a file

MachO Editing ---
    insert - Utils for inserting load commands into MachO Binaries
    edit - Utils for editing MachO Binaries
    lipo - Utilities for combining/separating slices in fat MachO files.

MachO Analysis ---
    dump - Tools to reconstruct certain files (headers, .tbds) from compiled MachOs
    file - Print very basic info about the MachO
    list - Dumps certain tables from MachOs
    info - Dump misc info about the target mach-o

Misc Utilities ---
    img4 - IMG4 Utilities

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


def _open(args):
    log.LOG_LEVEL = LogLevel.NONE
    screen = KToolScreen()
    screen.load_file(args.filename)


def insert(args):
    """
    ktool insert
    Utils for inserting load commands into mach-o binaries

    insert a LOAD_DYLIB command

        ktool insert --lc load --payload /Dylib/Install/Name/Here.dylib --out <output filename> [filename]

    commands currently supported:
        load: LOAD_DYLIB
        load-weak: LOAD_WEAK_DYLIB
        lazy-load: LAZY_LOAD_DYLIB
        load-upward: LOAD_UPWARD_DYLIB
    """

    require_args(args, always=['lc'])

    lc = None
    if args.lc == "load":
        lc = LOAD_COMMAND.LOAD_DYLIB
    elif args.lc == "load-weak" or args.lc == "load_weak":
        lc = LOAD_COMMAND.LOAD_WEAK_DYLIB
    elif args.lc in ["load_lazy", "load-lazy", "lazy-load", "lazy_load"]:
        lc = LOAD_COMMAND.LAZY_LOAD_DYLIB
    elif args.lc == "load-upward" or args.lc == "load_upward":
        lc = LOAD_COMMAND.LOAD_UPWARD_DYLIB

    patched_libraries = []

    with open(args.filename, 'rb') as fp:
        machofile = MachOFile(fp)
        for macho_slice in machofile.slices:
            library = Dyld.load(macho_slice)
            last_dylib_command_index = -1
            for i, cmd in enumerate(library.macho_header.load_commands):
                if isinstance(cmd, dylib_command):
                    last_dylib_command_index = i+1
            dylib_item = dylib.assemble()
            library.insert_lc_with_suf(dylib_command_t, lc, [dylib_item.raw], args.payload, last_dylib_command_index)
            library = Dyld.load(macho_slice)
            patched_libraries.append(library)

    with open(args.out, 'wb') as fd:
        if len(patched_libraries) > 1:
            slices = [library.slice for library in patched_libraries]
            fat_generator = FatMachOGenerator(slices)
            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())
        else:
            fd.write(patched_libraries[0].slice.full_bytes_for_slice())


def edit(args):
    """
    ktool edit
    Utils for editing MachO Binaries

    Modify the install name of a library

        ktool edit --iname [Desired Install Name] --out <Output Filename> [filename]

    """

    require_args(args, one_of=['iname', 'apad'])

    patched_libraries = []

    if args.iname:

        new_iname = args.iname

        with open(args.filename, 'rb') as fp:
            machofile = MachOFile(fp)
            for macho_slice in machofile.slices:
                library = Dyld.load(macho_slice)
                id_dylib_index = -1
                for i, cmd in enumerate(library.macho_header.load_commands):
                    if cmd.cmd == 0xD:
                        id_dylib_index = i
                        break
                dylib_item = dylib.assemble()
                library.rm_load_command(id_dylib_index)
                library = Dyld.load(macho_slice)
                library.insert_lc_with_suf(dylib_command_t, LOAD_COMMAND.ID_DYLIB, [dylib_item.raw], new_iname, id_dylib_index)
                library = Dyld.load(macho_slice)
                patched_libraries.append(library)

        with open(args.out, 'wb') as fd:
            if len(patched_libraries) > 1:
                slices = [library.slice for library in patched_libraries]
                fat_generator = FatMachOGenerator(slices)
                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())
            else:
                fd.write(patched_libraries[0].slice.full_bytes_for_slice())

    elif args.apad:
        add_padding = int(args.apad)
        with open(args.filename, 'rb') as fp:
            machofile = MachOFile(fp)
            library = Dyld.load(machofile.slices[0])
            library.increase_header_pad(add_padding)
            patched_libraries.append(library)

        with open(args.out, 'wb') as fd:
            if len(patched_libraries) > 1:
                slices = [library.slice for library in patched_libraries]
                fat_generator = FatMachOGenerator(slices)
                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())
            else:
                if patched_libraries[0].slice.use_patched_bytes:
                    print('using patched bytes')
                    fd.seek(0)
                    fd.write(patched_libraries[0].slice.patched_bytes)
                else:
                    fd.write(patched_libraries[0].slice.full_bytes_for_slice())


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>

    """

    require_args(args, one_of=['get_kbag', 'do_decrypt'])

    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]
    """

    require_args(args, one_of=['combine', 'extract'])

    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 actions for a file

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

    """

    require_args(args, one_of=['get_syms', 'get_binding', 'get_classes', 'get_protos', 'get_linked'])

    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]
    """

    require_args(args, one_of=['do_headers', 'do_tbd'])

    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.sort_headers:
            for objc_class in objc_lib.classlist:
                objc_class.methods = sorted(objc_class.methods, key=lambda h: h.signature)
                objc_class.properties = sorted(objc_class.properties, key=lambda h: h.name)

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