//
// This program demonstrates how to start up a script debugging without using mdm.exe.
// More information about the interfaces used in this sample can be found on MSDN.
//

#include <atlbase.h>

// Use PSAPI for a process listing
#include <psapi.h>
#pragma comment(lib, "psapi.lib")

// Use IELaunchURL for starting IE.
#include <iepmapi.h>
typedef HRESULT (STDAPICALLTYPE *IELAUNCHURL)(
    __in        LPCWSTR lpwstrUrl,
    __inout     PROCESS_INFORMATION *lpProcInfo,
    __in_opt    VOID *lpInfo
    );


// These header files and lib can be found in the Visual Studio SDK
#include "msdbg.h"
#include "activdbg.h"
#pragma comment(lib, "ad2de.lib")

// Engine filter which indicates that we are only interested in 'Script' code
const CONST_GUID_ARRAY ScriptEngineFilter = { 1, &guidScriptEng };


// Summary: 
//   Callback that the PDM will call when script code runs inside the IE process
//
class CMyCallback : public IDebugPortNotify2
{
private:
    volatile LONG m_refCount;
    const HANDLE m_hCodeRanEvent;

    CMyCallback(HANDLE hCodeRanEvent) : 
        m_refCount(1),
        m_hCodeRanEvent(hCodeRanEvent)
    {
    }

protected:
    ~CMyCallback()
    {
        CloseHandle(m_hCodeRanEvent);
    }

public:
    // Summary:
    //   Wait for either script code to start in the IE process, or for IE to exit.
    //
    // Arguments:
    //   hProcess: Handle to the IE process
    void Wait(HANDLE hProcess)
    {
        HANDLE handles[] = { hProcess, m_hCodeRanEvent };
        DWORD index = 0;
        CoWaitForMultipleHandles(0, INFINITE, _countof(handles), handles, &index);

        if (index == 0)
        {
            printf("WARNING: IE exited without running script code.\n");
        }
    }

// Factory method
public:
    // Summary:
    //   Create a new callback object.
    //
    // Arguments:
    //   ppNewObject: Created callback object.
    static HRESULT Create(CMyCallback** ppNewObject)
    {
        *ppNewObject = NULL;

        // Create a manual reset event which this class will signal when script code starts
        // in the IE process.
        HANDLE hEvent = CreateEvent(NULL, TRUE, FALSE, NULL);
        if (hEvent == NULL)
        {
            return HRESULT_FROM_WIN32(GetLastError());
        }

        *ppNewObject = new CMyCallback(hEvent);
        return S_OK;
    }

// IUnknown implementation
public:
    ULONG STDMETHODCALLTYPE AddRef()
    {
        return (ULONG)InterlockedIncrement(&m_refCount);
    }

    ULONG STDMETHODCALLTYPE Release()
    {
        ULONG refCount = (ULONG)InterlockedDecrement(&m_refCount);
        if (refCount == 0)
        {
            delete this;
        }

        return refCount;
    }

    HRESULT STDMETHODCALLTYPE QueryInterface(REFIID riid, void** ppv)
    {
        if (riid == __uuidof(IUnknown) || riid == __uuidof(IDebugPortNotify2))
        {
            *ppv = static_cast<IDebugPortNotify2*>(this);
            AddRef();
            return S_OK;
        }
        else
        {
            *ppv = NULL;
            return E_NOINTERFACE;
        }
    }

// IDebugPortNotify2 implementation
public:
    HRESULT STDMETHODCALLTYPE AddProgramNode( 
        __RPC__in_opt IDebugProgramNode2 *pProgramNode
        )
    {
        printf("Script code started inside the IE process.\n\n");

        // How to complete the launch:
        // 1. QI pProgramNode for IDebugProviderProgramNode2
        // 2. Call UnmarshalDebuggeeInterface passing in __uuidof().
        // 3. Call IRemoteDebugApplication::ConnectDebugger as before

        SetEvent(m_hCodeRanEvent);
        
        return S_OK;
    }
    
    HRESULT STDMETHODCALLTYPE RemoveProgramNode( 
        __RPC__in_opt IDebugProgramNode2 *pProgramNode
        )
    {
        return S_OK;
    }
};

// Summary: 
//   Launches IE with script debugging enabled
//
void LaunchIE()
{
    HRESULT hr;

    CComPtr<IDebugProgramProvider2> pPDM;
    hr = pPDM.CoCreateInstance(__uuidof(MsProgramProvider), NULL, CLSCTX_INPROC_SERVER);
    if (FAILED(hr))
    {
        printf("Failed to load pdm.dll (error code %08x). Ensure that pdm.dll is installed on this computer.\n", hr);
        exit(-1);
    }

    // Creae a callback object which will be passed to the PDM.
    CComPtr<CMyCallback> pMyCallback;
    hr = CMyCallback::Create(&pMyCallback);
    if (FAILED(hr))
    {
        printf("Failed to create callback object.\n");
        exit(-1);
    }

    // NOTE: This code relies on the IE7. Production code may wish to fall back to CreateProcess when this happens.

    HMODULE hIEFrame = LoadLibrary(TEXT("ieframe.dll"));
    if (hIEFrame == NULL)
    {
        printf("Failed to load ieframe.dll. Ensure that IE7 is installed.\n");
        exit(-1);
    }
    IELAUNCHURL pIELaunchUrl = (IELAUNCHURL)GetProcAddress(hIEFrame, "IELaunchURL");

    // Call into ieframe.dll to start IE. It is launched suspended.
    PROCESS_INFORMATION pi;
    IELAUNCHURLINFO launchInfo;
    launchInfo.cbSize = sizeof(launchInfo);
    launchInfo.dwCreationFlags = CREATE_SUSPENDED;
    hr = pIELaunchUrl(L"http://www.live.com/", &pi, &launchInfo);
    if (FAILED(hr))
    {
        printf("Failed to start IE (error code %08x).\n", hr);
        exit(-1);
    }
    
    // Use smart handles to close the returned handles when we exit this scope
    CHandle hIEProcess(pi.hProcess), hIEThread(pi.hThread);

    // Ask pdm.dll to tell us when script code runs. Note that because we aren't
    // supplying a port, the PDM will return E_NOINTERFACE. This is ignorable.
    const AD_PROCESS_ID processId = { AD_PROCESS_ID_SYSTEM, { pi.dwProcessId } };
    hr = pPDM->WatchForProviderEvents(
        PFLAG_DEBUGGEE | PFLAG_REASON_WATCH, // Ask the PDM to inform us when script code runs
        NULL,                                // The PDM implementation of this interface does not require the 'port' parameter
        processId,                           // the process id to query
        ScriptEngineFilter,                  // We are interested in script code
        GUID_NULL,                           // no launching engine
        pMyCallback                          // callback interface
        );
    if (FAILED(hr) && hr != E_NOINTERFACE)
    {
        printf("PDM failed to watch IE process for provider events (error code %08x).\n", hr);
        exit(-1);
    }

    // Now that we have setup our event sink, we can now resume the IE process.
    ResumeThread(pi.hThread);

    printf("IE process now started. Waiting for script code to run...\n");
    pMyCallback->Wait(pi.hProcess);

    // Unregister our event sink with pdm.dll
    hr = pPDM->WatchForProviderEvents(
        0,                                   // Tell the PDM that we want it to stop watching
        NULL,                                // The PDM implementation of this interface does not require the 'port' parameter
        processId,                           // the process id to query
        ScriptEngineFilter,                  // We are interested in script code
        GUID_NULL,                           // no launching engine
        pMyCallback                          // callback interface
        );

    return;
}




// Summary: 
//   Print out a listing of script processes
//
void ProcessListing()
{
    HRESULT hr;

    CComPtr<IDebugProgramProvider2> pPDM;
    hr = pPDM.CoCreateInstance(__uuidof(MsProgramProvider), NULL, CLSCTX_INPROC_SERVER);
    if (FAILED(hr))
    {
        printf("Failed to load pdm.dll (error code %08x). Ensure that pdm.dll is installed on this computer.\n", hr);
        exit(-1);
    }

    // For the purposes of simplicity, we cap the number of processes at 1024. Production 
    // code may wish to loop instead of having a maximum number of processes
    const DWORD MaxProcesses = 1024;

    ATL::CAutoVectorPtr<DWORD> processes;
    processes.Allocate(MaxProcesses);
    DWORD cbNeeded;
    
    if (!EnumProcesses(processes, MaxProcesses * sizeof(DWORD), &cbNeeded))
    {
        DWORD gle = GetLastError();
        printf("EROR: failed to obtain process listing (error code %d).\n", gle);
        exit(-1);
    }

    bool foundProcess = false;
    const DWORD NumProcesses = cbNeeded / sizeof(DWORD);

    for (DWORD cProcess = 0; cProcess < NumProcesses; cProcess++)
    {
        const AD_PROCESS_ID processId = { AD_PROCESS_ID_SYSTEM, { processes[cProcess] } };
        PROVIDER_PROCESS_DATA processData = {0};

        //
        // Ask the PDM for program nodes for this process.
        hr = pPDM->GetProviderProcessData(
            PFLAG_GET_PROGRAM_NODES,   // Ask for the program nodes to be returned
            NULL,                      // The PDM implementation of this interface does not require the 'port' parameter
            processId,                 // the process id to query
            ScriptEngineFilter,        // We are interested in script code
            &processData               // Output
            );

        if (hr == S_OK)
        {
            if (processData.ProgramNodes.dwCount > 0)
            {
                foundProcess = true;
                printf("process %d contains script code.\n", processId.ProcessId.dwProcessId);

                for (DWORD cProgramNode = 0; cProgramNode < processData.ProgramNodes.dwCount; cProgramNode++)
                {
                    processData.ProgramNodes.Members[cProgramNode]->Release();
                }
            }
        }

        // How to implement attach:
        // 1. Implementing attach also starts off with a call to GetProviderProcessData.
        //    In addition to the 'PFLAG_GET_PROGRAM_NODES' flag, 'PFLAG_DEBUGGEE' and
        //    'PFLAG_ATTACHED_TO_DEBUGGEE' should be passed
        // 2. QI the returned program nodes for IDebugProviderProgramNode2.
        // 3. Call UnmarshalDebuggeeInterface passing in __uuidof(IRemoteDebugApplication).
        // 4. Call IRemoteDebugApplication::ConnectDebugger as before.
        //
    }

    if (!foundProcess)
    {
        printf("No currently running processes has script debugging enabled.\n");
    }
}



int _tmain(int argc, _TCHAR* argv[])
{
    CoInitialize(NULL);

    // Demonstrate how to launch IE
    LaunchIE();

    // Demonstrate how to obtain a process listing
    ProcessListing();

    CoUninitialize();

	return 0;
}