Device Insertion Monitor (DIM)

// I call it a device insertion monitor, and it is used to monitor any media or USB devices that are inserted or plugged into the PC.
// This can also serve as a virus transmission tool since a virus can replicate itself upon media/device detection to be transmitted to // other victims via USB drives for example. But that is another subject that we can discuss later...
// For now, this program is capable of detecting CD/DVD media and USB devices and capturing fundamental information. I will work on it when //I get home, to print more data on the machines...
// A notable mention is that I used the RegisterDeviceNotification API to "record" USB devices.
// This is absolutely necessary since DBT_DEVICEARRIVAL / DBT_DEVICEREMOVECOMPLETE messages are sent by default for CD/DVD ROMs, but not for USBs.

#include <windows.h>
#include <dbt.h>
#include <assert.h>
#include <atlstr.h>
#include <Setupapi.h>
#include <GdiPlus.h>
using namespace Gdiplus;

#pragma comment(lib, "gdiplus")
#pragma comment(lib, "Setupapi")

char *dupcat(const char *s1, ...){
    int len;
    char *p, *q, *sn;
    va_list ap;

    len = strlen(s1);
    va_start(ap, s1);
    while (1) {
        sn = va_arg(ap, char *);
        if (!sn)
            break;
        len += strlen(sn);
    }
    va_end(ap);

    p = new char[len + 1];
    strcpy(p, s1);
    q = p + strlen(p);

    va_start(ap, s1);
    while (1) {
        sn = va_arg(ap, char *);
        if (!sn)
            break;
        strcpy(q, sn);
        q += strlen(q);
    }
    va_end(ap);

    return p;
}

//Global id for raw USB devices
GUID raw = {0xa5dcbf10, 0x6530, 0x11d2, 0x90, 0x1f, 0x00, 0xc0, 0x4f, 0xb9, 0x51, 0xed };

char FirstDriveFromMask(ULONG unitmask)
{
    char i;

    for (i = 0; i < 26; ++i)
    {
        if (unitmask & 0x1)
            break;
        unitmask = unitmask >> 1;
    }

    return(i + 'A');
}

void AddText(HWND edit, LPCTSTR Text)
{
    int len = GetWindowTextLength(edit);
    SendMessage(edit, EM_SETSEL, (WPARAM)len, (LPARAM)len);
    SendMessage(edit, EM_REPLACESEL, 0, (LPARAM)Text);
}

#define APP MAKEINTRESOURCE(101)
#define APPSMALL MAKEINTRESOURCE(102)
HWND hwnd, devType, devFriendlyName, devInfo;
LRESULT CALLBACK WndProc(HWND, UINT, WPARAM, LPARAM);
static char gszClassName[] = "Ath";
static HINSTANCE ghInstance = NULL;

/*
    The following code is not mine!
    It is copied and modified for my needs from this article http://www.codeproject.com/Articles/14500/Detecting-Hardware-Insertion-and-or-Removal
    Credits to Sam NG
*/

TCHAR *UpdateDevice(PDEV_BROADCAST_DEVICEINTERFACE pDevInf, WPARAM wParam)
{
    // dbcc_name:
    // \\?\USB#Vid_04e8&Pid_503b#0002F9A9828E0F06#{a5dcbf10-6530-11d2-901f-00c04fb951ed}
    // convert to
    // USB\Vid_04e8&Pid_503b\0002F9A9828E0F06
    assert(lstrlen(pDevInf->dbcc_name) > 4);
    CString szDevId = pDevInf->dbcc_name + 4;
    int idx = szDevId.ReverseFind(_T('#'));
    assert(-1 != idx);
    szDevId.Truncate(idx);
    szDevId.Replace(_T('#'), _T('\\'));
    szDevId.MakeUpper();

    CString szClass;
    idx = szDevId.Find(_T('\\'));
    assert(-1 != idx);
    szClass = szDevId.Left(idx);

    // if we are adding device, we only need present devices
    // otherwise, we need all devices
    DWORD dwFlag = DBT_DEVICEARRIVAL != wParam
        ? DIGCF_ALLCLASSES : (DIGCF_ALLCLASSES | DIGCF_PRESENT);
    HDEVINFO hDevInfo = SetupDiGetClassDevs(NULL, szClass, NULL, dwFlag);
    if (INVALID_HANDLE_VALUE == hDevInfo)
    {
        return NULL;
    }

    SP_DEVINFO_DATA* pspDevInfoData =
        (SP_DEVINFO_DATA*)HeapAlloc(GetProcessHeap(), 0, sizeof(SP_DEVINFO_DATA));
    pspDevInfoData->cbSize = sizeof(SP_DEVINFO_DATA);
    for (int i = 0; SetupDiEnumDeviceInfo(hDevInfo, i, pspDevInfoData); i++)
    {
        DWORD DataT;
        DWORD nSize = 0;
        TCHAR buf[MAX_PATH];

        if (!SetupDiGetDeviceInstanceId(hDevInfo, pspDevInfoData, buf, sizeof(buf), &nSize))
        {
            break;
        }

        if (szDevId == buf)
        {
            // device found
            if (SetupDiGetDeviceRegistryProperty(hDevInfo, pspDevInfoData,
                SPDRP_FRIENDLYNAME, &DataT, (PBYTE)buf, sizeof(buf), &nSize)) {
                // do nothing
            }
            else if (SetupDiGetDeviceRegistryProperty(hDevInfo, pspDevInfoData,
                SPDRP_DEVICEDESC, &DataT, (PBYTE)buf, sizeof(buf), &nSize)) {
                // do nothing
            }
            else {
                lstrcpy(buf, _T("Unknown"));
            }
            TCHAR *ret = new TCHAR[260];
            _tcscpy(ret, buf);
            return ret;
            // update UI
            // .....
            // .....
        }
    }

    if (pspDevInfoData) HeapFree(GetProcessHeap(), 0, pspDevInfoData);
    SetupDiDestroyDeviceInfoList(hDevInfo);
    return NULL;
}

void GDIPLUS(HDC hdc, int drvIndex = -1){
    Graphics graphics(hdc);
    wchar_t *title = L"Device insertion monitor";
    wchar_t *author = L"Athenian - Rohitab forums";

    //labels
    wchar_t *dev_type = L"Device type: ";
    wchar_t *dev_name = L"Device name: ";
    wchar_t *dev_info = L"Device information: ";

    FontFamily family(L"Verdana");
    FontFamily family2(L"Calibri");
    Font font(&family, 25, FontStyleRegular, UnitPixel);
    Font font2(&family2, 15, FontStyleRegular, UnitPixel);

    SolidBrush sbrush(Color::White);
    graphics.DrawString(title, wcslen(title), &font, PointF(20, 20), &sbrush);
    graphics.DrawString(author, wcslen(author), &font2, PointF(20, 60), &sbrush);

    //draw labels
    graphics.DrawString(dev_type, wcslen(dev_type), &font2, PointF(20, 120), &sbrush);
    graphics.DrawString(dev_name, wcslen(dev_name), &font2, PointF(20, 180), &sbrush);
    graphics.DrawString(dev_info, wcslen(dev_info), &font2, PointF(20, 240), &sbrush);
}

int WINAPI WinMain(HINSTANCE hInstance,
    HINSTANCE hPrevInstance,
    LPSTR lpCmdLine,
    int nCmdShow)
{
    WNDCLASSEX WndClass;
    MSG Msg;
    ghInstance = hInstance;

    WndClass.cbSize = sizeof(WNDCLASSEX);
    WndClass.style = NULL;
    WndClass.lpfnWndProc = WndProc;
    WndClass.cbClsExtra = 0;
    WndClass.cbWndExtra = 0;
    WndClass.hInstance = ghInstance;
    WndClass.hIcon = LoadIcon(hInstance, APP);
    WndClass.hCursor = LoadCursor(NULL, IDC_ARROW);
    WndClass.hbrBackground = CreateSolidBrush(RGB(20, 100, 200));
    WndClass.lpszMenuName = 0;
    WndClass.lpszClassName = gszClassName;
    WndClass.hIconSm = LoadIcon(hInstance, APPSMALL);

    RegisterClassEx(&WndClass);
    hwnd = CreateWindowEx(
        0,
        gszClassName,
        "Device monitor",
        WS_OVERLAPPEDWINDOW,
        CW_USEDEFAULT, CW_USEDEFAULT,
        800, 600,
        NULL, NULL,
        ghInstance,
        NULL);

    devType = CreateWindow(
        "edit",
        0,
        WS_CHILD | WS_VISIBLE,
        160, 120,
        350, 20,
        hwnd, 0,
        ghInstance,
        NULL);

    devFriendlyName = CreateWindow(
        "edit",
        0,
        WS_CHILD | WS_VISIBLE,
        160, 180,
        350, 20,
        hwnd, 0,
        ghInstance,
        NULL);

    devInfo = CreateWindow(
        "edit",
        0,
        WS_CHILD | WS_VISIBLE,
        160, 240,
        350, 200,
        hwnd, 0,
        ghInstance,
        NULL);

    ShowWindow(hwnd, 1);
    UpdateWindow(hwnd);

    SendMessage(devType, EM_SETREADONLY, 1, 0);
    SendMessage(devFriendlyName, EM_SETREADONLY, 1, 0);
    SendMessage(devInfo, EM_SETREADONLY, 1, 0);

    while (GetMessage(&Msg, NULL, 0, 0)) {
        TranslateMessage(&Msg);
        DispatchMessage(&Msg);
    }
    return Msg.wParam;
}

BOOL RegisterDevice(GUID InterfaceClassGuid, HWND __hWnd, HDEVNOTIFY *hDeviceNotify)
{
    DEV_BROADCAST_DEVICEINTERFACE NotificationFilter;

    ZeroMemory(&NotificationFilter, sizeof(NotificationFilter));
    NotificationFilter.dbcc_size = sizeof(DEV_BROADCAST_DEVICEINTERFACE);
    NotificationFilter.dbcc_devicetype = DBT_DEVTYP_DEVICEINTERFACE;
    NotificationFilter.dbcc_classguid = InterfaceClassGuid;

    *hDeviceNotify = RegisterDeviceNotification(
        __hWnd,
        &NotificationFilter,
        DEVICE_NOTIFY_WINDOW_HANDLE
        );

    if (NULL == *hDeviceNotify)
    {
        return FALSE;
    }

    return TRUE;
}

LRESULT CALLBACK WndProc(HWND hwnd, UINT Message, WPARAM wParam, LPARAM lParam) {
    HDC hdc;
    PAINTSTRUCT ps;
    //BRODCAST_HDR is used for generic identification of devices
    PDEV_BROADCAST_HDR lpdb = (PDEV_BROADCAST_HDR)lParam;
    GdiplusStartupInput gdiplusStartupInput;
    ULONG_PTR           gdiplusToken;
    GdiplusStartup(&gdiplusToken, &gdiplusStartupInput, NULL);

    HDEVNOTIFY notify;

    switch (Message) {

    case WM_CREATE:
            //For USB devices,the device type is DBT_DEVTYP_DEVICEINTERFACE
            //DBT_DEVICEARRIVAL/DBT_DEVICEREMOVECOMPLETE are sent by default for all
            //the devices,but not for USB type.
            //We need to register them with RegisterDeviceNotification

        if (!RegisterDevice(raw, hwnd, &notify))
            ExitProcess(0);
        break;

    case WM_DEVICECHANGE:
    {
                            switch (wParam)
                            {
                            case DBT_DEVICEARRIVAL:

                                // Check for USB devices
                                if (lpdb->dbch_devicetype == DBT_DEVTYP_DEVICEINTERFACE){
                                    PDEV_BROADCAST_DEVICEINTERFACE usb = (PDEV_BROADCAST_DEVICEINTERFACE)lpdb;
                                    SetWindowText(devType, "USB");
                                    MessageBox(hwnd, "An usb device has been inserted!", "USB", 0);
                                    TCHAR *name = UpdateDevice(usb, wParam);
                                    if (name)
                                        SetWindowText(devFriendlyName, name);
                                }

                                // Check whether a CD or DVD was inserted into a drive.
                                if (lpdb->dbch_devicetype == DBT_DEVTYP_VOLUME)
                                {
                                    PDEV_BROADCAST_VOLUME lpdbv = (PDEV_BROADCAST_VOLUME)lpdb;

                                    if (lpdbv->dbcv_flags & DBTF_MEDIA)
                                    {
                                        SetWindowText(devType, "CD / DVD");
                                        char *drive = (char *)FirstDriveFromMask(lpdbv->dbcv_unitmask);
                                        MessageBox(0, dupcat("Drive: - Media inserted!", 0) , "CD/DVD", 0);
                                    }
                                }
                                break;

                            case DBT_DEVICEREMOVECOMPLETE:
                                SetWindowText(devType, 0);
                                SetWindowText(devFriendlyName, 0);
                                SetWindowText(devInfo, 0);

                                // Check for USB devices
                                if (lpdb->dbch_devicetype == DBT_DEVTYP_DEVICEINTERFACE){
                                    PDEV_BROADCAST_DEVICEINTERFACE usb = (PDEV_BROADCAST_DEVICEINTERFACE)lpdb;
                                    MessageBox(hwnd, "An usb device has been removed!", "USB", 0);
                                }

                                // Check whether a CD or DVD was removed from a drive.
                                if (lpdb->dbch_devicetype == DBT_DEVTYP_VOLUME)
                                {
                                    PDEV_BROADCAST_VOLUME lpdbv = (PDEV_BROADCAST_VOLUME)lpdb;

                                    if (lpdbv->dbcv_flags & DBTF_MEDIA)
                                    {
                                        char *drive = (char *)FirstDriveFromMask(lpdbv->dbcv_unitmask);
                                        MessageBox(0, dupcat("Drive: - Media removed!", 0), "CD/DVD", 0);
                                    }
                                }
                                break;

                            default:
                                //Process other WM_DEVICECHANGE notifications for other
                                //devices or reasons.
                                ;
                            }
    }
        break;
    case WM_PAINT:
        hdc = BeginPaint(hwnd, &ps);
        GDIPLUS(hdc);
        EndPaint(hwnd, &ps);
        break;
    case WM_CLOSE:
        DestroyWindow(hwnd);
        break;
    case WM_DESTROY:
        PostQuitMessage(0);
        break;
    default:
        return DefWindowProc(hwnd, Message, wParam, lParam);
    }
    return 0;
}