# -*- coding: utf-8 -*-
from __future__ import unicode_literals, print_function

import binascii
import codecs
import os
import platform
import re
import subprocess
import threading
import time

from frida import FileMonitor

from frida_tools.model import Module, Function, ModuleFunction, ObjCMethod


def main():
    from colorama import Fore, Style
    import json

    from frida_tools.application import ConsoleApplication, input_with_timeout

    class TracerApplication(ConsoleApplication, UI):
        def __init__(self):
            super(TracerApplication, self).__init__(self._await_ctrl_c)
            self._palette = [Fore.CYAN, Fore.MAGENTA, Fore.YELLOW, Fore.GREEN, Fore.RED, Fore.BLUE]
            self._next_color = 0
            self._attributes_by_thread_id = {}
            self._last_event_tid = -1

        def _add_options(self, parser):
            pb = TracerProfileBuilder()
            def process_builder_arg(option, opt_str, value, parser, method, **kwargs):
                method(value)
            parser.add_option("-I", "--include-module", help="include MODULE", metavar="MODULE",
                    type='string', action='callback', callback=process_builder_arg, callback_args=(pb.include_modules,))
            parser.add_option("-X", "--exclude-module", help="exclude MODULE", metavar="MODULE",
                    type='string', action='callback', callback=process_builder_arg, callback_args=(pb.exclude_modules,))
            parser.add_option("-i", "--include", help="include FUNCTION", metavar="FUNCTION",
                    type='string', action='callback', callback=process_builder_arg, callback_args=(pb.include,))
            parser.add_option("-x", "--exclude", help="exclude FUNCTION", metavar="FUNCTION",
                    type='string', action='callback', callback=process_builder_arg, callback_args=(pb.exclude,))
            parser.add_option("-a", "--add", help="add MODULE!OFFSET", metavar="MODULE!OFFSET",
                    type='string', action='callback', callback=process_builder_arg, callback_args=(pb.include_relative_address,))
            parser.add_option("-T", "--include-imports", help="include program's imports",
                    action='callback', callback=process_builder_arg, callback_args=(pb.include_imports,))
            parser.add_option("-t", "--include-module-imports", help="include MODULE imports", metavar="MODULE",
                    type='string', action='callback', callback=process_builder_arg, callback_args=(pb.include_imports,))
            parser.add_option("-m", "--include-objc-method", help="include OBJC_METHOD", metavar="OBJC_METHOD",
                    type='string', action='callback', callback=process_builder_arg, callback_args=(pb.include_objc_method,))
            parser.add_option("-M", "--exclude-objc-method", help="exclude OBJC_METHOD", metavar="OBJC_METHOD",
                    type='string', action='callback', callback=process_builder_arg, callback_args=(pb.exclude_objc_method,))
            parser.add_option("-s", "--include-debug-symbol", help="include DEBUG_SYMBOL", metavar="DEBUG_SYMBOL",
                    type='string', action='callback', callback=process_builder_arg, callback_args=(pb.include_debug_symbol,))
            parser.add_option("-q", "--quiet", help="do not format output messages", action='store_true', default=False)
            parser.add_option("-d", "--decorate", help="add module name to generated onEnter log statement", action='store_true', default=False)
            parser.add_option("-S", "--init-session", help="path to JavaScript file used to initialize the session", metavar="PATH",
                    type='string', action='append', default=[])
            parser.add_option("-P", "--parameters", help="parameters as JSON, exposed as a global named 'parameters'", metavar="PARAMETERS_JSON",
                    type='string', action='store', default=None)
            parser.add_option("-o", "--output", help="dump messages to file", metavar="OUTPUT", type='string')
            self._profile_builder = pb

        def _usage(self):
            return "usage: %prog [options] target"

        def _initialize(self, parser, options, args):
            self._tracer = None
            self._targets = None
            self._profile = self._profile_builder.build()
            self._quiet = options.quiet
            self._decorate = options.decorate
            self._output = None
            self._output_path = options.output

            self._init_scripts = []
            for path in options.init_session:
                with codecs.open(path, 'rb', 'utf-8') as f:
                    source = f.read()
                self._init_scripts.append(InitScript(path, source))

            if options.parameters is not None:
                try:
                    params = json.loads(options.parameters)
                except Exception as e:
                    raise ValueError("failed to parse parameters argument as JSON: {}".format(e))
                if not isinstance(params, dict):
                    raise ValueError("failed to parse parameters argument as JSON: not an object")
                self._parameters = params
            else:
                self._parameters = {}

        def _needs_target(self):
            return True

        def _start(self):
            if self._output_path is not None:
                self._output = OutputFile(self._output_path)

            stage = 'early' if self._target[0] == 'file' else 'late'

            self._tracer = Tracer(self._reactor, FileRepository(self._reactor, self._decorate), self._profile,
                    self._init_scripts, log_handler=self._log)
            try:
                self._targets = self._tracer.start_trace(self._session, stage, self._parameters, self._runtime, self)
            except Exception as e:
                self._update_status("Failed to start tracing: {error}".format(error=e))
                self._exit(1)

        def _stop(self):
            self._tracer.stop()
            self._tracer = None
            if self._output is not None:
                self._output.close()
            self._output = None

        def _await_ctrl_c(self, reactor):
            while reactor.is_running():
                try:
                    input_with_timeout(0.5)
                except KeyboardInterrupt:
                    break

        def on_trace_progress(self, operation):
            if operation == 'resolve':
                self._update_status("Resolving functions...")
            elif operation == 'instrument':
                self._update_status("Instrumenting functions...")
            elif operation == 'ready':
                if len(self._targets) == 1:
                    plural = ""
                else:
                    plural = "s"
                self._update_status("Started tracing %d function%s. Press Ctrl+C to stop." % (len(self._targets), plural))

                self._resume()

        def on_trace_error(self, error):
            self._print(Fore.RED + Style.BRIGHT + "Error" + Style.RESET_ALL + ": " + error['message'])

        def on_trace_events(self, events):
            no_attributes = Style.RESET_ALL
            for timestamp, thread_id, depth, target_address, message in events:
                if self._output is not None:
                    self._output.append(message + "\n")
                elif self._quiet:
                    self._print(message)
                else:
                    indent = depth * "   | "
                    attributes = self._get_attributes(thread_id)
                    if thread_id != self._last_event_tid:
                        self._print("%s           /* TID 0x%x */%s" % (attributes, thread_id, Style.RESET_ALL))
                        self._last_event_tid = thread_id
                    self._print("%6d ms  %s%s%s%s" % (timestamp, attributes, indent, message, no_attributes))

        def on_trace_handler_create(self, function, handler, source):
            if self._quiet:
                return
            self._print("%s: Auto-generated handler at \"%s\"" % (function, source.replace("\\", "\\\\")))

        def on_trace_handler_load(self, function, handler, source):
            if self._quiet:
                return
            self._print("%s: Loaded handler at \"%s\"" % (function, source.replace("\\", "\\\\")))

        def _get_attributes(self, thread_id):
            attributes = self._attributes_by_thread_id.get(thread_id, None)
            if attributes is None:
                color = self._next_color
                self._next_color += 1
                attributes = self._palette[color % len(self._palette)]
                if (1 + int(color / len(self._palette))) % 2 == 0:
                    attributes += Style.BRIGHT
                self._attributes_by_thread_id[thread_id] = attributes
            return attributes

    app = TracerApplication()
    app.run()


class TracerProfileBuilder(object):
    _RE_REL_ADDRESS = re.compile("(?P<module>[^\s!]+)!(?P<offset>(0x)?[0-9a-fA-F]+)")

    def __init__(self):
        self._spec = []

    def include_modules(self, *module_name_globs):
        for m in module_name_globs:
            self._spec.append(('include', 'module', m))
        return self

    def exclude_modules(self, *module_name_globs):
        for m in module_name_globs:
            self._spec.append(('exclude', 'module', m))
        return self

    def include(self, *function_name_globs):
        for f in function_name_globs:
            self._spec.append(('include', 'function', f))
        return self

    def exclude(self, *function_name_globs):
        for f in function_name_globs:
            self._spec.append(('exclude', 'function', f))
        return self

    def include_relative_address(self, *address_rel_offsets):
        for f in address_rel_offsets:
            m = TracerProfileBuilder._RE_REL_ADDRESS.search(f)
            if m is None:
                continue
            self._spec.append(('include', 'relative_function', {
                'module': m.group('module'),
                'offset': int(m.group('offset'), base=16)
            }))
        return self

    def include_imports(self, *module_name_globs):
        for m in module_name_globs:
            self._spec.append(('include', 'imports', m))
        return self

    def include_objc_method(self, *function_name_globs):
        for f in function_name_globs:
            self._spec.append(('include', 'objc_method', f))
        return self

    def exclude_objc_method(self, *function_name_globs):
        for f in function_name_globs:
            self._spec.append(('exclude', 'objc_method', f))
        return self

    def include_debug_symbol(self, *function_name_globs):
        for f in function_name_globs:
            self._spec.append(('include', 'debug_symbol', f))
        return self

    def build(self):
        return TracerProfile(self._spec)


class TracerProfile(object):
    _BLACKLIST = set([
        "libSystem.B.dylib!dyld_stub_binder"
    ])

    def __init__(self, spec):
        self._spec = spec

    def resolve(self, session, runtime, log_handler=None):
        script = session.create_script(name="profile-resolver",
                                       source=self._create_resolver_script(),
                                       runtime=runtime)
        script.set_log_handler(log_handler)
        def on_message(message, data):
            print(message)
        script.on('message', on_message)
        script.load()
        try:
            data = script.exports.resolve(self._spec)
        finally:
            script.unload()

        modules = {}
        for module_id, m in data['modules'].items():
            module = Module(m['name'], int(m['base'], 16), m['size'], m['path'])
            modules[int(module_id)] = module

        working_set = []
        for target in data['targets']:
            objc = target.get('objc')
            if objc is not None:
                method = objc['method']
                of = ObjCMethod(method['type'], objc['className'], method['name'], int(target['address'], 16))
                working_set.append(of)
            else:
                name = target['name']
                absolute_address = int(target['address'], 16)
                module_id = target.get('module')
                if module_id is not None:
                    module = modules[module_id]
                    relative_address = absolute_address - module.base_address
                    exported = not target.get('private', False)
                    mf = ModuleFunction(module, name, relative_address, exported)
                    if not self._is_blacklisted(mf):
                        working_set.append(mf)
                else:
                    f = Function(name, absolute_address)
                    working_set.append(f)
        return working_set

    def _is_blacklisted(self, module_function):
        key = module_function.module.name + "!" + module_function.name
        return key in TracerProfile._BLACKLIST

    def _create_resolver_script(self):
        return r"""
rpc.exports = {
    resolve: function (spec) {
        var workingSet = spec.reduce(function (workingSet, item) {
            var operation = item[0];
            var scope = item[1];
            var param = item[2];
            switch (scope) {
                case 'module':
                    if (operation === 'include')
                        workingSet = includeModule(param, workingSet);
                    else if (operation === 'exclude')
                        workingSet = excludeModule(param, workingSet);
                    break;
                case 'function':
                    if (operation === 'include')
                        workingSet = includeFunction(param, workingSet);
                    else if (operation === 'exclude')
                        workingSet = excludeFunction(param, workingSet);
                    break;
                case 'relative_function':
                    if (operation === 'include')
                        workingSet = includeRelativeFunction(param, workingSet);
                    break;
                case 'imports':
                    if (operation === 'include')
                        workingSet = includeImports(param, workingSet);
                    break;
                case 'objc_method':
                    if (operation === 'include')
                        workingSet = includeObjCMethod(param, workingSet);
                    else if (operation === 'exclude')
                        workingSet = excludeObjCMethod(param, workingSet);
                    break;
                case 'debug_symbol':
                    if (operation === 'include')
                        workingSet = includeDebugSymbol(param, workingSet);
                    break;
            }
            return workingSet;
        }, {});

        var modules = {};
        var targets = [];
        for (var address in workingSet) {
            if (workingSet.hasOwnProperty(address)) {
                var target = workingSet[address];
                var moduleId = target.module;
                if (moduleId !== undefined && !modules.hasOwnProperty(moduleId)) {
                    var m = allModules()[moduleId];
                    delete m._cachedFunctionExports;
                    modules[moduleId] = m;
                }
                targets.push(target);
            }
        }
        return {
            modules: modules,
            targets: targets
        };
    }
};

function includeModule(pattern, workingSet) {
    moduleResolver().enumerateMatches('exports:' + pattern + '!*').forEach(function (m) {
        workingSet[m.address.toString()] = moduleExportFromMatch(m);
    });
    return workingSet;
}

function excludeModule(pattern, workingSet) {
    moduleResolver().enumerateMatches('exports:' + pattern + '!*').forEach(function (m) {
        delete workingSet[m.address.toString()];
    });
    return workingSet;
}

function includeFunction(pattern, workingSet) {
    var parseResult = parseExportsFunctionPattern(pattern);
    moduleResolver().enumerateMatches('exports:' + parseResult.module + '!' + parseResult.function).forEach(function (m) {
        workingSet[m.address.toString()] = moduleExportFromMatch(m);
    });
    return workingSet;
}

function excludeFunction(pattern, workingSet) {
    var parseResult = parseExportsFunctionPattern(pattern);
    moduleResolver().enumerateMatches('exports:' + parseResult.module + '!' + parseResult.function).forEach(function (m) {
        delete workingSet[m.address.toString()];
    });
    return workingSet;
}

function parseExportsFunctionPattern(pattern) {
    var res = pattern.split('!');

    var m, f;
    if (res.length === 1) {
        m = '*';
        f = res[0];
    } else {
        m = (res[0] === '') ? '*' : res[0];
        f = (res[1] === '') ? '*' : res[1];
    }

    return {
        module: m,
        function: f
    };
}

function includeRelativeFunction(func, workingSet) {
    var relativeToModule = func.module;
    var modules = allModules();
    for (var moduleIndex = 0; moduleIndex !== modules.length; moduleIndex++) {
        var module = modules[moduleIndex];
        if (module.path === relativeToModule || module.name === relativeToModule) {
            var relativeAddress = ptr(func.offset);
            var absoluteAddress = module.base.add(relativeAddress);
            workingSet[absoluteAddress] = {
                name: 'sub_' + relativeAddress.toString(16),
                address: absoluteAddress,
                module: moduleIndex,
                private: true
            };
        }
    }
    return workingSet;
}

function includeImports(pattern, workingSet) {
    var matches;
    if (pattern === null) {
        var mainModule = allModules()[0].path;
        matches = moduleResolver().enumerateMatches('imports:' + mainModule + '!*');
    } else {
        matches = moduleResolver().enumerateMatches('imports:' + pattern + '!*');
    }

    matches.map(moduleExportFromMatch).forEach(function (e) {
        workingSet[e.address.toString()] = e;
    });

    return workingSet;
}

function includeObjCMethod(pattern, workingSet) {
    objcResolver().enumerateMatches(pattern).forEach(function (m) {
        workingSet[m.address.toString()] = objcMethodFromMatch(m);
    });
    return workingSet;
}

function excludeObjCMethod(pattern, workingSet) {
    objcResolver().enumerateMatches(pattern).forEach(function (m) {
        delete workingSet[m.address.toString()];
    });
    return workingSet;
}

function includeDebugSymbol(pattern, workingSet) {
    DebugSymbol.findFunctionsMatching(pattern).forEach(function (addr) {
        workingSet[addr.toString()] = debugSymbolFromAddress(addr);
    });
    return workingSet;
}

var cachedModuleResolver = null;
function moduleResolver() {
    if (cachedModuleResolver === null)
        cachedModuleResolver = new ApiResolver('module');
    return cachedModuleResolver;
}

var cachedObjcResolver = null;
function objcResolver() {
    if (cachedObjcResolver === null) {
        try {
            cachedObjcResolver = new ApiResolver('objc');
        } catch (e) {
            throw new Error('Objective-C runtime is not available');
        }
    }
    return cachedObjcResolver;
}

var cachedModules = null;
function allModules() {
    if (cachedModules === null) {
        cachedModules = Process.enumerateModules();
        cachedModules._idByPath = cachedModules.reduce(function (mappings, module, index) {
            mappings[module.path] = index;
            return mappings;
        }, {});
        cachedModules._idByName = cachedModules.reduce(function (mappings, module, index) {
            mappings[module.name] = index;
            return mappings;
        }, {});
    }
    return cachedModules;
}

function moduleExportFromMatch(m) {
    var encodedName = m.name;
    var delimiterIndex = encodedName.indexOf('!');
    var modulePath = encodedName.substring(0, delimiterIndex);
    var functionName = encodedName.substring(delimiterIndex + 1);
    return {
        name: functionName,
        address: m.address,
        module: allModules()._idByPath[modulePath]
    };
}

function objcMethodFromMatch(m) {
    var encodedName = m.name;
    var methodType = encodedName[0];
    var delimiterIndex = encodedName.indexOf(' ', 3);
    var className = encodedName.substring(2, delimiterIndex);
    var methodName = encodedName.substring(delimiterIndex + 1, encodedName.length - 1);
    return {
        objc: {
            className: className,
            method: {
                type: methodType,
                name: methodName
            }
        },
        address: m.address
    };
}

function debugSymbolFromAddress(address) {
    var symbol = DebugSymbol.fromAddress(address);
    return {
        name: symbol.name,
        address: symbol.address,
        module: allModules()._idByName[symbol.moduleName]
    };
}
"""


class Tracer(object):
    def __init__(self, reactor, repository, profile, init_scripts=[], log_handler=None):
        self._reactor = reactor
        self._repository = repository
        self._profile = profile
        self._script = None
        self._init_scripts = init_scripts
        self._log_handler = log_handler

    def start_trace(self, session, stage, parameters, runtime, ui):
        def on_create(*args):
            ui.on_trace_handler_create(*args)
        self._repository.on_create(on_create)

        def on_load(*args):
            ui.on_trace_handler_load(*args)
        self._repository.on_load(on_load)

        def on_update(function, handler, source):
            self._script.exports.update([{
                'name': function.name,
                'absolute_address': hex(function.absolute_address),
                'handler': handler
            }])
        self._repository.on_update(on_update)

        def on_message(message, data):
            self._reactor.schedule(lambda: self._process_message(message, data, ui))

        ui.on_trace_progress('resolve')
        working_set = self._profile.resolve(session, runtime, log_handler=self._log_handler)
        ui.on_trace_progress('instrument')
        self._script = session.create_script(name="tracer",
                                             source=self._create_trace_script(),
                                             runtime=runtime)
        self._script.set_log_handler(self._log_handler)
        self._script.on('message', on_message)
        self._script.load()

        agent = self._script.exports

        init_scripts = [{ 'filename': script.filename, 'source': script.source } for script in self._init_scripts]
        agent.init(stage, parameters, init_scripts)

        for chunk in [working_set[i:i+1000] for i in range(0, len(working_set), 1000)]:
            targets = [{
                    'name': function.name,
                    'absolute_address': hex(function.absolute_address),
                    'handler': self._repository.ensure_handler(function)
                } for function in chunk]
            agent.add(targets)

        self._repository.commit_handlers()

        self._reactor.schedule(lambda: ui.on_trace_progress('ready'))

        return working_set

    def stop(self):
        if self._script is not None:
            try:
                self._script.unload()
            except:
                pass
            self._script = None

    def _create_trace_script(self):
        return """\
var started = Date.now();
var stage;
var parameters;
var state = {};

(function () {

var handlers = {};
var pending = [];
var timer = null;

rpc.exports = {
    init: function (s, p, initScripts) {
        stage = s;
        parameters = p;

        initScripts.forEach(function (script) {
            try {
                (1, eval)(script.source);
            } catch (e) {
                throw new Error('Unable to load ' + script.filename + ': ' + e.stack);
            }
        });
    },
    dispose: function () {
        flush();
    },
    add: function (targets) {
        targets.forEach(function (target) {
            var h = [parseHandler(target)];
            var name = target.name;
            var targetAddress = target.absolute_address;
            target = null;

            handlers[targetAddress] = h;

            function invokeCallback(callback, context, param) {
                if (callback === undefined)
                    return;

                var timestamp = Date.now() - started;
                var threadId = context.threadId;
                var depth = context.depth;

                function log(message) {
                    emit([timestamp, threadId, depth, targetAddress, message]);
                }

                callback.call(context, log, param, state);
            }

            try {
                Interceptor.attach(ptr(targetAddress), {
                    onEnter: function (args) {
                        invokeCallback(h[0].onEnter, this, args);
                    },
                    onLeave: function (retval) {
                        invokeCallback(h[0].onLeave, this, retval);
                    }
                });
            } catch (e) {
                send({
                    from: '/targets',
                    name: '+error',
                    payload: {
                        message: 'Skipping "' + name + '": ' + e.message
                    }
                });
            }
        });
    },
    update: function (targets) {
        targets.forEach(function (target) {
            handlers[target.absolute_address][0] = parseHandler(target);
        });
    }
};

function emit(event) {
    pending.push(event);

    if (timer === null)
        timer = setTimeout(flush, 50);
}

function flush() {
    if (timer !== null) {
        clearTimeout(timer);
        timer = null;
    }

    if (pending.length === 0)
        return;

    var items = pending;
    pending = [];

    send({
        from: '/events',
        name: '+add',
        payload: {
            items: items
        }
    });
}

function parseHandler(target) {
    try {
        return (1, eval)('(' + target.handler + ')');
    } catch (e) {
        send({
            from: '/targets',
            name: '+error',
            payload: {
                message: 'Invalid handler for "' + target.name + '": ' + e.message
            }
        });
        return {};
    }
}

})();
"""

    def _process_message(self, message, data, ui):
        handled = False
        if message['type'] == 'send':
            stanza = message['payload']
            if isinstance(stanza, dict):
                if stanza['from'] == "/events":
                    if stanza['name'] == '+add':
                        events = [(timestamp, thread_id, depth, int(target_address.rstrip("L"), 16), message) for timestamp, thread_id, depth, target_address, message in stanza['payload']['items']]
                        ui.on_trace_events(events)
                        handled = True
                elif stanza['from'] == "/targets" and stanza['name'] == '+error':
                    ui.on_trace_error(stanza['payload'])
                    handled = True
        if not handled:
            print(message)


class Repository(object):
    def __init__(self):
        self._on_create_callback = None
        self._on_load_callback = None
        self._on_update_callback = None
        self._decorate = False

    def ensure_handler(self, function):
        raise NotImplementedError("not implemented")

    def commit_handlers(self):
        pass

    def on_create(self, callback):
        self._on_create_callback = callback

    def on_load(self, callback):
        self._on_load_callback = callback

    def on_update(self, callback):
        self._on_update_callback = callback

    def _notify_create(self, function, handler, source):
        if self._on_create_callback is not None:
            self._on_create_callback(function, handler, source)

    def _notify_load(self, function, handler, source):
        if self._on_load_callback is not None:
            self._on_load_callback(function, handler, source)

    def _notify_update(self, function, handler, source):
        if self._on_update_callback is not None:
            self._on_update_callback(function, handler, source)

    def _create_stub_handler(self, function, decorate):
        if isinstance(function, ObjCMethod):
            display_name = function.display_name()
            state = {"index": 2}
            def objc_arg(m):
                index = state["index"]
                r = ":' + args[%d] + ' " % index
                state["index"] = index + 1
                return r

            log_str = "'" + re.sub(r':', objc_arg, display_name) + "'"
            if log_str.endswith("' ]'"):
                log_str = log_str[:-3] + "]'"
        else:
            display_name = function.name

            for man_section in (2, 3):
                args = []
                try:
                    with open(os.devnull, 'w') as devnull:
                        man_argv = ["man"]
                        if platform.system() != "Darwin":
                            man_argv.extend(["-E", "UTF-8"])
                        man_argv.extend(["-P", "col -b", str(man_section), function.name])
                        output = subprocess.check_output(man_argv, stderr=devnull)
                    match = re.search(r"^SYNOPSIS(?:.|\n)*?((?:^.+$\n)* {5}\w+[ \*\n]*" + function.name + r"\((?:.+\,\s*?$\n)*?(?:.+\;$\n))(?:.|\n)*^DESCRIPTION", output.decode('UTF-8', errors='replace'), re.MULTILINE)
                    if match:
                        decl = match.group(1)

                        for argm in re.finditer(r"[\(,]\s*(.+?)\s*\b(\w+)(?=[,\)])", decl):
                            typ = argm.group(1)
                            arg = argm.group(2)
                            if arg == "void":
                                continue
                            if arg == "...":
                                args.append("\", ...\" +");
                                continue

                            read_ops = ""
                            annotate_pre = ""
                            annotate_post = ""

                            normalized_type = re.sub(r"\s+", "", typ)
                            if normalized_type.endswith("*restrict"):
                                normalized_type = normalized_type[:-8]
                            if normalized_type in ("char*", "constchar*"):
                                read_ops = ".readUtf8String()"
                                annotate_pre = "\""
                                annotate_post = " + '\"'"

                            arg_index = len(args)

                            args.append("'%(arg_delimiter)s%(arg_name)s=%(annotate_pre)s' + args[%(arg_index)s]%(read_ops)s%(annotate_post)s +" % {
                                "arg_name": arg,
                                "arg_index": arg_index,
                                "arg_delimiter": ", " if arg_index > 0 else "",
                                "read_ops": read_ops,
                                "annotate_pre": annotate_pre,
                                "annotate_post": annotate_post
                            })
                        break
                except Exception as e:
                    pass

            if decorate:
                module_string = " [%s]" %  function.module.name
            else:
                module_string = ""

            if len(args) == 0:
                log_str = "'%(name)s()%(module_string)s'" % { "name": function.name, "module_string" : module_string }
            else:
                indent_outer = "    "
                indent_inner = "      "
                log_str = "'%(name)s(' +\n%(indent_inner)s%(args)s\n%(indent_outer)s')%(module_string)s'" % {
                    "name": function.name,
                    "args": ("\n" + indent_inner).join(args),
                    "indent_outer": indent_outer,
                    "indent_inner": indent_inner,
                    "module_string": module_string
                }

        return """\
/*
 * Auto-generated by Frida. Please modify to match the signature of %(display_name)s.
 * This stub is currently auto-generated from manpages when available.
 *
 * For full API reference, see: http://www.frida.re/docs/javascript-api/
 */

{
  /**
   * Called synchronously when about to call %(display_name)s.
   *
   * @this {object} - Object allowing you to store state for use in onLeave.
   * @param {function} log - Call this function with a string to be presented to the user.
   * @param {array} args - Function arguments represented as an array of NativePointer objects.
   * For example use args[0].readUtf8String() if the first argument is a pointer to a C string encoded as UTF-8.
   * It is also possible to modify arguments by assigning a NativePointer object to an element of this array.
   * @param {object} state - Object allowing you to keep state across function calls.
   * Only one JavaScript function will execute at a time, so do not worry about race-conditions.
   * However, do not use this to store function arguments across onEnter/onLeave, but instead
   * use "this" which is an object for keeping state local to an invocation.
   */
  onEnter: function (log, args, state) {
    log(%(log_str)s);
  },

  /**
   * Called synchronously when about to return from %(display_name)s.
   *
   * See onEnter for details.
   *
   * @this {object} - Object allowing you to access state stored in onEnter.
   * @param {function} log - Call this function with a string to be presented to the user.
   * @param {NativePointer} retval - Return value represented as a NativePointer object.
   * @param {object} state - Object allowing you to keep state across function calls.
   */
  onLeave: function (log, retval, state) {
  }
}
""" % {"display_name": display_name, "log_str": log_str}


class MemoryRepository(Repository):
    def __init__(self):
        super(MemoryRepository, self).__init__()
        self._handlers = {}

    def ensure_handler(self, function):
        handler = self._handlers.get(function)
        if handler is None:
            handler = self._create_stub_handler(function, False)
            self._handlers[function] = handler
            self._notify_create(function, handler, "memory")
        else:
            self._notify_load(function, handler, "memory")
        return handler


class FileRepository(Repository):
    def __init__(self, reactor, decorate):
        super(FileRepository, self).__init__()
        self._reactor = reactor
        self._handler_by_address = {}
        self._handler_by_file = {}
        self._changed_files = set()
        self._last_change_id = 0
        self._repo_dir = os.path.join(os.getcwd(), "__handlers__")
        self._repo_monitors = {}
        self._decorate = decorate

    def ensure_handler(self, function):
        entry = self._handler_by_address.get(function.absolute_address)
        if entry is not None:
            (function, handler, handler_file) = entry
            return handler

        handler = None
        handler_files_to_try = []

        if isinstance(function, ModuleFunction):
            module_dir = os.path.join(self._repo_dir, to_filename(function.module.name))
            module_handler_file = os.path.join(module_dir, to_handler_filename(function.name))
            handler_files_to_try.append(module_handler_file)

        any_module_handler_file = os.path.join(self._repo_dir, to_handler_filename(function.name))
        handler_files_to_try.append(any_module_handler_file)

        for handler_file in handler_files_to_try:
            if os.path.isfile(handler_file):
                with codecs.open(handler_file, 'r', 'utf-8') as f:
                    handler = f.read()
                self._notify_load(function, handler, handler_file)
                break

        if handler is None:
            handler = self._create_stub_handler(function, self._decorate)
            handler_file = handler_files_to_try[0]
            handler_dir = os.path.dirname(handler_file)
            if not os.path.isdir(handler_dir):
                os.makedirs(handler_dir)
            with open(handler_file, 'w') as f:
                f.write(handler)
            self._notify_create(function, handler, handler_file)

        entry = (function, handler, handler_file)
        self._handler_by_address[function.absolute_address] = entry
        self._handler_by_file[handler_file] = entry

        self._ensure_monitor(handler_file)

        return handler

    def _ensure_monitor(self, handler_file):
        handler_dir = os.path.dirname(handler_file)
        monitor = self._repo_monitors.get(handler_dir)
        if monitor is None:
            monitor = FileMonitor(handler_dir)
            monitor.on('change', self._on_change)
            self._repo_monitors[handler_dir] = monitor

    def commit_handlers(self):
        for monitor in self._repo_monitors.values():
            monitor.enable()

    def _on_change(self, changed_file, other_file, event_type):
        if changed_file not in self._handler_by_file or event_type == 'changes-done-hint':
            return
        self._changed_files.add(changed_file)
        self._last_change_id += 1
        change_id = self._last_change_id
        self._reactor.schedule(lambda: self._sync_handlers(change_id), delay=0.05)

    def _sync_handlers(self, change_id):
        if change_id != self._last_change_id:
            return
        changes = self._changed_files.copy()
        self._changed_files.clear()
        for changed_handler_file in changes:
            (function, old_handler, handler_file) = self._handler_by_file[changed_handler_file]
            with codecs.open(handler_file, 'r', 'utf-8') as f:
                new_handler = f.read()
            changed = new_handler != old_handler
            if changed:
                entry = (function, new_handler, handler_file)
                self._handler_by_address[function.absolute_address] = entry
                self._handler_by_file[handler_file] = entry
                self._notify_update(function, new_handler, handler_file)


class InitScript(object):
    def __init__(self, filename, source):
        self.filename = filename
        self.source = source


class OutputFile(object):
    def __init__(self, filename):
        self._fd = codecs.open(filename, 'wb', 'utf-8')

    def close(self):
        self._fd.close()

    def append(self, message):
        self._fd.write(message)
        self._fd.flush()


class UI(object):
    def on_trace_progress(self, operation):
        pass

    def on_trace_error(self, error):
        pass

    def on_trace_events(self, events):
        pass

    def on_trace_handler_create(self, function, handler, source):
        pass

    def on_trace_handler_load(self, function, handler, source):
        pass


def to_filename(name):
    result = ""
    for c in name:
        if c.isalnum() or c == ".":
            result += c
        else:
            result += "_"
    return result


def to_handler_filename(name):
    full_filename = to_filename(name)
    if len(full_filename) <= 41:
        return full_filename + ".js"
    crc = binascii.crc32(full_filename.encode())
    return full_filename[0:32] + "_%08x.js" % crc


if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        pass
