//------------------------------------------------------------
// IDAPython - Python plugin for Interactive Disassembler Pro
//
// Copyright (c) 2004-2009 Gergely Erdelyi <dyce@d-dome.net>
//
// All rights reserved.
//
// For detailed copyright information see the file COPYING in
// the root of the distribution archive.
//------------------------------------------------------------
// python.cpp - Main plugin code
//------------------------------------------------------------
#include <Python.h>

/* This define fixes the redefinition of ssize_t */
#ifdef HAVE_SSIZE_T
#define _SSIZE_T_DEFINED 1
#endif

#include <stdio.h>
#include <string.h>
#ifdef __LINUX__
#include <dlfcn.h>
#endif
#include <ida.hpp>
#include <idp.hpp>
#include <ieee.h>
#include <bytes.hpp>
#include <diskio.hpp>
#include <expr.hpp>
#include <loader.hpp>
#include <kernwin.hpp>
#include <netnode.hpp>

#ifdef __cplusplus
extern "C"
#endif

/* Python-style version tuple comes from the makefile */
/* Only the serial and status is set here */
#define VER_SERIAL 0
#define VER_STATUS "final"

#define IDAPYTHON_RUNFILE      0
#define IDAPYTHON_RUNSTATEMENT 1
#define IDAPYTHON_SCRIPTBOX    2


#define IDAPYTHON_DATA_STATEMENT 0

#ifdef __EA64__
#define PYTHON_DIR_NAME "python64"
#else
#define PYTHON_DIR_NAME "python"
#endif

void init_idaapi(void);
void idaapi run(int arg);
static int initialized = 0;

/* This is a simple tracing code for debugging purposes. */
/* It might evolve into a tracing facility for user scripts. */
/* #define ENABLE_PYTHON_PROFILING */

#ifdef ENABLE_PYTHON_PROFILING
#include "compile.h"
#include "frameobject.h"

int tracefunc(PyObject *obj, _frame *frame, int what, PyObject *arg)
{
    PyObject *str;

    /* Catch line change events. */
    /* Print the filename and line number */
    if (what == PyTrace_LINE)
    {
        str = PyObject_Str(frame->f_code->co_filename);
        if (str)
        {
            msg("PROFILING: %s:%d\n", PyString_AsString(str), frame->f_lineno);
            Py_DECREF(str);
        }
    }
    return 0;
}
#endif

/* Helper routines to make Python script execution breakable from IDA */
static int ninsns = 0;      // number of times trace function was called
static bool box_displayed;  // has the wait box been displayed?
static time_t start_time;   // the start time of the execution
static int script_timeout = 2;

/* Exported to the Python environment */
void set_script_timeout(int timeout)
{
    script_timeout = timeout;
}

/* This callback is called on various interpreter events */
int break_check(PyObject *obj, _frame *frame, int what, PyObject *arg)
{
    if (wasBreak())
        /* User pressed Cancel in the waitbox; send KeyboardInterrupt exception */
        PyErr_SetInterrupt();
    else if (!box_displayed && ++ninsns > 10)
    {
        /* We check the timer once every 10 calls */
        ninsns = 0;
        if (time(NULL) - start_time > script_timeout) /* Timeout elapsed? */
        {
            show_wait_box("Running Python script");
            box_displayed = true;
        }
    }
#ifdef ENABLE_PYTHON_PROFILING
    return tracefunc(obj, frame, what, arg);
#else
    return 0;
#endif
}

/* Prepare for Python execution */
void begin_execution()
{
    ninsns = 0;
    box_displayed = false;
    start_time = time(NULL);
    PyEval_SetTrace(break_check, NULL);
}

/* Called after Python execution finishes */
void end_execution()
{
    if (box_displayed)
        hide_wait_box();
#ifdef ENABLE_PYTHON_PROFILING
    PyEval_SetTrace(tracefunc, NULL);
#else
    PyEval_SetTrace(NULL, NULL);
#endif
}

/* Simple Python statement runner function for IDC */
static const char idc_runpythonstatement_args[] = { VT_STR, 0 };
static error_t idaapi idc_runpythonstatement(idc_value_t *argv, idc_value_t *res)
{
    begin_execution();
    res->num = PyRun_SimpleString(argv[0].str);
    end_execution();
    return eOk;
}

/* QuickFix for the FILE* incompatibility problem */
int ExecFile(const char *FileName)
{
    PyObject* PyFileObject = PyFile_FromString((char*)FileName, "r");

    if (!PyFileObject)
        return 0;

    if (PyRun_SimpleFile(PyFile_AsFile(PyFileObject), FileName) == 0)
    {
        Py_DECREF(PyFileObject);
        return 1;
    }
    else
    {
        Py_DECREF(PyFileObject);
        return 0;
    }
}

/* Check for the presence of a file in IDADIR/python */
bool CheckFile(char *filename)
{
    char filepath[MAXSTR+1];

    qmakepath(filepath, MAXSTR, idadir(PYTHON_DIR_NAME), filename, NULL);
    if (!qfileexist(filepath))
    {
        warning("IDAPython: Missing required file %s", filename);
        return false;
    }

    return true;
}

/* Execute the Python script from the plugin */
/* Default hotkey: Alt-9 */
void IDAPython_RunScript(char *script)
{
    char statement[MAXSTR+32];
    char slashpath[MAXSTR+1];
    char *scriptpath;

    int i;

    if (script)
        scriptpath = script;
    else
    {
        scriptpath = askfile_c(0, "*.py", "Python file to run");
        if (!scriptpath)
            return;
    }

    /* Make a copy of the path with '\\' => '/' */
    for (i=0; scriptpath[i]; i++)
    {
        if (scriptpath[i] == '\\')
            slashpath[i] = '/';
        else
            slashpath[i] = scriptpath[i];
    }
    slashpath[i] = '\0';

    /* Add the script't path to sys.path */
    qsnprintf(statement, sizeof(statement), "runscript(\"%s\")", slashpath);
    begin_execution();
    PyRun_SimpleString(statement);
    end_execution();

    /* Error handling */
    if (PyErr_Occurred())
        PyErr_Print();

}

/* Execute Python statement(s) from an editor window */
/* Default hotkey: Alt-8 */
void IDAPython_RunStatement(void)
{
    char statement[4096];
    netnode history;

    /* Get the existing or create a new netnode in the database */
    history.create("IDAPython_Data");

    /* Fetch the previous statement */
    if (history.supval(IDAPYTHON_DATA_STATEMENT, statement, sizeof(statement)) == -1)
        statement[0] = '\0';

    if (asktext(sizeof(statement), statement, statement, "Enter Python expressions"))
    {
        begin_execution();
        PyRun_SimpleString(statement);
        end_execution();
        /* Store the statement to the database */
        history.supset(IDAPYTHON_DATA_STATEMENT, statement);
    }
}

/* History of previously executed scripts */
/* Default hotkey: Alt-7 */
void IDAPython_ScriptBox(void)
{
    PyObject *module;
    PyObject *dict;
    PyObject *scriptbox;
    PyObject *pystr;

    /* Get globals() */
    /* These two should never fail */
    module = PyImport_AddModule("__main__");
    dict = PyModule_GetDict(module);

    scriptbox = PyDict_GetItemString(dict, "scriptbox");

    if (!scriptbox)
    {
        warning("INTERNAL ERROR: ScriptBox_instance missing! Broken init.py?");
        return;
    }

    pystr = PyObject_CallMethod(scriptbox, "run", "");

    if (pystr)
    {
        /* If the return value is string use it as path */
        if (PyObject_TypeCheck(pystr, &PyString_Type))
        {
            begin_execution();
            ExecFile(PyString_AsString(pystr));
            end_execution();
        }
        Py_DECREF(pystr);
    }
    else
    {
        /* Print the exception info */
        if (PyErr_Occurred())
            PyErr_Print();
    }
}

bool idaapi IDAPython_Menu_Callback(void *ud)
{
    run((int)ud);
    return true;
}

/* Return a formatted error or just print it to the console */
static void handle_python_error(char *errbuf, size_t errbufsize)
{
    PyObject *result;
    PyObject *ptype, *pvalue, *ptraceback;

    if ( errbufsize > 0 )
        errbuf[0] = '\0';

    if (PyErr_Occurred())
    {
        PyErr_Fetch(&ptype, &pvalue, &ptraceback);
        result = PyObject_Repr(pvalue);
        if (result)
        {
            qsnprintf(errbuf, errbufsize, "ERROR: %s", PyString_AsString(result));
            PyErr_Clear();
            Py_XDECREF(ptype);
            Py_XDECREF(pvalue);
            Py_XDECREF(ptraceback);
        }
        else
            PyErr_Print();
    }
}

/* Convert return value from Python to IDC or report about an error */
static bool return_python_result(idc_value_t *rv,
                                 PyObject *result,
                                 char *errbuf,
                                 size_t errbufsize)
{
    if (errbufsize > 0)
        errbuf[0] = '\0';

    if (result == NULL)
    {
        handle_python_error(errbuf, errbufsize);
        return false;
    }

    if (PyInt_Check(result))
    {
        rv->num = PyInt_AsLong(result);
        rv->vtype = VT_LONG;
        Py_XDECREF(result);
        return true;
    }

    if (PyString_Check(result))
    {
        rv->str = (char *)qalloc(PyString_Size(result)+1);
        if (!rv->str)
            return false;
        qstrncpy(rv->str, PyString_AsString(result), MAXSTR);
        rv->vtype = VT_STR;
        Py_XDECREF(result);
        return true;
    }

    if (PyFloat_Check(result))
    {
        double dresult = PyFloat_AsDouble(result);
        ieee_realcvt((void *)&dresult, rv->e, 3);
        rv->vtype = VT_FLOAT;
        Py_XDECREF(result);
        return true;
    }

    qsnprintf(errbuf, errbufsize, "ERROR: bad return value");
    return false;
}

/* Compile callback for Python external language evaluator */
bool idaapi IDAPython_extlang_compile(const char *name,
				      ea_t /*current_ea*/,
				      const char *expr,
				      char *errbuf,
				      size_t errbufsize)
{
    PyObject *main = PyImport_AddModule("__main__");  QASSERT(main != NULL);
    PyObject *globals = PyModule_GetDict(main);       QASSERT(globals != NULL);

    PyCodeObject *code = (PyCodeObject *)Py_CompileString(expr, "<string>", Py_eval_input);
    if (code == NULL)
    {
        handle_python_error(errbuf, errbufsize);
        return false;
    }

    // set the desired function name
    Py_XDECREF(code->co_name);
    code->co_name = PyString_FromString(name);

    // create a function out of code
    PyObject *func = PyFunction_New((PyObject *)code, globals);
    if (func == NULL)
    {
    ERR:
        handle_python_error(errbuf, errbufsize);
        Py_XDECREF(code);
        return false;
    }

    int err = PyDict_SetItemString(globals, name, func);
    if (err)
        goto ERR;

    return true;
}

/* Run callback for Python external language evaluator */
bool idaapi IDAPython_extlang_run(const char *name,
				  int nargs,
				  const idc_value_t args[],
				  idc_value_t *result,
				  char *errbuf,
				  size_t errbufsize)
{
    // convert arguments to python
    qvector<PyObject *> pargs;

    for (int i=0; i<nargs; i++)
    {
        double dresult;
        PyObject *pa;
        switch (args[i].vtype)
        {
        case VT_LONG:
            pa = PyInt_FromLong(args[i].num);
            break;
        case VT_STR:
            pa = PyString_FromString(args[i].str);
            break;
        case VT_FLOAT:
            ieee_realcvt(&dresult, (ushort *)args[i].e, 013);
            pa = PyFloat_FromDouble(dresult);
            break;
        default:
            qsnprintf(errbuf, errbufsize, "arg#%d has wrong type %d", i, args[i].vtype);
            return false;
        }
        pargs.push_back(pa);
    }

    PyObject *main = PyImport_AddModule("__main__");  QASSERT(main != NULL);
    PyObject *globals = PyModule_GetDict(main);       QASSERT(globals != NULL);
    PyObject *func = PyDict_GetItemString(globals, name);

    if (func == NULL)
    {
        qsnprintf(errbuf, errbufsize, "undefined function %s", name);
        return false;
    }

    PyCodeObject *code = (PyCodeObject *)PyFunction_GetCode(func);
    PyObject *pres = PyEval_EvalCodeEx(code, globals, NULL, &pargs[0], nargs,
                                       NULL, 0, NULL, 0, NULL);

    // free argument objects
    for  (int i=0; i<nargs; i++)
        Py_XDECREF(pargs[i]);

    return return_python_result(result, pres, errbuf, errbufsize);
}


/* Compile callback for Python external language evaluator */
bool idaapi IDAPython_extlang_compile_file(const char *name,
				           char *errbuf,
				           size_t errbufsize)
{
    PyObject *main = PyImport_AddModule("__main__");  QASSERT(main != NULL);
    PyObject *globals = PyModule_GetDict(main);       QASSERT(globals != NULL);

    if (!ExecFile(name))
    {
        handle_python_error(errbuf, errbufsize);
        return false;
    }

    return true;
}

/* Calculator callback for Python external language evaluator */
bool idaapi IDAPython_extlang_calcexpr(ea_t /*current_ea*/,
                                       const char *expr,
                                       idc_value_t *rv,
                                       char *errbuf,
                                       size_t errbufsize)
{
    PyObject *result;
    PyObject *module = PyImport_AddModule("__main__");

    if (module == NULL)
        return false;

    PyObject *globals = PyModule_GetDict(module);

    begin_execution();
    result = PyRun_String(expr, Py_eval_input, globals, globals);
    end_execution();

    VarFree(rv);

    return return_python_result(rv, result, errbuf, errbufsize);
}

extlang_t extlang_python =
{
    sizeof(extlang_t),
    0,
    "Python",
    IDAPython_extlang_compile,
    IDAPython_extlang_run,
    IDAPython_extlang_calcexpr,
    IDAPython_extlang_compile_file,
    "py"
};

void enable_extlang_python(bool enable)
{
    if (enable)
        register_extlang(&extlang_python);
    else
        register_extlang(NULL);
}

#if IDA_SDK_VERSION >= 540
/* Execute a line in the Python CLI */
bool idaapi IDAPython_cli_execute_line(const char *line)
{
    begin_execution();
    PyRun_SimpleString(line);
    end_execution();
    return true;
}

cli_t cli_python =
{
    sizeof(cli_t),
    0,
    "Python",
    "Python - IDAPython plugin",
    "Enter any Python expression",
    IDAPython_cli_execute_line,
    NULL,
    NULL
};

/* Control the Python CLI status */
void enable_python_cli(bool enable)
{
    if (enable)
        install_command_interpreter(&cli_python);
    else
        remove_command_interpreter(&cli_python);
}
#endif

/* Install python menu items */
static void install_python_menus()
{
  /* Add menu items for all the functions */
  /* Different paths are used for the GUI version */
  add_menu_item("File/IDC command...", "P~y~thon command...",
    "Alt-8", SETMENU_APP,
    (menu_item_callback_t *)IDAPython_Menu_Callback,
    (void *)IDAPYTHON_RUNSTATEMENT);

  /* Add Load Python file menu item*/
  bool result = add_menu_item("File/Load file/IDC file...", "P~y~thon file...",
    "Alt-9", SETMENU_APP,
    (menu_item_callback_t *)IDAPython_Menu_Callback,
    (void *)IDAPYTHON_RUNFILE);
  if (!result)
    add_menu_item("File/IDC command...", "P~y~thon file...",
    "Alt-9", SETMENU_APP,
    (menu_item_callback_t *)IDAPython_Menu_Callback,
    (void *)IDAPYTHON_RUNFILE);

  /* Add View Python Scripts menu item*/
  result = add_menu_item("View/Open subviews/Show strings", "Python S~c~ripts",
    "Alt-7", SETMENU_APP,
    (menu_item_callback_t *)IDAPython_Menu_Callback,
    (void *)IDAPYTHON_SCRIPTBOX);
  if (!result)
    add_menu_item("View/Open subviews/Problems", "Python S~c~ripts",
    "Alt-7", SETMENU_APP,
    (menu_item_callback_t *)IDAPython_Menu_Callback,
    (void *)IDAPYTHON_SCRIPTBOX);
}

/* we install the menu later because the text version crashes if
add_menu_item is called too early */
static int idaapi menu_installer_cb(void *, int code, va_list)
{
  if ( code != ui_ready_to_run )
    return 0;

  install_python_menus();
  unhook_from_notification_point(HT_UI, menu_installer_cb, NULL);
  return 0;
}

/* Initialize the Python environment */
bool IDAPython_Init(void)
{
    char *options;
    char tmp[MAXSTR+64];
    bool result = true;

    /* Already initialized? */
    if (initialized == 1)
        return true;

    /* Check for the presence of essential files */
    initialized = 0;

    result &= CheckFile("idc.py");
    result &= CheckFile("init.py");
    result &= CheckFile("idaapi.py");
    result &= CheckFile("idautils.py");
    if (!result)
        return false;

#ifdef __LINUX__
    /* Export symbols from libpython to resolve imported module deps */
    qsnprintf(tmp, sizeof(tmp), "libpython%d.%d.so",
              PY_MAJOR_VERSION,
              PY_MINOR_VERSION);
    if (!dlopen(tmp, RTLD_NOLOAD | RTLD_GLOBAL))
    {
        warning("IDAPython: dlopen(%s) failed", tmp);
        return false;
    }
#endif

    /* Start the interpreter */
    Py_Initialize();
    if (!Py_IsInitialized())
    {
        warning("IDAPython: Py_Initialize() failed");
        return false;
    }

    /* Init the SWIG wrapper */
    init_idaapi();

    /* Set IDAPYTHON_VERSION in Python */
    qsnprintf(tmp, sizeof(tmp), "IDAPYTHON_VERSION=(%d, %d, %d, '%s', %d)", \
              VER_MAJOR,
              VER_MINOR,
              VER_PATCH,
              VER_STATUS,
              VER_SERIAL);
    begin_execution();
    PyRun_SimpleString(tmp);
    end_execution();

    /* Pull in the Python side of init */
    qmakepath(tmp, MAXSTR, idadir(PYTHON_DIR_NAME), "init.py", NULL);
    if (!ExecFile(tmp))
    {
        warning("IDAPython: error executing init.py");
        return false;
    }

#ifdef ENABLE_PYTHON_PROFILING
    PyEval_SetTrace(tracefunc, NULL);
#endif

    /* Batch-mode operation: */
    /* A script specified on the command line is run */
    options = (char *)get_plugin_options("IDAPython");
    if (options)
        IDAPython_RunScript(options);

#ifdef PLUGINFIX
    hook_to_notification_point(HT_UI, menu_installer_cb, NULL);
#else
    install_python_menus();
#endif
    /* Register a RunPythonStatement() function for IDC */
    set_idc_func("RunPythonStatement", idc_runpythonstatement, idc_runpythonstatement_args);

#if IDA_SDK_VERSION >= 540
    /* Enable the CLI by default */
    enable_python_cli(true);
#endif

    initialized = 1;

    return true;
}

/* Cleaning up Python */
void IDAPython_Term(void)
{
#ifdef PLUGINFIX
    unhook_from_notification_point(HT_UI, menu_installer_cb, NULL);
#endif
    /* Remove the menu items before termination */
    del_menu_item("File/Load file/Python file...");
    del_menu_item("File/Python file...");
    del_menu_item("File/Python command...");
    del_menu_item("View/Open subviews/Python Scripts");

#if IDA_SDK_VERSION >= 540
    /* Remove the CLI */
    enable_python_cli(false);
#endif

    /* Remove the extlang */
    register_extlang(NULL);

    /* Shut the interpreter down */
    Py_Finalize();

    initialized = 0;
}

/* Plugin init routine */
int idaapi init(void)
{
    if (IDAPython_Init())
        return PLUGIN_KEEP;
    else
        return PLUGIN_SKIP;
}

/* Plugin term routine */
void idaapi term(void)
{
    IDAPython_Term();
}

/* Plugin hotkey entry point */
void idaapi run(int arg)
{
    try
    {
        switch (arg)
        {
        case 0:
            IDAPython_RunScript(NULL);
            break;
            ;;
        case 1:
            IDAPython_RunStatement();
            break;
            ;;
        case 2:
            IDAPython_ScriptBox();
            break;
            ;;
        case 3:
            enable_extlang_python(true);
            break;
            ;;
        case 4:
            enable_extlang_python(false);
            break;
            ;;
        default:
            warning("IDAPython: unknown plugin argument %d", arg);
            break;
            ;;
        }
    }
    catch(...)
    {
        warning("Exception in Python interpreter. Reloading...");
        IDAPython_Term();
        IDAPython_Init();
    }
}

//--------------------------------------------------------------------------
// PLUGIN DESCRIPTION BLOCK
//--------------------------------------------------------------------------
char comment[]       = "IDAPython";
char help[]          = "IDA Python Plugin\n";
char wanted_name[]   = "IDAPython";
char wanted_hotkey[] = "Alt-9";

extern "C"
{
    plugin_t PLUGIN = {
        IDP_INTERFACE_VERSION,
#ifdef PLUGINFIX
        PLUGIN_FIX,    // plugin flags
#else
        0,             // plugin flags
#endif
        init,          // initialize
        term,          // terminate. this pointer may be NULL.
        run,           // invoke plugin
        comment,       // long comment about the plugin
                       // it could appear in the status line
                       // or as a hint
        help,          // multiline help about the plugin
        wanted_name,   // the preferred short name of the plugin
        wanted_hotkey  // the preferred hotkey to run the plugin
    };
}