#include "stdafx.h"

#include <wdbgexts.h>

#include <vector>
#include <string>

#include <boost/python/module.hpp>
#include <boost/python/def.hpp>
#include <boost/tokenizer.hpp>
#include <boost/python/overloads.hpp>

#include "dbgext.h"
#include "dbgprint.h"
#include "dbgreg.h"
#include "dbgtype.h"
#include "dbgmodule.h"
#include "dbgsym.h"
#include "dbgmem.h"
#include "dbgsystem.h"
#include "dbgcmd.h"
#include "dbgdump.h"
#include "dbgexcept.h"
#include "dbgeventcb.h"
#include "dbgsession.h"
#include "dbgcallback.h"
#include "dbgpath.h"
#include "dbginput.h"
#include "dbgprocess.h"
#include "dbgsynsym.h"

/////////////////////////////////////////////////////////////////////////////////

// óêàçàòåëü íà òåêóùéè èíòåðôåéñ
DbgExt    *dbgExt = NULL;


/////////////////////////////////////////////////////////////////////////////////

class WindbgGlobalSession 
{
public:

    WindbgGlobalSession() {
    
        boost::python::import( "pykd" );
        
        main = boost::python::import("__main__");
        
        // ïåðåíàïðàâëåíèå ñòàíäàðòíûõ ïîòîêîâ ÂÂ
        boost::python::object       sys = boost::python::import( "sys");
        
        dbgOut                      dout;
        sys.attr("stdout") = boost::python::object( dout );

        dbgIn                       din;
        sys.attr("stdin") = boost::python::object( din );
    }
    
    boost::python::object
    global() {
        return main.attr("__dict__");
    }

private:
   
    boost::python::object       main;
   
};   

WindbgGlobalSession     *windbgGlobalSession = NULL; 

/////////////////////////////////////////////////////////////////////////////////
BOOST_PYTHON_FUNCTION_OVERLOADS( dprint, DbgPrint::dprint, 1, 2 )
BOOST_PYTHON_FUNCTION_OVERLOADS( dprintln, DbgPrint::dprintln, 1, 2 )

BOOST_PYTHON_FUNCTION_OVERLOADS( loadBytes, loadArray<unsigned char>, 2, 3 )
BOOST_PYTHON_FUNCTION_OVERLOADS( loadWords, loadArray<unsigned short>, 2, 3 )
BOOST_PYTHON_FUNCTION_OVERLOADS( loadDWords, loadArray<unsigned long>, 2, 3 )
BOOST_PYTHON_FUNCTION_OVERLOADS( loadQWords, loadArray<unsigned __int64> , 2, 3 )
BOOST_PYTHON_FUNCTION_OVERLOADS( loadSignBytes, loadArray<char> , 2, 3 )
BOOST_PYTHON_FUNCTION_OVERLOADS( loadSignWords, loadArray<short> , 2, 3 )
BOOST_PYTHON_FUNCTION_OVERLOADS( loadSignDWords, loadArray<long> , 2, 3 )
BOOST_PYTHON_FUNCTION_OVERLOADS( loadSignQWords, loadArray<__int64>, 2, 3 )

BOOST_PYTHON_FUNCTION_OVERLOADS( compareMemoryOver, compareMemory, 3, 4 )

BOOST_PYTHON_MODULE( pykd )
{
    boost::python::def( "go", &setExecutionStatus<DEBUG_STATUS_GO> );
    boost::python::def( "trace", &setExecutionStatus<DEBUG_STATUS_STEP_INTO> );
    boost::python::def( "step", &setExecutionStatus<DEBUG_STATUS_STEP_OVER> );   
    boost::python::def( "expr", &evaluate ); 
    boost::python::def( "createSession", &dbgCreateSession );   // deprecated
    boost::python::def( "isSessionStart", &dbgIsSessionStart );
    boost::python::def( "symbolsPath", &dbgSymPath );
    boost::python::def( "dprint", &DbgPrint::dprint, dprint( boost::python::args( "str", "dml" ), ""  ) );
    boost::python::def( "dprintln", &DbgPrint::dprintln, dprintln( boost::python::args( "str", "dml" ), ""  ) );
    boost::python::def( "loadDump", &dbgLoadDump );
    boost::python::def( "startProcess", &startProcess );
    boost::python::def( "dbgCommand", &dbgCommand );
    boost::python::def( "isValid", &isOffsetValid );
    boost::python::def( "is64bitSystem", &is64bitSystem );
    boost::python::def( "isKernelDebugging", &isKernelDebugging );
    boost::python::def( "ptrSize", ptrSize );
    boost::python::def( "reg", &loadRegister );
    boost::python::def( "typedVar", &loadTypedVar );
    boost::python::def( "typedVarList", &loadTypedVarList );
    boost::python::def( "typedVarArray", &loadTypedVarArray );
    boost::python::def( "containingRecord", &containingRecord );
    boost::python::def( "getTypeClass", &getTypeClass );
    boost::python::def( "sizeof", &sizeofType );
    boost::python::def( "loadModule", &loadModule );
    boost::python::def( "findSymbol", &findSymbolForAddress );
    boost::python::def( "getOffset", &findAddressForSymbol );
    boost::python::def( "findModule", &findModule );
    boost::python::def( "addr64", &addr64 );
    boost::python::def( "loadBytes", &loadArray<unsigned char>, loadBytes( boost::python::args( "address", "number",  "phyAddr"  ), "" ) );
    boost::python::def( "loadWords", &loadArray<unsigned short>, loadWords( boost::python::args( "address", "number",  "phyAddr"  ), "" ) );
    boost::python::def( "loadDWords", &loadArray<unsigned long>, loadDWords( boost::python::args( "address", "number",  "phyAddr"  ), "" ) );
    boost::python::def( "loadQWords", &loadArray<unsigned __int64>, loadQWords( boost::python::args( "address", "number",  "phyAddr"  ), "" ) );
    boost::python::def( "loadSignBytes", &loadArray<char>, loadSignBytes( boost::python::args( "address", "number",  "phyAddr"  ), "" ) );
    boost::python::def( "loadSignWords", &loadArray<short>, loadSignWords( boost::python::args( "address", "number",  "phyAddr"  ), "" ) );
    boost::python::def( "loadSignDWords", &loadArray<long>, loadSignDWords( boost::python::args( "address", "number",  "phyAddr"  ), "" ) );
    boost::python::def( "loadSignQWords", &loadArray<__int64>, loadSignQWords( boost::python::args( "address", "number",  "phyAddr"  ), "" ) );
    boost::python::def( "loadPtrs", &loadPtrArray );
    boost::python::def( "loadUnicodeString", &loadUnicodeStr );
    boost::python::def( "loadAnsiString", &loadAnsiStr );   
    boost::python::def( "loadCStr", &loadCStr );
    boost::python::def( "loadWStr", &loadWStr );
    boost::python::def( "loadLinkedList", &loadLinkedList ); 
    boost::python::def( "ptrByte", &loadByPtr<unsigned char> );
    boost::python::def( "ptrSignByte", &loadByPtr<char> );
    boost::python::def( "ptrWord", &loadByPtr<unsigned short> );
    boost::python::def( "ptrSignWord", &loadByPtr<short> );
    boost::python::def( "ptrDWord", &loadByPtr<unsigned long> );
    boost::python::def( "ptrSignDWord", &loadByPtr<long> );
    boost::python::def( "ptrQWord", &loadByPtr<unsigned __int64> );
    boost::python::def( "ptrSignQWord", &loadByPtr<__int64> );
    boost::python::def( "ptrPtr", &loadPtrByPtr );   
    boost::python::def( "ptrMWord", &loadMWord );
    boost::python::def( "ptrSignMWord", &loadSignMWord );
    boost::python::def( "compareMemory", &compareMemory, compareMemoryOver( boost::python::args( "addr1", "addr2", "length", "phyAddr" ), "" ) );
    boost::python::def( "getCurrentStack", &getCurrentStack );
    boost::python::def( "locals", &getLocals );
    boost::python::def( "reloadModule", &reloadModule );
    boost::python::def( "getPdbFile", &getPdbFile );
    boost::python::def( "getImplicitThread", &getImplicitThread );
    boost::python::def( "setImplicitThread", &setImplicitThread );
    boost::python::def( "getThreadList", &getThreadList );
    boost::python::def( "getCurrentProcess", &getCurrentProcess );
    boost::python::def( "setCurrentProcess", &setCurrentProcess );
    boost::python::def( "getProcessorMode", &getProcessorMode );
    boost::python::def( "setProcessorMode", &setProcessorMode );
    boost::python::def( "addSynSymbol", &addSyntheticSymbol );
    boost::python::def( "delAllSynSymbols", &delAllSyntheticSymbols);
    boost::python::def( "delSynSymbol", &delSyntheticSymbol );
    boost::python::def( "delSynSymbolsMask", &delSyntheticSymbolsMask);
    boost::python::class_<typeClass, boost::shared_ptr<typeClass> >( "typeClass" )
        .def("sizeof", &typeClass::size )
        .def("offset", &typeClass::getOffset )
        .def("__str__", &typeClass::print);
    boost::python::class_<typedVarClass, boost::python::bases<typeClass>, boost::shared_ptr<typedVarClass> >( "typedVarClass" )
        .def("getAddress", &typedVarClass::getAddress );
    boost::python::class_<dbgModuleClass>( "dbgModuleClass" )
        .def("begin", &dbgModuleClass::getBegin )
        .def("end", &dbgModuleClass::getEnd )
        .def("name", &dbgModuleClass::getName )
        .def("contain", &dbgModuleClass::contain )
        .def("image", &dbgModuleClass::getImageSymbolName )
        .def("pdb", &dbgModuleClass::getPdbName )
        .def("addSynSymbol", &dbgModuleClass::addSyntheticSymbol )
        .def("delAllSynSymbols", &dbgModuleClass::delAllSyntheticSymbols )
        .def("delSynSymbol", &dbgModuleClass::delSyntheticSymbol )
        .def("delSynSymbolsMask", &dbgModuleClass::delSyntheticSymbolsMask )
        .def("__getattr__", &dbgModuleClass::getOffset )
		.def("__str__", &dbgModuleClass::print );
    boost::python::class_<dbgExtensionClass>( 
            "ext",
            "windbg extension",
             boost::python::init<const char*>( boost::python::args("path"), "__init__  dbgExtensionClass" ) ) 
        .def("call", &dbgExtensionClass::call )
		.def("__str__", &dbgExtensionClass::print );
    boost::python::class_<dbgStackFrameClass>( "dbgStackFrameClass", "dbgStackFrameClass" )
        .def_readonly( "instructionOffset", &dbgStackFrameClass::InstructionOffset )
        .def_readonly( "returnOffset", &dbgStackFrameClass::ReturnOffset )
        .def_readonly( "frameOffset", &dbgStackFrameClass::FrameOffset )
        .def_readonly( "stackOffset", &dbgStackFrameClass::StackOffset )
        .def_readonly( "frameNumber", &dbgStackFrameClass::FrameNumber )
		.def( "__str__", &dbgStackFrameClass::print );
    boost::python::class_<dbgOut>( "windbgOut", "windbgOut" )
        .def( "write", &dbgOut::write );
    boost::python::class_<dbgIn>( "windbgIn", "windbgIn" )
        .def( "readline", &dbgIn::readline );                
    boost::python::class_<dbgBreakpointClass>( 
         "bp",
         "break point",
         boost::python::init<ULONG64>( boost::python::args("offset"), "__init__  dbgBreakpointClass" ) ) 
        .def( "set", &dbgBreakpointClass::set )
        .def( "remove", &dbgBreakpointClass::remove )
		.def( "__str__", &dbgBreakpointClass::print );
        
}    

/////////////////////////////////////////////////////////////////////////////////

HRESULT
CALLBACK
DebugExtensionInitialize(
    OUT PULONG  Version,
    OUT PULONG  Flags )
{
    *Version = DEBUG_EXTENSION_VERSION( 1, 0 );
    *Flags = 0;

    PyImport_AppendInittab("pykd", initpykd ); 

    Py_Initialize();

    windbgGlobalSession = new WindbgGlobalSession();

    return setDbgSessionStarted();
}


VOID
CALLBACK
DebugExtensionUninitialize()
{
    DbgEventCallbacks::Stop();

    delete windbgGlobalSession;
    windbgGlobalSession = NULL;

    Py_Finalize();
}


void
SetupDebugEngine( IDebugClient4 *client, DbgExt *dbgExt  )
{
    client->QueryInterface( __uuidof(IDebugClient), (void **)&dbgExt->client );
    client->QueryInterface( __uuidof(IDebugClient4), (void **)&dbgExt->client4 );
    
    
    client->QueryInterface( __uuidof(IDebugControl), (void **)&dbgExt->control );
    client->QueryInterface( __uuidof(IDebugControl4), (void **)&dbgExt->control4 );
    
    client->QueryInterface( __uuidof(IDebugRegisters), (void **)&dbgExt->registers );
    
    client->QueryInterface( __uuidof(IDebugSymbols), (void ** )&dbgExt->symbols );
    client->QueryInterface( __uuidof(IDebugSymbols2), (void ** )&dbgExt->symbols2 );    
    client->QueryInterface( __uuidof(IDebugSymbols3), (void ** )&dbgExt->symbols3 );      
    
    client->QueryInterface( __uuidof(IDebugDataSpaces), (void **)&dbgExt->dataSpaces );
    client->QueryInterface( __uuidof(IDebugDataSpaces4), (void **)&dbgExt->dataSpaces4 );
    
    client->QueryInterface( __uuidof(IDebugAdvanced2), (void **)&dbgExt->advanced2 );
    
    client->QueryInterface( __uuidof(IDebugSystemObjects), (void**)&dbgExt->system );
    client->QueryInterface( __uuidof(IDebugSystemObjects2), (void**)&dbgExt->system2 );
}

/////////////////////////////////////////////////////////////////////////////////    
    
HRESULT 
CALLBACK
py( PDEBUG_CLIENT4 client, PCSTR args)
{

    PyThreadState   *globalInterpreter = PyThreadState_Swap( NULL );
    PyThreadState   *localInterpreter = Py_NewInterpreter();

    try {
    
        DbgExt      ext = { 0 };
        SetupDebugEngine( client, &ext );
        dbgExt = &ext;

        boost::python::import( "pykd" ); 

        boost::python::object       main =  boost::python::import("__main__");

        boost::python::object       global(main.attr("__dict__"));
        
        // ïåðåíàïðàâëåíèå ñòàíäàðòíûõ ïîòîêîâ ÂÂ
        boost::python::object       sys = boost::python::import("sys");
        
        dbgOut                      dout;
        sys.attr("stdout") = boost::python::object( dout );

        dbgIn                       din;
        sys.attr("stdin") = boost::python::object( din );   
       
        // ðàçáîð ïàðàìåòðîâ
        typedef  boost::escaped_list_separator<char>    char_separator_t;
        typedef  boost::tokenizer< char_separator_t >   char_tokenizer_t;  
        
        std::string                 argsStr( args );
        
        char_tokenizer_t            token( argsStr , char_separator_t( "", " \t", "\"" ) );
        std::vector<std::string>    argsList;
        
        for ( char_tokenizer_t::iterator   it = token.begin(); it != token.end(); ++it )
        {
            if ( *it != "" )
                argsList.push_back( *it );
        }            
            
        if ( argsList.size() == 0 )
            return S_OK;      
            
        char    **pythonArgs = new char* [ argsList.size() ];
     
        for ( size_t  i = 0; i < argsList.size(); ++i )
            pythonArgs[i] = const_cast<char*>( argsList[i].c_str() );
            
        PySys_SetArgv( (int)argsList.size(), pythonArgs );

        delete[]  pythonArgs;            

        
        // íàéòè ïóòü ê ôàéëó
        std::string     fullFileName;
        std::string     filePath;
        DbgPythonPath   dbgPythonPath;
        
        if ( dbgPythonPath.findPath( argsList[0], fullFileName, filePath ) )
        {
            DWORD       oldCurDirLen = GetCurrentDirectoryA( 0, NULL );

	    std::vector<char> oldCurDirCstr(oldCurDirLen);
            
            GetCurrentDirectoryA( oldCurDirLen, &oldCurDirCstr[0] );
            
            SetCurrentDirectoryA( filePath.c_str() );
            
            try {                  
            
                boost::python::object       result;
        
                result =  boost::python::exec_file( fullFileName.c_str(), global, global );
                
            }                
            catch( boost::python::error_already_set const & )
            {
                // îøèáêà â ñêðèïòå
                PyObject    *errtype = NULL, *errvalue = NULL, *traceback = NULL;
                
                PyErr_Fetch( &errtype, &errvalue, &traceback );
                
                if(errvalue != NULL) 
                {
                    PyObject *s = PyObject_Str(errvalue);
                    
                    dbgExt->control->Output( DEBUG_OUTPUT_ERROR, "%s/n", PyString_AS_STRING( s ) );

                    Py_DECREF(s);
                }

                Py_XDECREF(errvalue);
                Py_XDECREF(errtype);
                Py_XDECREF(traceback);        
            }  
            
            SetCurrentDirectoryA( &oldCurDirCstr[0] );
        }
        else
        {
      		dbgExt->control->Output( DEBUG_OUTPUT_ERROR, "script file not found\n" );
        }           
    }
   
    catch(...)
    {           
    }     
    
    Py_EndInterpreter( localInterpreter ); 
    PyThreadState_Swap( globalInterpreter );
    
    return S_OK;  
}

/////////////////////////////////////////////////////////////////////////////////  

HRESULT 
CALLBACK
pycmd( PDEBUG_CLIENT4 client, PCSTR args )
{
    try {


        DbgExt      ext = { 0 };

        SetupDebugEngine( client, &ext );
        dbgExt = &ext;

        if ( !std::string( args ).empty() )
        {
            try
            {
                boost::python::exec( args, windbgGlobalSession->global(), windbgGlobalSession->global() );
            }
            catch( boost::python::error_already_set const & )
            {
                // îøèáêà â ñêðèïòå
                PyObject    *errtype = NULL, *errvalue = NULL, *traceback = NULL;
                
                PyErr_Fetch( &errtype, &errvalue, &traceback );
                
                if(errvalue != NULL) 
                {
                    PyObject *s = PyObject_Str(errvalue);
                    
                    dbgExt->control->Output( DEBUG_OUTPUT_ERROR, "%s\n", PyString_AS_STRING( s )  );
                    
                    Py_DECREF(s);
                }

                Py_XDECREF(errvalue);
                Py_XDECREF(errtype);
                Py_XDECREF(traceback);        
            }  
        }
        else
        {
            char        str[100];
            ULONG       inputSize;
            bool        stopInput = false;

            do {
            
                std::string     output;
                
                dbgExt->control->Output( DEBUG_OUTPUT_NORMAL, ">>>" );
                    
                do {    
                                
                    OutputReader        outputReader( dbgExt->client );                    
                    
                    HRESULT   hres = dbgExt->control->Input( str, sizeof(str), &inputSize );
                
                    if ( FAILED( hres ) || std::string( str ) == "" )
                    {
                       stopInput = true;
                       break;
                    }                       
                    
                } while( FALSE );                    
                
                if ( !stopInput )
                    try {
                        boost::python::exec( str, windbgGlobalSession->global(), windbgGlobalSession->global() );
                    }
                    catch( boost::python::error_already_set const & )
                    {
                        // îøèáêà â ñêðèïòå
                        PyObject    *errtype = NULL, *errvalue = NULL, *traceback = NULL;
                        
                        PyErr_Fetch( &errtype, &errvalue, &traceback );
                        
                        if(errvalue != NULL) 
                        {
                            PyObject *s = PyObject_Str(errvalue);
                            
                            dbgExt->control->Output( DEBUG_OUTPUT_ERROR, "%s/n", PyString_AS_STRING( s ) );
                            
                            Py_DECREF(s);
                        }

                        Py_XDECREF(errvalue);
                        Py_XDECREF(errtype);
                        Py_XDECREF(traceback);        
                    }  
                    
            } while( !stopInput );                                
        }
    }
  
    catch(...)
    {           
    }     
    
    return S_OK;          
            
}

///////////////////////////////////////////////////////////////////////////////// 

HRESULT 
CALLBACK
pythonpath( PDEBUG_CLIENT4 client, PCSTR args )
{
    DbgExt      ext = { 0 };

    SetupDebugEngine( client, &ext );  
    dbgExt = &ext;

    //DbgPrint::dprintln( dbgPythonPath.getStr() );

    return S_OK;
}

/////////////////////////////////////////////////////////////////////////////////