// WiFiSnifferDlg.cpp : implementation file
//

#include "stdafx.h"
#include "WiFiSniffer.h"
#include "WiFiSnifferDlg.h"

#ifdef _DEBUG
#define new DEBUG_NEW
#endif


// CWiFiSnifferDlg dialog




CWiFiSnifferDlg::CWiFiSnifferDlg(CWnd* pParent /*=NULL*/)
	: CDialog(CWiFiSnifferDlg::IDD, pParent)
{
	m_hIcon = AfxGetApp()->LoadIcon(IDR_MAINFRAME);
}

void CWiFiSnifferDlg::DoDataExchange(CDataExchange* pDX)
{
	CDialog::DoDataExchange(pDX);
	DDX_Control(pDX, IDC_LIST_IFACES, lvIfaces);
	DDX_Control(pDX, IDC_STATIC_MAC, laMac);
	DDX_Control(pDX, IDC_STATIC_CHANNEL_NUM, laChannel);
	DDX_Control(pDX, IDC_STATIC_PHY, laPhy);
	DDX_Control(pDX, IDC_COMBO_PHY, cbPhy);
	DDX_Control(pDX, IDC_EDIT_CHANNEL, edChannel);
	DDX_Control(pDX, IDC_LIST_PACKES, lbRaw);
	DDX_Control(pDX, IDC_CHECK_DO_NOT_CHANGE_MODE, cbDoNotChangeMode);
	DDX_Control(pDX, IDC_STATIC_IFACE_MODE, laIfaceMode);
}

BEGIN_MESSAGE_MAP(CWiFiSnifferDlg, CDialog)
	ON_WM_PAINT()
	ON_WM_QUERYDRAGICON()
	//}}AFX_MSG_MAP
	ON_WM_DESTROY()
	ON_BN_CLICKED(IDC_BUTTON_ENUM_IFACES, &CWiFiSnifferDlg::OnBnClickedButtonEnumIfaces)
	ON_BN_CLICKED(IDC_BUTTON_START, &CWiFiSnifferDlg::OnBnClickedButtonStart)
	ON_BN_CLICKED(IDC_BUTTON_STOP, &CWiFiSnifferDlg::OnBnClickedButtonStop)
	ON_BN_CLICKED(IDC_BUTTON_SET_PHY, &CWiFiSnifferDlg::OnBnClickedButtonSetPhy)
	ON_BN_CLICKED(IDC_BUTTON_CLEAR, &CWiFiSnifferDlg::OnBnClickedButtonClear)
END_MESSAGE_MAP()


// CWiFiSnifferDlg message handlers

BOOL CWiFiSnifferDlg::OnInitDialog()
{
	CDialog::OnInitDialog();

	// Set the icon for this dialog.  The framework does this automatically
	//  when the application's main window is not a dialog
	SetIcon(m_hIcon, TRUE);			// Set big icon
	SetIcon(m_hIcon, FALSE);		// Set small icon

	// TODO: Add extra initialization here

	lvIfaces.InsertColumn(0, _T("Friendly name"), 0, 300);
	lvIfaces.InsertColumn(1, _T("Id"), 0, 300);
	lvIfaces.SetExtendedStyle(lvIfaces.GetExtendedStyle() | LVS_EX_FULLROWSELECT);

	cbPhy.AddString(_T("802.11a"));
	cbPhy.AddString(_T("802.11b"));
	cbPhy.AddString(_T("802.11g"));
	cbPhy.AddString(_T("802.11n"));
	cbPhy.SetCurSel(0);

	edChannel.SetWindowText(_T("6"));

	__hook(&CwclWiFiSniffer::AfterOpen, &Sniffer, &CWiFiSnifferDlg::AfterOpen);
	__hook(&CwclWiFiSniffer::BeforeClose, &Sniffer, &CWiFiSnifferDlg::BeforeClose);
	__hook(&CwclWiFiSniffer::OnRawFrameReceived, &Sniffer, &CWiFiSnifferDlg::OnRawFrameReceived);

	return TRUE;  // return TRUE  unless you set the focus to a control
}

// If you add a minimize button to your dialog, you will need the code below
//  to draw the icon.  For MFC applications using the document/view model,
//  this is automatically done for you by the framework.

void CWiFiSnifferDlg::OnPaint()
{
	if (IsIconic())
	{
		CPaintDC dc(this); // device context for painting

		SendMessage(WM_ICONERASEBKGND, reinterpret_cast<WPARAM>(dc.GetSafeHdc()), 0);

		// Center icon in client rectangle
		int cxIcon = GetSystemMetrics(SM_CXICON);
		int cyIcon = GetSystemMetrics(SM_CYICON);
		CRect rect;
		GetClientRect(&rect);
		int x = (rect.Width() - cxIcon + 1) / 2;
		int y = (rect.Height() - cyIcon + 1) / 2;

		// Draw the icon
		dc.DrawIcon(x, y, m_hIcon);
	}
	else
	{
		CDialog::OnPaint();
	}
}

// The system calls this function to obtain the cursor to display while the user drags
//  the minimized window.
HCURSOR CWiFiSnifferDlg::OnQueryDragIcon()
{
	return static_cast<HCURSOR>(m_hIcon);
}


void CWiFiSnifferDlg::OnDestroy()
{
	CDialog::OnDestroy();

	// TODO: Add your message handler code here

	Sniffer.Close();
	__unhook(&Sniffer);
}

void CWiFiSnifferDlg::RefreshChannel()
{
	unsigned long Ch;
	int Res = Sniffer.GetChannel(Ch);
	CString s;
	if (Res == WCL_E_SUCCESS)
	{
		s.Format(_T("%d"), Ch);
		laChannel.SetWindowText(s);
	}
	else
	{
		s.Format(_T("Error: 0x%.8X"), Res);
		laChannel.SetWindowText(s);
	}
}

void CWiFiSnifferDlg::RefreshPhy()
{
	wclWiFiSnifferPhy Phy;
	int Res = Sniffer.GetPhy(Phy);
	if (Res == WCL_E_SUCCESS)
	{
		switch (Phy)
		{
			case ph802_11a:
				laPhy.SetWindowText(_T("802.11a"));
				break;
			case ph802_11b:
				laPhy.SetWindowText(_T("802.11b"));
				break;
			case ph802_11g:
				laPhy.SetWindowText(_T("802.11g"));
				break;
			case ph802_11n:
				laPhy.SetWindowText(_T("802.11n"));
				break;
			default:
				laPhy.SetWindowText(_T("UNKNOWN"));
				break;
		}
	}
	else
	{
		CString s;
		s.Format(_T("Error: 0x%.8X"), Res);
		laPhy.SetWindowText(s);
	}
}

void CWiFiSnifferDlg::ShowError(CString s, const int Error)
{
#pragma warning(disable:4840)
	s.Format(_T("%s: 0x%.8X"), s, Error);
#pragma warning(default:4840)
	AfxMessageBox(s);
}

CString GuidToString(const GUID& Id)
{
	CString res;
	
	LPOLESTR lpszClsID;
	StringFromCLSID(Id, &lpszClsID);
	res = CString(lpszClsID);
	CoTaskMemFree(lpszClsID);
	return res;
}

void CWiFiSnifferDlg::OnBnClickedButtonEnumIfaces()
{
	lvIfaces.DeleteAllItems();
	
	int Res = Client.Open();
	if (Res != WCL_E_SUCCESS)
		ShowError(_T("Unable to open WiFi Client"), Res);
	else
	{
		wclWiFiInterfaces Ifaces;
		Res = Client.EnumInterfaces(Ifaces);
		if (Res != WCL_E_SUCCESS)
			ShowError(_T("Unable to enumerate WiFi Interfaces"), Res);
		else
		{
			if (Ifaces.size() == 0)
				AfxMessageBox(_T("No WiFi Interfaces were found"));
			else
			{
				for (size_t i = 0; i < Ifaces.size(); i++)
				{
					lvIfaces.InsertItem((int)i, Ifaces[i].Description.c_str());
					lvIfaces.SetItemText((int)i, 1, GuidToString(Ifaces[i].Id));
				}
			}
		}
		Client.Close();
	}
}

GUID StringToGuid(const CString& s)
{
	GUID Guid;
	CLSIDFromString((LPOLESTR)const_cast<LPTSTR>((LPCTSTR)s), &Guid);
	return Guid;
}

void CWiFiSnifferDlg::OnBnClickedButtonStart()
{
	POSITION Pos = lvIfaces.GetFirstSelectedItemPosition();
	if (Pos == NULL)
		AfxMessageBox(_T("Select interface"));
	else
	{
		GUID Id = StringToGuid(lvIfaces.GetItemText(lvIfaces.GetNextSelectedItem(Pos), 1));
		try
		{
			Sniffer.DoNotChangeMode = (cbDoNotChangeMode.GetCheck() == TRUE);
			int Res = Sniffer.Open(Id);
			if (Res != WCL_E_SUCCESS)
				ShowError(_T("Unable start capturing"), Res);
		}
		catch(wclException* e)
		{
			MessageBoxA(NULL, e->what(), "ERROR", 0);
		}
	}
}

void CWiFiSnifferDlg::OnBnClickedButtonStop()
{
	int Res = Sniffer.Close();
	if (Res != WCL_E_SUCCESS)
		ShowError(_T("Unable to stop capturing"), Res);
}

void CWiFiSnifferDlg::OnBnClickedButtonSetPhy()
{
	wclWiFiSnifferPhy Phy = (wclWiFiSnifferPhy)cbPhy.GetCurSel();
	int Res = Sniffer.SetPhy(Phy);
	if (Res != WCL_E_SUCCESS)
		ShowError(_T("Unable to change PHY"), Res);
	else
		RefreshPhy();
}

void CWiFiSnifferDlg::AfterOpen(void* Sender)
{
	UNREFERENCED_PARAMETER(Sender);

	RefreshChannel();
	RefreshPhy();
	__int64 Mac;
	int Res = Sniffer.GetMacAddr(Mac);
	if (Res == WCL_E_SUCCESS)
	{
		CString s;
		s.Format(_T("%.4X%.8X"), static_cast<INT32>((Mac >> 32) & 0x00000FFFF),
			static_cast<INT32>(Mac) & 0xFFFFFFFF);
		laMac.SetWindowText(s);
		switch (Sniffer.Mode)
		{
			case omUnknown:
				laIfaceMode.SetWindowText(_T("Unknown"));
				break;
			case omStation:
				laIfaceMode.SetWindowText(_T("Station"));
				break;
			case omAccessPoint:
				laIfaceMode.SetWindowText(_T("Access point"));
				break;
			case omExtensibleStation:
				laIfaceMode.SetWindowText(_T("Extensible Station"));
				break;
			case omExtensibleAccessPoint:
				laIfaceMode.SetWindowText(_T("Extensible Access Point"));
				break;
			case omWiFiDirectDevice:
				laIfaceMode.SetWindowText(_T("WiFiDirect Device"));
				break;
			case omWiFiDirectGroupOwner:
				laIfaceMode.SetWindowText(_T("WiFiDirect Group Owner"));
				break;
			case omWiFiDirectClient:
				laIfaceMode.SetWindowText(_T("WiFiDirect Client"));
				break;
			case omManufacturing:
				laIfaceMode.SetWindowText(_T("Manufacturing"));
				break;
			case omNetworkMonitor:
				laIfaceMode.SetWindowText(_T("Network Monitor"));
				break;
			default:
				laIfaceMode.SetWindowText(_T("Unknown"));
				break;
		}
	}
	else
	{
		CString s;
		s.Format(_T("Error: 0x%.8X"), Res);
		laMac.SetWindowText(s);
	}
}

void CWiFiSnifferDlg::BeforeClose(void* Sender)
{
	UNREFERENCED_PARAMETER(Sender);

	laMac.SetWindowText(_T("000000000000"));
	laPhy.SetWindowText(_T("NONE"));
	laChannel.SetWindowText(_T("0"));
	laIfaceMode.SetWindowText(_T("Unknown"));
}
void CWiFiSnifferDlg::OnBnClickedButtonClear()
{
	lbRaw.ResetContent();
}

void CWiFiSnifferDlg::OnRawFrameReceived(void* Sender, const void* const Buffer, 
	const unsigned long Size)
{
	UNREFERENCED_PARAMETER(Sender);

	CString s = _T("");
	for (unsigned long i = 0; i < Size; i++)
	{
		CString a;
		a.Format(_T("%.2X"), ((PBYTE)Buffer)[i]);
		s = s + a + _T(" ");
		if (s.GetLength() == 96)
		{
			lbRaw.AddString(s);
			s = _T("");
		}
	}
	if (s.GetLength() > 0)
		lbRaw.AddString(s);
}