diff --git a/pykd_ext/dllmain.cpp b/pykd_ext/dllmain.cpp index bc42683..6284baf 100644 --- a/pykd_ext/dllmain.cpp +++ b/pykd_ext/dllmain.cpp @@ -9,13 +9,13 @@ BOOL APIENTRY DllMain( HMODULE hModule, LPVOID lpReserved ) { - //if (!pinHandle) - //{ - // GetModuleHandleEx( - // GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS | GET_MODULE_HANDLE_EX_FLAG_PIN, - // (LPCTSTR)hModule, - // &pinHandle); - //} + if (!pinHandle) + { + GetModuleHandleEx( + GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS | GET_MODULE_HANDLE_EX_FLAG_PIN, + (LPCTSTR)hModule, + &pinHandle); + } switch (ul_reason_for_call) { diff --git a/pykd_ext/pyinterpret.cpp b/pykd_ext/pyinterpret.cpp index d987f3d..36735bb 100644 --- a/pykd_ext/pyinterpret.cpp +++ b/pykd_ext/pyinterpret.cpp @@ -52,6 +52,9 @@ public: bool isPy3; + void checkPykd(); + void deactivate(); + PyObject* PyType_Type; PyObject* PyProperty_Type; PyObject* Py_None; @@ -121,8 +124,11 @@ public: HMODULE m_handlePython; PyThreadState* m_globalState; PythonInterpreter* m_globalInterpreter; + bool m_pykdInit; }; + + class PythonInterpreter { public: @@ -227,12 +233,19 @@ public: return m_modules.find(std::make_pair(majorVersion, minorVersion)) != m_modules.end(); } + void checkPykd() + { + m_currentInterpter->m_module->checkPykd(); + } + void stopAllInterpreter() { for (auto m : m_modules) - delete m.second; - - m_modules.clear(); + { + m_currentInterpter = m.second->m_globalInterpreter; + m.second->deactivate(); + } + m_currentInterpter = 0; } private: @@ -404,23 +417,80 @@ PyModule::PyModule(int majorVesion, int minorVersion) *reinterpret_cast(&PyGILState_Release) = GetProcAddress(m_handlePython, "PyGILState_Release"); *reinterpret_cast(&PyDescr_NewMethod) = GetProcAddress(m_handlePython, "PyDescr_NewMethod"); - Py_Initialize(); PyEval_InitThreads(); + m_globalState = PyEval_SaveThread(); } PyModule::~PyModule() { - PyEval_RestoreThread(m_globalState); + assert(0); - Py_Finalize(); + //if (m_globalInterpreter) + //{ + // delete m_globalInterpreter; + // m_globalInterpreter = 0; + //} - FreeLibrary(m_handlePython); + //PyThreadState_Swap(m_globalState); + + //Py_Finalize(); + + //FreeLibrary(m_handlePython); } +void PyModule::deactivate() +{ + if (m_globalInterpreter) + { + delete m_globalInterpreter; + m_globalInterpreter = 0; + } + + PyThreadState_Swap(m_globalState); + + if (m_pykdInit) + { + PyObject* mainName = isPy3 ? PyUnicode_FromString("__main__") : PyString_FromString("__main__"); + PyObject* mainMod = PyImport_Import(mainName); + PyObject* globals = PyObject_GetAttrString(mainMod, "__dict__"); + PyObject* result = PyRun_String("__import__('pykd').deinitialize()\n", Py_file_input, globals, globals); + + if (mainName) Py_DecRef(mainName); + if (mainMod) Py_DecRef(mainMod); + if (globals) Py_DecRef(globals); + if (result) Py_DecRef(result); + } + + m_globalState = PyEval_SaveThread(); +} + +void PyModule::checkPykd() +{ + PyObject* pykdMod = PyImport_ImportModule("pykd"); + + if (!pykdMod) + throw std::exception("Pykd package is not installed.You can install it by command \"!pip install pykd\""); + + PyObject* mainName = isPy3 ? PyUnicode_FromString("__main__") : PyString_FromString("__main__"); + PyObject* mainMod = PyImport_Import(mainName); + PyObject* globals = PyObject_GetAttrString(mainMod, "__dict__"); + PyObject* result = PyRun_String("__import__('pykd').initialize()\n", Py_file_input, globals, globals); + + if (mainName) Py_DecRef(mainName); + if (mainMod) Py_DecRef(mainMod); + if (globals) Py_DecRef(globals); + if (result) Py_DecRef(result); + if (pykdMod) Py_DecRef(pykdMod); + + m_pykdInit = true; +} + + + PythonInterpreter* activateInterpreter(bool global, int majorVersion, int minorVersion) { return PythonSingleton::get()->getInterpreter(majorVersion, minorVersion, global); @@ -441,6 +511,11 @@ void stopAllInterpreter() PythonSingleton::get()->stopAllInterpreter(); } +void checkPykd() +{ + PythonSingleton::get()->checkPykd(); +} + void __stdcall Py_IncRef(PyObject* object) { PythonSingleton::get()->currentInterpreter()->m_module->Py_IncRef(object); diff --git a/pykd_ext/pyinterpret.h b/pykd_ext/pyinterpret.h index 4bff860..5ec828f 100644 --- a/pykd_ext/pyinterpret.h +++ b/pykd_ext/pyinterpret.h @@ -24,6 +24,8 @@ bool isInterpreterLoaded(int majorVersion, int minorVersion); void stopAllInterpreter(); +void checkPykd(); + class AutoInterpreter { diff --git a/pykd_ext/windbgext.cpp b/pykd_ext/windbgext.cpp index f87c9c2..8884c5e 100644 --- a/pykd_ext/windbgext.cpp +++ b/pykd_ext/windbgext.cpp @@ -100,7 +100,7 @@ VOID CALLBACK DebugExtensionUninitialize() { - // stopAllInterpreter(); + stopAllInterpreter(); } ////////////////////////////////////////////////////////////////////////////// @@ -200,6 +200,13 @@ help( ////////////////////////////////////////////////////////////////////////////// +static const char consoleScript[] = + "import pykd\n" + "from pykd import *\n" + "import code\n" + "code.InteractiveConsole(globals()).interact()\n"; + + extern "C" HRESULT CALLBACK @@ -239,11 +246,13 @@ py( PyObjectRef mainMod = PyImport_Import(mainName); PyObjectRef globals = PyObject_GetAttrString(mainMod, "__dict__"); + checkPykd(); + InterruptWatch interruptWatch(client); if (opts.args.empty()) { - PyObjectRef result = PyRun_String("__import__('code').InteractiveConsole(__import__('__main__').__dict__).interact()\n", Py_file_input, globals, globals); + PyObjectRef result = PyRun_String(consoleScript, Py_file_input, globals, globals); } else {