// SnifferDlg.cpp : implementation file
//

#include "stdafx.h"
#include "Sniffer.h"
#include "SnifferDlg.h"

#ifdef _DEBUG
#define new DEBUG_NEW
#endif


inline CString GetNiceString(LPCTSTR pString)
{
	const int numDots = 40 - _tcslen(pString);

	CString dotString(_T('.'),numDots);


	CString s = CString(pString) + dotString + CString(_T(" "));

	return s;
}


// CAboutDlg dialog used for App About

class CAboutDlg : public CDialog
{
public:
	CAboutDlg();

// Dialog Data
	enum { IDD = IDD_ABOUTBOX };

	protected:
	virtual void DoDataExchange(CDataExchange* pDX);    // DDX/DDV support

// Implementation
protected:
	DECLARE_MESSAGE_MAP()
};

CAboutDlg::CAboutDlg() : CDialog(CAboutDlg::IDD)
{
}

void CAboutDlg::DoDataExchange(CDataExchange* pDX)
{
	CDialog::DoDataExchange(pDX);
}

BEGIN_MESSAGE_MAP(CAboutDlg, CDialog)
END_MESSAGE_MAP()


// CSnifferDlg dialog



CSnifferDlg::CSnifferDlg(CWnd* pParent /*=NULL*/)
	: CDialog(CSnifferDlg::IDD, pParent)
	, m_SelInterface(_T(""))
{
	m_hIcon = AfxGetApp()->LoadIcon(IDR_MAINFRAME);

	Init();
}

void CSnifferDlg::DoDataExchange(CDataExchange* pDX)
{
	CDialog::DoDataExchange(pDX);
	DDX_Control(pDX, IDC_CMB_INTERFACES, m_CmbInterfaces);
	DDX_Control(pDX, IDC_FRM_INTERFACES, m_FrmInterfaces);
	DDX_Control(pDX, IDC_STARTSNIFFING, m_BtnStartSniffing);
	DDX_Control(pDX, IDC_TV_PACKETS, m_TVPackets);
}

BEGIN_MESSAGE_MAP(CSnifferDlg, CDialog)
	ON_WM_SYSCOMMAND()
	ON_WM_PAINT()
	ON_WM_QUERYDRAGICON()
	//}}AFX_MSG_MAP
	ON_WM_SIZE()
	ON_BN_CLICKED(IDC_STARTSNIFFING, OnBnClickedStartsniffing)
	ON_WM_CLOSE()
END_MESSAGE_MAP()


// CSnifferDlg message handlers

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

	// Add "About..." menu item to system menu.

	// IDM_ABOUTBOX must be in the system command range.
	ASSERT((IDM_ABOUTBOX & 0xFFF0) == IDM_ABOUTBOX);
	ASSERT(IDM_ABOUTBOX < 0xF000);

	CMenu* pSysMenu = GetSystemMenu(FALSE);
	if (pSysMenu != NULL)
	{
		CString strAboutMenu;
		strAboutMenu.LoadString(IDS_ABOUTBOX);
		if (!strAboutMenu.IsEmpty())
		{
			pSysMenu->AppendMenu(MF_SEPARATOR);
			pSysMenu->AppendMenu(MF_STRING, IDM_ABOUTBOX, strAboutMenu);
		}
	}

	// 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
	for ( int i = 0 ; i < m_Interfaces.GetSize(); i ++ )
	{
		m_CmbInterfaces.AddString( m_Interfaces.GetAt(i) );
	}

	m_CmbInterfaces.SetCurSel( 0 );

	m_IL.Create(16,16,ILC_COLOR32 | ILC_MASK, 0,1);

	m_IL.Add( LoadIcon(AfxGetInstanceHandle(), MAKEINTRESOURCE(IDI_COMP2COMP ) ) );
	m_IL.Add( LoadIcon(AfxGetInstanceHandle(), MAKEINTRESOURCE(IDI_IPHeader ) ) );
	m_IL.Add( LoadIcon(AfxGetInstanceHandle(), MAKEINTRESOURCE(IDI_TCPHEADER ) ) );

	m_TVPackets.SetImageList( &m_IL,TVSIL_NORMAL );

//	ShowWindow( SW_SHOWMAXIMIZED );
	ResizeForm();

	LOGFONT logFont;
	CFont* pFont = new CFont();

	RtlZeroMemory( &logFont, sizeof(logFont) );

	_tcscpy(logFont.lfFaceName, _T("Lucida Console"));
	logFont.lfHeight = 12;

	pFont->CreateFontIndirect( &logFont );

	m_TVPackets.SetFont( pFont, TRUE );
	



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

void CSnifferDlg::OnSysCommand(UINT nID, LPARAM lParam)
{
	if ((nID & 0xFFF0) == IDM_ABOUTBOX)
	{
		CAboutDlg dlgAbout;
		dlgAbout.DoModal();
	}
	else
	{
		CDialog::OnSysCommand(nID, lParam);
	}
}

// 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 CSnifferDlg::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 CSnifferDlg::OnQueryDragIcon()
{
	return static_cast<HCURSOR>(m_hIcon);
}

// Starts the sniffing process
bool CSnifferDlg::StartSniffing(void)
{
	SOCKET        s = INVALID_SOCKET;
	WSABUF        wbuf = {0};
	DWORD         dwBytesRet = 0, dwFlags = 0;
	unsigned int  optval = 0;
	char         *rcvbuf=NULL;
	int           rc = 0, err;


	
	if ( m_ulFilterMask & (FILTER_MASK_SOURCE_ADDRESS | FILTER_MASK_SOURCE_PORT) )
	{
		printf("Source address filter     : ");
		//PrintAddress((SOCKADDR *)&m_saSourceAddress, sizeof(m_saSourceAddress));
		printf("\n");
	}
	if ( m_ulFilterMask & (FILTER_MASK_DESTINATION_ADDRESS | FILTER_MASK_DESTINATION_PORT) )
	{
		printf("Destination address filter: ");
		//PrintAddress((SOCKADDR *)&g_saDestinationAddress, sizeof(g_saDestinationAddress));
		printf("\n");
	}

	//
	// Create a raw socket for receiving IP datagrams
	//
	s = WSASocket(AF_INET, SOCK_RAW, IPPROTO_IP, NULL, 0, WSA_FLAG_OVERLAPPED);
	if (s == INVALID_SOCKET)
	{
		printf("WSASocket() failed: %d\n", WSAGetLastError());
		return false;
	}

	//
	// This socket MUST be bound before calling the ioctl
	//

	sockaddr_in sa;
	sa.sin_family = AF_INET;
	sa.sin_addr.s_addr = inet_addr(m_SelInterface);
	sa.sin_port = INADDR_ANY;//htons(7000);

	rc = bind(s, (SOCKADDR *)&sa, sizeof(sa));
	if (rc == SOCKET_ERROR)
	{
		printf("bind() failed: %d\n", WSAGetLastError());
		if (INVALID_SOCKET != s)
		{
			closesocket(s);
			s = INVALID_SOCKET;
		}
		WSACleanup();
		return false;
	}
	printf("Binding to: ");
	//PrintAddress((SOCKADDR *)&g_saLocalInterface, sizeof(g_saLocalInterface));
	printf("\n");

	//
	// Set the SIO_RCVALLxxx ioctl
	//
	optval = 1;
	rc = WSAIoctl(s, SIO_RCVALL, &optval, sizeof(optval),
		NULL, 0, &dwBytesRet, NULL, NULL);
	if (rc == SOCKET_ERROR)
	{
		printf("WSAIotcl(0x%x) failed: %d\n", SIO_RCVALL,
			(err = WSAGetLastError()));
		if (err == WSAEINVAL)
		{
			printf("NOTE: IPv6 does not currently support the SIO_RCVALL* ioctls\n");
		}

		if (INVALID_SOCKET != s)
		{
			closesocket(s);
			s = INVALID_SOCKET;
		}
		WSACleanup();
		return false;
	}

	//
	// Allocate a buffer for receiving data
	//
	rcvbuf = (char *)HeapAlloc(GetProcessHeap(), 0, MAX_IP_SIZE);
	if (rcvbuf == NULL)
	{
		fprintf(stderr, "HeapAlloc failed: %d\n", GetLastError());
		if (INVALID_SOCKET != s)
		{
			closesocket(s);
			s = INVALID_SOCKET;
		}
		WSACleanup();
		return false;
	}

	//
	// Start receiving IP datagrams until interrupted
	// 
	while ( !m_bExit )
	{
		wbuf.len = MAX_IP_SIZE;
		wbuf.buf = rcvbuf;
		dwFlags  = 0;

		rc = WSARecv(s, &wbuf, 1, &dwBytesRet, &dwFlags, NULL, NULL);
		if (rc == SOCKET_ERROR)
		{
			printf("WSARecv() failed: %d\n", WSAGetLastError());
			break;
		}

		DecodePacket(rcvbuf , dwBytesRet );


	}
	//
	// Cleanup
	//

	if (rcvbuf) 
		HeapFree(GetProcessHeap(), 0, rcvbuf);

	if (INVALID_SOCKET != s)
	{
		closesocket(s);
		s = INVALID_SOCKET;
	}
	WSACleanup();

	return true;
}

bool CSnifferDlg::Init(void)
{
	try
	{
		m_bExit = false;
		m_hThread = NULL;
		m_dwThreadID = 0;

		m_ulFilterMask = 0;

		m_Logger.Open( _T("Logger.txt"),CFile::modeCreate | CFile::modeWrite | CFile::shareDenyNone);

		int           rc = 0;
		WSADATA       wsd;
		//
		// Load Winsock
		//
		if ((rc = WSAStartup(MAKEWORD(2,2), &wsd)) != 0)
		{
			printf("WSAStartup() failed: %d\n", rc);
			return false;
		}

		GetInterfaces( m_Interfaces );


		return true;
	}
	catch(...)
	{
		return false;
	}

}

// Decodes the packet
bool CSnifferDlg::DecodePacket(char* pData, DWORD dwSize)
{

	try
	{
	
		if (dwSize < sizeof(char))
			return false;
		
		// Check the IP version
		const int ip_version = HI_BYTE(*pData);
		
		CString source, destination;

		if ( ip_version == 4)
		{
			IPV4_HDR* pHeader;
			int headerLen = 0;

		
			// Verify the buffer is large enough
			if (dwSize  < sizeof(IPV4_HDR))
				return false;
			
			// Get length of IPv4 header to determine where next protocol header begins
			headerLen = LO_BYTE(*pData) * 4;
			
			pHeader = (IPV4_HDR *)pData;

			ConvertAddress( pHeader->ip_srcaddr, source );
			ConvertAddress( pHeader->ip_destaddr, destination);

			switch (pHeader->ip_protocol)
			{