Detect Driver

#include "stdafx.h"
#include "FileForm.h"
#include "winioctl.h"
#include "DetectDriver_Ioctl.h"
#include "DriverWork.h"
#include "PackageParser.h"
#include "FolderDlg.h"

#include <stdexcept>
#include <string>
#include <list>
#include <sstream>

#define TIMER_QUERY_STATUS_ID     0xAA
#define TIMER_QUERY_STATUS_ELAPSE 300 // 300 ms

IMPLEMENT_DYNCREATE(FileForm,ListForm)

BEGIN_MESSAGE_MAP(FileForm, ListForm)
    ON_WM_TIMER()
END_MESSAGE_MAP()

FileForm::FileForm()
    : isRunning_(false)
    , totalCount_(0)
{}

void FileForm::InitializeListCtrl(CListCtrl& listCtrl)
{
	listCtrl.InsertColumn(0,_T("File path"),LVCFMT_LEFT,350,0);

    pListCtrl_ = &listCtrl;
}

void ParseOutputString(wchar_t* buffer,
                       size_t bufSize,
                       WStringList* hiddenProcList)
{
    std::wstring inputString(buffer,bufSize);
    utils::UnPackByEOL(inputString,hiddenProcList);
}

void QueryHiddenFiles(WStringList* hiddenProcList)
{
    std::vector<char> data;
    utils::DriverWork::ExchangeWithExceptionTranslation(
        _T("\\\\.\\DetectDriver"),
        IOCTL_FILES_QUERY_HIDEN,
        &data);

    ParseOutputString((wchar_t*)&data.front(),data.size()/2,hiddenProcList);
}

void QueryStatus(ULONG* status,
                 ULONG* flag,
                 ULONG* totalCount)
{
    std::vector<char> data;
    utils::DriverWork::ExchangeWithExceptionTranslation(
        _T("\\\\.\\DetectDriver"),
        IOCTL_FILES_QUERY_STATUS,
        &data);
    
    if( data.size() != 3*sizeof(ULONG) )
        throw std::exception("Error: Driver return wrong status size");
    
    memcpy(status    ,&data.front()                ,sizeof(ULONG));
    memcpy(flag      ,&data.front()+1*sizeof(ULONG),sizeof(ULONG));
    memcpy(totalCount,&data.front()+2*sizeof(ULONG),sizeof(ULONG));
}

void InsertToList(CListCtrl& listCtrl,const WStringList& hiddenProcList)
{
    WStringList::const_iterator it = hiddenProcList.begin();
    for( ; it != hiddenProcList.end() ; ++it )
    {
        int nIndex = listCtrl.GetItemCount();
        listCtrl.InsertItem(nIndex,it->c_str());
    }
}

void FileForm::OnButtonClick(CListCtrl& listCtrl)
{
    if( !isRunning_ )
    {
        CFolderDialog dirDlg;

        if( dirDlg.DoModal() != IDOK )
            return;

        std::wstring dirPath = dirDlg.GetFolderPath();
        if( dirPath.empty() )
            return;

        // Start
        std::vector<char> data;
        data.assign( (char*)dirPath.c_str(), (char*)dirPath.c_str()+dirPath.size()*2);

        utils::DriverWork::ExchangeWithExceptionTranslation(
            _T("\\\\.\\DetectDriver"),
            IOCTL_FILES_START_DETECT,
            &data);

        listCtrl.DeleteAllItems();

        isRunning_ = true;
        SetTimer(TIMER_QUERY_STATUS_ID,TIMER_QUERY_STATUS_ELAPSE,NULL);   

        SetStatus(L"Started");
        SetDlgItemText(IDC_BUTTON_START_STOP,_T("Stop"));
    }
    else
    {
        // Stop
        std::vector<char> data;
        utils::DriverWork::ExchangeWithExceptionTranslation(
            _T("\\\\.\\DetectDriver"),
            IOCTL_FILES_STOP_DETECT,
            &data);

        isRunning_ = false;
        KillTimer(TIMER_QUERY_STATUS_ID);
        
        std::wstring counters;
        ConvertCountersToString(&counters);
        SetStatus(L"Stopped. "+counters);

        SetDlgItemText(IDC_BUTTON_START_STOP,_T("Start"));
    }
}

void UpdateListCtrlImpl(CListCtrl& listCtrl)
{
    WStringList hiddenProcList;
    QueryHiddenFiles(&hiddenProcList);

    InsertToList(listCtrl,hiddenProcList);
}

void FileForm::ConvertCountersToString(std::wstring* counters)
{
    *counters = L"Directory processed: ";

    std::wstringstream stream;
    stream << totalCount_;
    
    std::wstring totalCountStr;
    stream >> totalCountStr;

    *counters += totalCountStr;
}

void FileForm::OnTimerImpl()
{
    ULONG status,flag;
    QueryStatus(&status,&flag,&totalCount_);

    std::wstring counters;
    ConvertCountersToString(&counters);

    switch(status)
    {
    case FILES_DETECT_STATUS_RUNNING:
        SetStatus(L"Running. " + counters);
        break;
    case FILES_DETECT_STATUS_FINISHED:
        SetStatus(L"Finished. " + counters);
        SetDlgItemText(IDC_BUTTON_START_STOP,_T("Start"));
        isRunning_ = false;
        break;
    case FILES_DETECT_STATUS_ABORTED:
        SetStatus(L"Aborted. " + counters);
        break;
    }

    switch(flag)
    {
    case FILES_DETECT_FLAG_NEW_HIDDEN_FILE:
        UpdateListCtrlImpl(*pListCtrl_);
        break;
    case FILES_DETECT_FLAG_ERROR:
        {
            std::vector<char> data;
            utils::DriverWork::ExchangeWithExceptionTranslation(
                _T("\\\\.\\DetectDriver"),
                IOCTL_FILES_QUERY_ERROR,
                &data);
        }
        break;
    }
}

void FileForm::OnTimer(UINT_PTR nIDEvent)
{
    if( nIDEvent == TIMER_QUERY_STATUS_ID )
    {
        KillTimer(TIMER_QUERY_STATUS_ID);
        
        try
        {
            OnTimerImpl();
            
            // Resume checking driver only if all calls is succeed
            // And detection is in progress
            if( isRunning_ )
                SetTimer(TIMER_QUERY_STATUS_ID,TIMER_QUERY_STATUS_ELAPSE,NULL);
        }
        catch(const std::exception& ex)
        {
            ::MessageBoxA(this->GetSafeHwnd(),ex.what(),"Error",MB_ICONERROR|MB_OK);
        }
    }
    ListForm::OnTimer(nIDEvent);
}