#include "wifi.h"

// For listing network adapters
#include <iphlpapi.h>
#pragma comment(lib,"iphlpapi.lib")
// For NDIS ioctl's
#include <winioctl.h>
#include <nuiouser.h>

CWifi::CWifi(unsigned long bufferLength)
: m_pIoCtlBuffer(0)
, m_ulBufferLength(bufferLength)
, m_pDeviceName(0)
{
    m_ulBufferLength+=sizeof(NDISUIO_QUERY_OID);
    m_pIoCtlBuffer = new unsigned char[m_ulBufferLength];
    FindDeviceName();
}
CWifi::~CWifi() {
    if (m_pIoCtlBuffer) {
        delete[] m_pIoCtlBuffer;
        m_pIoCtlBuffer=0;
    }
    if (m_pDeviceName) {
        delete[] m_pDeviceName;
        m_pDeviceName=0;
    }
}
void CWifi::FindDeviceName() {
    // If we already found the device, do nothing
    if (m_pDeviceName) return;
 
    // Enumerate all network adapters
    PIP_ADAPTER_INFO pAdapterInfo;
    PIP_ADAPTER_INFO pAdapter = 0;
 
    ULONG ulOutBufLen = sizeof (IP_ADAPTER_INFO);
    pAdapterInfo = (IP_ADAPTER_INFO *) malloc(sizeof (IP_ADAPTER_INFO));
    if (!pAdapterInfo) return;
    if (GetAdaptersInfo(pAdapterInfo,&ulOutBufLen)==ERROR_BUFFER_OVERFLOW) {
        free(pAdapterInfo);
        pAdapterInfo = (IP_ADAPTER_INFO *) malloc(ulOutBufLen);
        if (!pAdapterInfo) return;
    }
 
    if (NO_ERROR==GetAdaptersInfo(pAdapterInfo, &ulOutBufLen)) {
        pAdapter = pAdapterInfo;
        while (pAdapter) {
 
            // only consider Ethernet adapters
            if (pAdapter->Type==MIB_IF_TYPE_ETHERNET) {
 
                // Open the NDIS driver in order to make a
                // IOCTL_NDISUIO_NIC_STATISTICS ioctl.
                HANDLE ndisAccess=CreateFile(
                    TEXT("UIO1:"),
                    0,0,0,
                    OPEN_EXISTING,
                    FILE_ATTRIBUTE_NORMAL|FILE_FLAG_OVERLAPPED,
                    INVALID_HANDLE_VALUE
                );
                if (ndisAccess && ndisAccess!=INVALID_HANDLE_VALUE) {
                    // Allocate a new unicode string with the adapter name
                    int len=0;
                    while (pAdapter->AdapterName[len])
                        ++len;
                    wchar_t* unicodeName = new wchar_t[len+1];
                    for (int c=0; c<len; ++c)
                        unicodeName[c]=(wchar_t) pAdapter->AdapterName[c];
                    unicodeName[len]=0;
                    // Make the ioctl call
                    NIC_STATISTICS nicStatistics = {0};
                    nicStatistics.ptcDeviceName = unicodeName;
                    nicStatistics.Size = sizeof(NIC_STATISTICS);
                    DWORD dwBytesWritten = 0;
                    bool success=(DeviceIoControl(
                        ndisAccess,
                        IOCTL_NDISUIO_NIC_STATISTICS,
                        0,0,
                        &nicStatistics,sizeof(NIC_STATISTICS),
                        &dwBytesWritten,
                        0
                    )==TRUE)?true:false;
                    CloseHandle(ndisAccess);
                    // Only consider WIFI adapters
                    success=(success&&(nicStatistics.PhysicalMediaType
                                        ==NdisPhysicalMediumWirelessLan));
                    if (success) {
                        // Save the newly found device name and bail out
                        m_pDeviceName = unicodeName;
                        break;
                    }
                    else {
                        // Free the temporary string
                        delete[] unicodeName;
                    }
                }
            }
 
            // Iterate to next adapter
            pAdapter = pAdapter->Next;
        }
    }
    if (pAdapterInfo) free(pAdapterInfo);
}
 
bool CWifi::InitNewScan() {
    if (!m_pDeviceName) return false;
    HANDLE ndisAccess=CreateFile(
        TEXT("UIO1:"),
        0,0,0,
        OPEN_EXISTING,
        FILE_ATTRIBUTE_NORMAL|FILE_FLAG_OVERLAPPED,
        INVALID_HANDLE_VALUE
    );
    if (!ndisAccess || ndisAccess==INVALID_HANDLE_VALUE)
        return false;
 
    NDISUIO_SET_OID* query = (NDISUIO_SET_OID*) m_pIoCtlBuffer;
    query->Oid = OID_802_11_BSSID_LIST_SCAN;
    query->ptcDeviceName = m_pDeviceName;
 
    DWORD dwBytesWritten = 0;
    bool success=(DeviceIoControl(
        ndisAccess,
        IOCTL_NDISUIO_SET_OID_VALUE,
        query,m_ulBufferLength,
        query,m_ulBufferLength,
        &dwBytesWritten,
        0
    )==TRUE)?true:false;
    CloseHandle(ndisAccess);
    return success;
}
 
bool CWifi::Scan(IWifiListener& listener) {
    if (!Enumerate(listener)) return false;
    if (!InitScan()) return false;
    return true;
}
 
bool CWifi::Enumerate(IWifiListener& listener) {
    if (!m_pDeviceName) return false;
    HANDLE ndisAccess=CreateFile(
        TEXT("UIO1:"),
        0,0,0,
        OPEN_EXISTING,
        FILE_ATTRIBUTE_NORMAL|FILE_FLAG_OVERLAPPED,
        INVALID_HANDLE_VALUE
    );
    if (!ndisAccess || ndisAccess==INVALID_HANDLE_VALUE)
        return false;
 
    NDISUIO_QUERY_OID* query = (NDISUIO_QUERY_OID*) m_pIoCtlBuffer;
    query->Oid = OID_802_11_BSSID_LIST;
    query->ptcDeviceName = m_pDeviceName;
 
    DWORD dwBytesWritten = 0;
    bool success=(DeviceIoControl(
        ndisAccess,
        IOCTL_NDISUIO_QUERY_OID_VALUE,
        query,m_ulBufferLength,
        query,m_ulBufferLength,
        &dwBytesWritten,
        0
    )==TRUE)?true:false;
    CloseHandle(ndisAccess);
    if (success) {
        NDIS_802_11_BSSID_LIST* list;
        list = (NDIS_802_11_BSSID_LIST*)(&query->Data[0]);
        unsigned long count = list->NumberOfItems;
        NDIS_WLAN_BSSID* ap = &list->Bssid[0];
        unsigned long index=0;
        while(count--) {
            listener.OnBSSID(*ap);
            ap=(NDIS_WLAN_BSSID*)((unsigned char*)ap+ap->Length);
        }
    }
    return success;
}