#!/usr/bin/env python3

import mmap
import os
import argparse
import pprint
from kdump.headers import *
from kdump.classlist import *
from kdump.sdkbuilder import *


def main():
    parser = argparse.ArgumentParser(description="kdump")
    parser.add_argument('--map', dest='do_map', action='store_true')
    parser.add_argument('--classes', dest='do_classes', action='store_true')
    parser.add_argument('--binding', dest='do_binding', action='store_true')
    parser.add_argument('--class', dest='getclass')
    parser.add_argument('--sdk', dest='cache_location')
    parser.add_argument('--sdkout', dest='sdk_out')
    parser.add_argument('--headers', dest='do_headers', action='store_true')
    parser.add_argument('--out', dest='outdir')
    parser.add_argument('filename')
    parser.set_defaults(do_map=False, do_classes=False, do_binding=False, getclass=None, outdir=None, do_headers=False, cache_location=None, sdk_out=None)
    args = parser.parse_args()

    if args.cache_location:
        assert args.sdk_out is not None
        cachein = args.cache_location
        cacheout = args.sdk_out
        SDKBuilder(cachein, cacheout)
        return

    file = None
    print('\r[*] Loading File...', end="")
    fd = open(args.filename, "rb")
    file = mmap.mmap(fd.fileno(), 0, access=mmap.ACCESS_READ)
    print('\r[+] Loading File...')

    print('\r[*] Processing MachO Structure...', end="")
    macho = MachO(file)
    print('\r[+] Processing MachO Structure...')

    if macho.name == "":
        macho.name = args.filename

    if args.do_binding:
        pprint.pprint(macho.binding_actions)

    if args.do_map:
        print(macho.vm)

    print('\r[*] Processing Objective C Classes...', end="")
    classlist = Classlist(macho, args.getclass)
    print('\r[+] Processing Objective C Classes...')

    if args.do_classes:
        print("--------")
        print("Classes:")
        for objcclass in classlist.list:
            print(objcclass)
        print("--------")

    if args.getclass is not None:
        for objcclass in classlist.list:
            if str(objcclass) == args.getclass:
                print(Header(macho, classlist, objcclass))

    if args.do_headers:
        print('\r[*] Producing Output Headers...', end="")
        if args.outdir is None:
            raise AssertionError("Missing --out flag (--out <directory>), specifies directory to place headers")
        generator = HeaderGenerator(macho, classlist)
        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))

        print('\r[+] Producing Output Headers...')
    print('\n[+] Program Finished')
    file.close()
    fd.close()


if __name__ == "__main__":
    #import cProfile
    #import pstats

    #profile = cProfile.Profile()
    #profile.runcall(main)
    #ps = pstats.Stats(profile)
    #ps.sort_stats('cumtime', 'calls')
    #ps.print_stats(10)
    main()