TnbConsoleHooker.h

[詳解]

#pragma once
#include <Wincon.h>
#ifdef __TNBPROCESSHANDLE_H__
 #ifndef _TnbPROCESS_PipeInOut_ENABLE
  #error Please define _TnbPROCESS_PipeInOut_ENABLE first rather than you include TnbProcessHandle.h. 
 #endif
 #ifndef _TnbPROCESS_EnumProcesses_ENABLE
  #error Please define _TnbPROCESS_EnumProcesses_ENABLE first rather than you include TnbProcessHandle.h. 
 #endif
#endif
 
#ifndef _TnbDOXYGEN //Document作成用シンボル
 #define _TnbPROCESS_PipeInOut_ENABLE
 #define _TnbPROCESS_EnumProcesses_ENABLE
#endif
#include "TnbProcessHandle.h"
#include "TnbStrVector.h"
#include "TnbStrOperator.h"
 
 
 
//TNB Library
namespace TNB
{
 
 
 
class CConsoleHooker : CCopyImpossible
{
public:
 
    CConsoleHooker(void) : m_processId(0), m_isHooked(false), m_hStdOutput(NULL), m_hStdInput(NULL)
    {
        m_rep.SetKind(_T("ConsoleHooker"));
    }
 
    ~CConsoleHooker(void)
    {
        Unhook();
    }
 
    bool IsHooked(void) const
    {
        return m_isHooked || m_proc.IsAlive();
    }
 
    bool IsAlived(void) const
    {
        return (m_processId != 0 && ::GetProcessVersion(m_processId) != 0);
    }
 
    DWORD GetProcessId(void) const
    {
        return m_processId;
    }
 
    bool Hook(DWORD processId)
    {
        m_processName.Format(_T("%d"), processId);
        return m_Hook(processId);
    }
 
    bool Hook(LPCTSTR lpszProcessName)
    {
        m_processName = lpszProcessName;
        DWORD id = 0;
        CWorkMemT<DWORD> ids;
        CTickCount tc;
        do
        {
            if ( CProcessHandle::EnumProcesses(ids, lpszProcessName) )
            {
                if ( ids.GetSize() > 0 )
                {
                    id = ids[0];
                    break;
                }
            }
        }
        while ( ! tc.IsPassed(2000) );
        return m_Hook(id);
    }
 
    bool Execute(LPCTSTR lpszProcessName, bool isShow = false)
    {       
        Unhook();
        m_pipeOut.Open();
        m_pipeIn.Open();
        if ( m_proc.Create(lpszProcessName, &m_pipeIn, &m_pipeOut, NULL, isShow) )
        {
            m_processId = m_proc.GetProcessId();
            CStr s = lpszProcessName;
            s = m_processName.FindCut(' ');
            if ( s.IsEmpty() )
            {
                s = lpszProcessName;
            }
            int r = ToInt(s.ReverseFindOneOf(_T("\\/")));
            m_processName = (r < 0) ? s : s.Mid(r + 1);
            return true;
        }
        return false;
    }
 
    void Terminate(void)
    {
        Unhook();
        m_proc.Terminate();
    }
 
    void Unhook(void)
    {
        m_pipeIn.Close();
        m_pipeOut.Close();
        if ( m_isHooked )
        {
            ::FreeConsole();
            m_hStdOutput = NULL;
            m_hStdInput = NULL;
            m_isHooked = false;
        }
    }
 
    bool WaitForExit(DWORD to)
    {
        CTickCount tc;
        do
        {
            if ( ::GetProcessVersion(m_processId) == 0 )
            {
                return true;
            }
            ::Sleep(1);
        }
        while ( ! tc.IsPassed(to) );
        return false;
    }
 
    bool SendString(LPCTSTR lpsz)
    {
        if ( m_isHooked )
        {
            CStr s = lpsz;
            size_t len = s.GetLength();
            loop ( i, len )
            {
                if ( ! m_SendChar(s[i]) )
                {
                    return false;
                }
            }
            return true;
        }
        CAscii a = lpsz;
        a.Replace('\r', '\n');
        LPCSTR P = a;
        size_t l = a.GetLength();
        return (m_pipeIn.Send(l, P) == l);
    }
 
    CStrVector GetString(void)
    {
        CStrVector vs;
        if ( m_isHooked )
        {
            m_WaitIdle();
            CONSOLE_SCREEN_BUFFER_INFO sb;
            if ( ::GetConsoleScreenBufferInfo(m_hStdOutput, &sb) )
            {
                size_t len = sb.dwCursorPosition.Y * sb.dwSize.X + sb.dwCursorPosition.X;
                //
                COORD coord = { 0, 0 };
                size_t sz = len;
                CWorkMemT<TCHAR> buf(sb.dwSize.X + 1);
                DWORD dw;
                while ( sz > 0 )
                {
                    buf.Fill(0);
                    size_t z = min(sb.dwSize.X, ToInt(sz));
                    ::ReadConsoleOutputCharacter(m_hStdOutput, buf, ToDword(z), coord, &dw);
                    coord.Y++;
                    sz -= z;
                    CStr s = buf.Ref();
                    s.TrimRight(' ');
                    vs.Add(s);
                    TRACE2("%s [%s]\n", m_processName, s);
                    m_rep.ToViewer(CStr::Fmt(_T("%s [%s]\n"), m_processName, s));
                }
                coord.Y = 0;
                ::FillConsoleOutputCharacter(m_hStdOutput, ' ', ToDword(len), coord, &dw);
                ::SetConsoleCursorPosition(m_hStdOutput, coord);
            }
        }
        else
        {
            size_t l = m_pipeOut.GetReceivingDataSize();
            if ( l > 0 && l != INVALID_SIZE )
            {
                CWorkMemT<char> work(l + 1);
                work.Fill(0); //終端
                m_pipeOut.Receive(l, work);
                CAsciiVector va = CAsciiOperator::SeparateLine(work);
                loop ( i, va )
                {
                    vs.Add(CStr(va[i]));
                    TRACE2("%s [%s]\n", m_processName, CStr(va[i]));
                    m_rep.ToViewer(CStr::Fmt(_T("%s [%s]\n"), m_processName, CStr(va[i])));
                }
            }
        }
        return vs;
    }
 
    bool WaitString(LPCTSTR lpsz, DWORD to)
    {
        m_lastStrings.RemoveAll();
        CStr a = lpsz;
        CTickCount tc;
        do
        {
            if ( ::GetProcessVersion(m_processId) == 0 )
            {
                return false;
            }
            CStrVector vs = GetString();
            m_lastStrings.Append(vs);
            loop ( i, vs )
            {
                if ( vs[i].Find(a) != INVALID_INDEX )
                {
                    return true;
                }
            }
            ::Sleep(1);
        }
        while ( ! tc.IsPassed(to) );
        return false;
    }
 
    int WaitStringEx(const CStrVector& words, DWORD to)
    {
        m_lastStrings.RemoveAll();
        CTickCount tc;
        do
        {
            if ( ::GetProcessVersion(m_processId) == 0 )
            {
                return -2;
            }
            CStrVector vs = GetString();
            m_lastStrings.Append(vs);
            loop ( j, words )
            {
                loop ( i, vs )
                {
                    if ( vs[i].Find(words[j]) != INVALID_INDEX )
                    {
                        return ToInt(j);
                    }
                }
            }
            ::Sleep(1);
        }
        while ( ! tc.IsPassed(to) );
        return -1;
    }
 
    int WaitStringEx(LPCTSTR lpszWord1, LPCTSTR lpszWord2, DWORD to)
    {
        CStrVector vs;
        vs.Add(lpszWord1);
        vs.Add(lpszWord2);
        return WaitStringEx(vs, to);
    }
 
    const CStrVector& GetLastString(void) const
    {
        return m_lastStrings;
    }
 
    void SetLastString(const CStrVector& vs)
    {
        m_lastStrings = vs;
    }
 
private:
    bool m_Hook(DWORD processId)
    {
        Unhook();
        if ( ::AttachConsole(processId) )
        {
            m_hStdOutput = ::GetStdHandle(STD_OUTPUT_HANDLE);
            m_hStdOutput = reinterpret_cast<HANDLE>(7);
            m_hStdInput = ::GetStdHandle(STD_INPUT_HANDLE);
            if ( m_hStdOutput != NULL && m_hStdInput != NULL )
            {
                m_processId = processId;
                m_isHooked = true;
                return true;
            }
        }
        _GetLastError("AttachConsole");
        return false;
    }
    bool m_WaitIdle(DWORD to = 500)
    {
        ASSERT( m_isHooked );
        CONSOLE_SCREEN_BUFFER_INFO sb1;
        CONSOLE_SCREEN_BUFFER_INFO sb;
        if ( ! ::GetConsoleScreenBufferInfo(m_hStdOutput, &sb1) )
        {
            _GetLastError("GetConsoleScreenBufferInfo");
            return false;
        }
        while ( true )
        {
            CTickCount tc;
            bool isIdle = true;
            do
            {
                ::Sleep(100);
                if ( ! ::GetConsoleScreenBufferInfo(m_hStdOutput, &sb) )
                {
                    return false;
                }
                if ( sb1.dwCursorPosition.X != sb.dwCursorPosition.X || sb1.dwCursorPosition.Y != sb.dwCursorPosition.Y )
                {
                    isIdle = false;
                    break;
                }
            }
            while ( ! tc.IsPassed(to) );
            if ( isIdle )
            {
                return true;
            }
            sb1 = sb;
        }
        return false;
    }
    bool m_SendChar(TCHAR c)
    {
        ASSERT( m_isHooked );
        INPUT_RECORD ir[2] = { 0 };
        ir[0].EventType = KEY_EVENT;
        #ifndef _UNICODE
            ir[0].Event.KeyEvent.uChar.AsciiChar = c;
        #else
            ir[0].Event.KeyEvent.uChar.UnicodeChar = c;
        #endif
        ir[1] = ir[0];
        ir[0].Event.KeyEvent.bKeyDown = true;
        DWORD dw;
        if ( ::WriteConsoleInput(m_hStdInput, ir, 2, &dw) )
        {
            return dw == 2;
        }
        return false;
    }
    CStr            m_processName;  
    DWORD           m_processId;    
    bool            m_isHooked;     
    HANDLE          m_hStdOutput;   
    HANDLE          m_hStdInput;    
    CPipe           m_pipeIn;       
    CPipe           m_pipeOut;      
    CProcessHandle  m_proc;         
    CStrVector      m_lastStrings;  
    CReportViewerEx m_rep;          
};
 
 
 
}; // TNBLIB