#!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 json
import os
import os.path
import threading
import urllib.request
from argparse import ArgumentParser
from collections import namedtuple
from enum import Enum

from kimg4.img4 import IM4P
from packaging.version import Version

from kmacho import LOAD_COMMAND, dylib_command, dylib, Struct
from ktool import (
    Dyld,
    TBDGenerator,
    FatMachOGenerator,
    HeaderGenerator,
    MachOFileType,
    MachOFile,
    ObjCLibrary,
    TapiYAMLWriter,
    log,
    LogLevel
)
from ktool.exceptions import *
from ktool.util import KTOOL_VERSION
from ktool.window import KToolScreen, external_hard_fault_teardown, Table

UPDATE_AVAILABLE = False


def get_terminal_size():
    # We use this instead of shutil.get_terminal_size, because when output is being piped, it returns column width 80
    # We want to make sure if output is being piped (for example, to grep), that no wrapping occurs, so greps will
    #   always display all relavant info on a single line. This also helps if it's being piped into a file, for processing
    #   purposes among everything else.
    try:
        return os.get_terminal_size()
    except OSError:
        return os.terminal_size((5000, 5000))


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
    FiletypeError = 2


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):
    """
    This is a quick macro to enforce argument requirements for different commands.

    If a check fails, it'll print usage for the subcommand and exit the program.

    :param args: Parsed argument object
    :param always: Arguments that *must* be passed
    :param one_of: At least one of these arguments must be passed, and must evaluate as True
    :return:
    """
    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=1)
    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_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_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_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('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('--slice', dest='slice_index', type=int,
                             help="Specify Index of Slice (in FAT MachO) to examine")
    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('--cmds', dest='get_lcs', action='store_true', help="Print Load Commands")
    parser_list.add_argument('filename', nargs='?', default='')
    parser_list.set_defaults(func=_list, get_lcs=False, get_classes=False, get_protos=False, get_linked=False,
                             slice_index=0)

    parser_symbols = subparsers.add_parser('symbols', help='Print various symbols')
    parser_symbols.add_argument('--imports', dest='get_imports', action='store_true', help='Print Imports')
    parser_symbols.add_argument('--imp-acts', dest='get_actions', action='store_true', help='Print Raw Binding Imports')
    parser_symbols.add_argument('--symtab', dest='get_symtab', action='store_true', help='Print out the symtab')
    parser_symbols.add_argument('--exports', dest='get_exports', action='store_true', help='Print exports')
    parser_symbols.add_argument('filename', nargs='?', default='')
    parser_symbols.set_defaults(func=symbols, get_imports=False, get_actions=False, get_exports=False, get_symtab=False,
                                slice_index=0)

    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:
        try:
            args.func(args)
        except UnsupportedFiletypeException:
            exit_with_error(KToolError.FiletypeError, f'{args.filename} is not a valid MachO Binary')
        except FileNotFoundError:
            exit_with_error(KToolError.ArgumentError, f'{args.filename} doesn\'t exist')

    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
    list - Print various lists (Classlist, etc.)
    symbols - Print various tables (Symbols, binding info)
    info - Print misc info about the target mach-o

Misc Utilities ---
    file - Print very basic info about the MachO
    img4 - IMG4 Utilities

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


def _open(args):
    try:
        log.LOG_LEVEL = LogLevel.NONE
        screen = KToolScreen()
        screen.load_file(args.filename)
    except KeyboardInterrupt:
        external_hard_fault_teardown()
        print('Hard Faulted. This was likely due to a curses error causing a freeze while rendering.')


def symbols(args):
    """
    ----------
    List symbol imports/exports

    Print the list of imported symbols

        ktool symbols --imports [filename]

    Print the list of exported symbols

        ktool symbols --exports [filename]

    Print the symbol table

        ktool symbols --symtab [filename]
    """

    require_args(args, one_of=['get_imports', 'get_actions', 'get_exports', 'get_symtab'])

    if args.get_exports:
        with open(args.filename, 'rb') as fd:
            machofile = MachOFile(fd)
            library = Dyld.load(machofile.slices[args.slice_index])

            table = Table()
            table.titles = ['Address', 'Symbol']

            for i in library.exports.nodes:
                table.rows.append([hex(i.offset), i.text])

            print(table.render(get_terminal_size().columns))

    if args.get_symtab:
        with open(args.filename, 'rb') as fd:
            machofile = MachOFile(fd)
            library = Dyld.load(machofile.slices[args.slice_index])

            table = Table()
            table.titles = ['Address', 'Name']

            for sym in library.symbol_table.table:
                table.rows.append([hex(sym.addr), sym.fullname])

            print(table.render(get_terminal_size().columns))

    if args.get_imports:
        with open(args.filename, 'rb') as fd:
            machofile = MachOFile(fd)
            library = Dyld.load(machofile.slices[args.slice_index])

            syms = {}
            symbol = namedtuple('symbol', ['name', 'library', 'from_table'])

            for sym in library.binding_table.symbol_table:
                try:
                    syms[sym.fullname] = symbol(sym.fullname, library.linked[int(sym.ordinal) - 1].install_name, '')
                except:
                    pass
            for sym in library.weak_binding_table.symbol_table:
                try:
                    syms[sym.fullname] = symbol(sym.fullname, library.linked[int(sym.ordinal) - 1].install_name, 'Weak')
                except:
                    pass
            for sym in library.lazy_binding_table.symbol_table:
                try:
                    syms[sym.fullname] = symbol(sym.fullname, library.linked[int(sym.ordinal) - 1].install_name, 'Lazy')
                except:
                    pass

            table = Table()
            table.titles = ['Symbol', 'Library', 'Type']

            for _, sym in syms.items():
                table.rows.append([sym.name, sym.library, sym.from_table])

            print(table.render(get_terminal_size().columns))

    elif args.get_actions:
        with open(args.filename, 'rb') as fd:
            machofile = MachOFile(fd)
            library = Dyld.load(machofile.slices[args.slice_index])
            print('\nBinding Info'.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
            print('\nWeak Binding Info'.ljust(60, '-') + '\n')
            for sym in library.weak_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('\nLazy Binding Info'.ljust(60, '-') + '\n')
            for sym in library.lazy_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


def insert(args):
    """
    ----------
    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 = Struct.create_with_values(dylib, [0x18, 0x2, 0x010000, 0x010000])
            library.insert_lc_with_suf(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):
    """
    ----------
    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(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)


def img4(args):
    """
    ----------
    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):
    """
    ----------
    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):
    """
    ----------
    Tools for printing various lists

    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 a list of Load Commands and their data

        ktool list --cmds [filename]

    """

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

    with open(args.filename, 'rb') as fd:
        machofile = MachOFile(fd)
        library = Dyld.load(machofile.slices[args.slice_index])
        print(f'\n{args.filename} '.ljust(60, '-') + '\n')
        if args.get_lcs:
            for lc in library.macho_header.load_commands:
                print(str(lc))
        elif args.get_classes:
            objc_lib = ObjCLibrary(library)
            for obj_class in objc_lib.classlist:
                print(f'{obj_class.name}')
        elif args.get_protos:
            objc_lib = ObjCLibrary(library)
            for objc_proto in objc_lib.protolist:
                print(f'{objc_proto.name}')
        elif 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]
    """
    with open(args.filename, 'rb') as fp:
        machofile = MachOFile(fp)
        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):
    """
    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]

    """
    with open(args.filename, 'rb') as fp:
        machofile = MachOFile(fp)
        library = Dyld.load(machofile.slices[args.slice_index], load_symtab=False, load_binding=False)
        if args.get_vm:
            print(library.vm)
        else:
            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}')



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:
        with open(args.filename, 'rb') as fp:
            macho_file = MachOFile(fp)
            library = Dyld.load(macho_file.slices[args.slice_index])
            if library.name == "":
                library.name = os.path.basename(args.filename)
            objc_lib = ObjCLibrary(library)

            if args.sort_headers:
                for objc_class in objc_lib.classlist:
                    objc_class.methods.sort(key=lambda h: h.signature)
                    objc_class.properties.sort(key=lambda h: h.name)
                    if objc_class.metaclass is not None:
                        objc_class.metaclass.methods.sort(key=lambda h: h.signature)

                for objc_proto in objc_lib.protolist:
                    objc_proto.methods.sort(key=lambda h: h.signature)
                    objc_proto.opt_methods.sort(key=lambda h: h.signature)

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

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


if __name__ == "__main__":
    main()
