Implementing command completion for IDAPython

In this blog post we are going to illustrate how to use the command line interpreter (CLI) interface from Python and how to write a basic command completion functionality for the Python CLI.

Understanding the CLI structure

IDA Pro SDK allows programmers to implement command line interpreters with the use of the cli_t structure:

struct cli_t
{
  size_t size;                  // Size of this structure
  int32 flags;                  // Feature bits
  const char *sname;            // Short name (displayed on the button)
  const char *lname;            // Long name (displayed in the menu)
  const char *hint;             // Hint for the input line
  // callback: the user pressed Enter
  bool (idaapi *execute_line)(const char *line);
  // callback: the user pressed Tab (optional)
  bool (idaapi *complete_line)(
        qstring *completion,
        const char *prefix,
        int n,
        const char *line,
        int prefix_start);
  // callback: a keyboard key has been pressed (optional)
  bool (idaapi *keydown)(
        qstring *line,
        int *p_x,
        int *p_sellen,
        int *vk_key,
        int shift);
};

For example, the IDAPython plugin defines its CLI like this:

static const cli_t cli_python =
{
    sizeof(cli_t),
    0,
    "Python",
    "Python - IDAPython plugin",
    "Enter any Python expression",
    IDAPython_cli_execute_line,
    NULL, // No completion support
    NULL // No keydown support
};

And registers/unregisters it with:

void install_command_interpreter(const cli_t *cp);
void remove_command_interpreter(const cli_t *cp);

The CLI in Python

The CLI functionality has been added in IDAPython 1.4.1. Let us suppose that we want to reimplement the Python CLI in Python (rather than in C++), we would do it like this:

class pycli_t(idaapi.cli_t):
    flags = 0
    sname = "PyPython"
    lname = "Python - PyCLI"
    hint  = "Enter any Python statement"
    def OnExecuteLine(self, line):
        """
        The user pressed Enter.
        @param line: typed line(s)
        @return Boolean: True-executed line, False-ask for more lines
        """
        try:
            exec(line, globals(), globals())
        except Exception, e:
            print str(e) + "\n" + traceback.format_exc()
        return True

Summary:

1. Subclass idaapi.cli_t
2. Define the CLI short name, long name and hint
3. Declare the OnExecuteLine handler and use Python’s exec() to execute a line

To try this code, simply instantiate a CLI and register it:

pycli = pycli_t()
pycli.register()

And later to unregister the CLI:

pycli.unregister()

Adding command completion

In order to add command completion, we need to implement the OnCompleteLine() callback:

class cli_t:
    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 prefix_start (string)
        @param n: completion number (int)
        @param line: the current line (string)
        @param prefix_start: the index where PREFIX starts in LINE (int)
        @return: None or a String with the completion value
        """
        print "OnCompleteLine: pfx=%s n=%d line=%s pfx_st=%d" % (prefix, n, line, prefix_start)
        return None

This callback is invoked everytime the user presses TAB in the CLI. The callback receives:

• prefix: the prefix at the cursor position
• n: the completion number. If this callback succeeded the first time (when n == 0), then IDA will call this callback again with n=1 (and so on) asking for the next possible completion value
• line: the whole line
• prefix_start: the index of the prefix in the line

For example typing “os.path.ex” and pressing TAB would call:

OnCompleteLine(prefix="ex", n=0, line="print os.path.ex", prefix_start=14)

For demonstration purposes, here is a very simple algorithm to guess what could complete the “os.path.ex” expression:

1. Parse the identifier: Since IDA passes the line, prefix and the prefix start index, we need to get the whole expression including the prefix (we need “os.path.ex”).
   This can be done by scanning backwards from prefix_start until we encounter a non identifier character:

   def parse_identifier(line, prefix, prefix_start):
       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)]

2. Fetch the attributes with dir(): Given the full identifier name (example “os.path.ex”), we split the name by ‘.’ and use a getattr() loop starting from the __main__ module up to the last split value:

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

3. Implementing OnCompleteLine(): We parse the identifier and get a list of possible completion values only if n==0, otherwise we return the next possible value in the list we previously built:

   def OnCompleteLine(self, prefix, n, line, prefix_start):
       # new search?
       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]

With this approach we could complete something like this:

f = open('somefile.txt', 'r')
f.re^TAB => f.read, f.readinto, f.readline or f.readlines

Or:

print idau^TAB.GetRe^TAB => print idautils.GetRegisterList

The sample script can be downloaded from here and a win32 build of the latest IDAPython plugin (with completion integrated in the plugin) can be downloaded from here.