# -----------------------------------------------------------------------
# This is an example illustrating how to implement a CLI
# It features a very basic command completion capability
#
# (c) Hex-Rays
#
from idaapi import data_type_t, data_format_t, NW_OPENIDB, NW_CLOSEIDB, NW_TERMIDA, NW_REMOVE, COLSTR, cli_t
import sys
import traceback
import __builtin__

# -----------------------------------------------------------------------
class pycli_t(cli_t):
    flags = 0
    sname = "pycli"
    lname = "Python CLI"
    hint  = "pycli hint"

    def __init__(self):
        cli_t.__init__(self)
        self.n = 0
        self.completion = None


    @staticmethod
    def parse_identifier(line, prefix, prefix_start):
        """Parse a line and extracts"""
        id_start = prefix_start
        while id_start > 0:
            ch = line[id_start]
            if not ch.isalpha() and ch != '.' and ch != '_':
                id_start += 1
                break
            id_start -= 1

        return line[id_start:prefix_start + len(prefix)]


    @staticmethod
    def get_completion(id, prefix):
        try:
            parts = id.split('.')
            m = sys.modules['__main__']
            c = len(parts)
            for i in xrange(0, c-1):
                m = getattr(m, parts[i])
        except Exception, e:
            return None
        else:
            # search in the module
            completion = [x for x in dir(m) if x.startswith(prefix)]

            # no completion found? looking from the global scope? then try the builtins
            if not completion and c == 1:
                completion = [x for x in dir(__builtin__) if x.startswith(prefix)]

            return completion if len(completion) else None


    def OnExecuteLine(self, line):
        try:
            exec(line, globals(), globals())
        except Exception, e:
            print str(e) + "\n" + traceback.format_exc()
        return True


    def OnCompleteLine(self, prefix, n, line, prefix_start):
        """
        The user pressed Tab. Find a completion number N for prefix PREFIX

        This callback is optional.

        @param prefix: Line prefix at x (string)
        @param n: completion number (int)
        @param line: the current line (string)
        @param x: the index where PREFIX starts in LINE (int)

        @return: None if no completion could be generated otherwise a String with the completion suggestion
        """
        if n == 0:
            self.n = n
            id = self.parse_identifier(line, prefix, prefix_start)
            self.completion = self.get_completion(id, prefix)
        return None if (self.completion is None) or (n >= len(self.completion)) else self.completion[n]


# -----------------------------------------------------------------------
# notifywhen handler. we use it to properly register/unregister the cli upon
# database open/close
def nw_handler(code, old=0):
    if code == NW_OPENIDB:
        pycli.register()
    elif code == NW_CLOSEIDB:
        pycli.unregister()
    elif code == NW_TERMIDA:
        idaapi.notify_when(NW_TERMIDA | NW_OPENIDB | NW_CLOSEIDB | NW_REMOVE, nw_handler)

# -----------------------------------------------------------------------
# Already installed?
try:
    pycli
    # remove previous CLI
    pycli.unregister()
    # remove previous handler
    nw_handler(NW_TERMIDA)
except:
    pass
finally:
    pycli = pycli_t()

# register CLI
if pycli.register():
    print "CLI installed"
    # install the nwhandler
    idaapi.notify_when(NW_TERMIDA | NW_OPENIDB | NW_CLOSEIDB, nw_handler)
else:
    del pycli
    print "Failed to install CLI"