/**
 * @file PCAPCapture.cpp
 * @author Tomas Urban
 * @version 0.4
 * @date 23/07/2009
 */
#include "PCAPCapture.h"
#include <boost/thread/thread.hpp>
#include <boost/tokenizer.hpp>
#include "Logger/Logger.h"

PCAPCapture::PCAPCapture() :
	m_fp(NULL),
	m_nHandleCount(0),
	m_bRunning(false),
	m_dumpFiles(0),
	m_threadRunning(0)
{
}

PCAPCapture::~PCAPCapture() 
{
	CloseDevice();	
}

pcap_t * PCAPCapture::OpenPcapSource(const std::string sSource)
{
	char errbuf[PCAP_ERRBUF_SIZE];
	return pcap_open(sSource.c_str(), SNAP_LEN, PCAP_OPENFLAG_PROMISCUOUS,
		1000, NULL, errbuf);
}

void PCAPCapture::CloseDevice()
{
	if (m_fp != NULL)
	{
		StopCapture();
		CloseCaptureFile();
		for (int i = 0; i < m_nHandleCount; i++)
			pcap_close(m_fp[i]);
		delete m_fp;
		m_fp = NULL;
		m_threadRunning = NULL;
		m_nHandleCount = 0;
	}
}
void PCAPCapture::SetPcapHandles(pcap_t ** srcBuf, int nLen)
{
	CloseDevice();
	m_fp = new pcap_t*[nLen];
	m_nHandleCount = nLen;
	memcpy(m_fp, srcBuf, nLen * sizeof(pcap_t*));
	m_threadRunning = new bool[nLen];
	memset(m_threadRunning, 0, nLen * sizeof(bool));
}

void PCAPCapture::DispatcherHandler(u_char *temp1, 
                        const struct pcap_pkthdr *header, const u_char *pkt_data)
{	
	PCAPCapture::CapturingThread * pThread = reinterpret_cast<PCAPCapture::CapturingThread*>(temp1);
	PCAPCapture * pObj = pThread->m_pCapture;
	int nIndex = pThread->m_nHandleIndex;
	if (pObj->m_dumpFiles)
	{
		boost::mutex::scoped_lock lock(pObj->m_mutex);
		pcap_dump(reinterpret_cast <u_char*>(pObj->m_dumpFiles[nIndex]), header, pkt_data);
	}
	CapturedData cd;
	cd.SetData(header->caplen, reinterpret_cast <const char*> (pkt_data));
	cd.SetTimestamp(header->ts);
	if (pObj->m_fp)
		pObj->ProcessCapturedData(&cd);
}

PCAPCapture::CapturingThread::CapturingThread(PCAPCapture * pCapture, int nHandleIndex) :
	m_pCapture(pCapture),
	m_nHandleIndex(nHandleIndex)
{
}

void PCAPCapture::CapturingThread::operator()()
{
	std::string s = "Packet capture started";
	std::string s2 = "";
	if (m_pCapture->m_nHandleCount > 1)
	{
		s2 = " (source #";
		s2 += boost::lexical_cast<std::string>(m_nHandleIndex + 1);
		s2 += ")";
		s += s2;
	}
	Logger::Instance().LogDebug(s);

	pcap_loop(m_pCapture->m_fp[m_nHandleIndex], 0, PCAPCapture::DispatcherHandler, 
		reinterpret_cast <u_char*> (this));

	s = "Packet capture ended";
	s += s2;
	Logger::Instance().LogDebug(s);
	{
		boost::mutex::scoped_lock(m_pCapture->m_mutex);
		m_pCapture->m_threadRunning[m_nHandleIndex] = false;
		if (m_pCapture->m_bRunning)
			m_pCapture->m_threadClosed.notify_one();
	}
}

bool PCAPCapture::StartCapture()
{
	if (m_bRunning)
		return true;
	m_bRunning = true;
	if (m_fp)		
	{
		for (int i = 0; i < m_nHandleCount; i++)
		{
			m_threadRunning[i] = true;
			boost::thread thread(CapturingThread(this, i));
		}
	}
	return true;
}

bool PCAPCapture::StopCapture()
{	
	if (!m_bRunning)
		return true;
	m_bRunning = false;
	if (m_fp)
	{
		for (int i = 0; i < m_nHandleCount; i++)
		{
			boost::mutex::scoped_lock cond(m_mutex);
			if (m_threadRunning[i])
			{
				pcap_breakloop(m_fp[i]);
				m_threadClosed.wait(cond);
			}
		}
	}
	return true;
}

bool PCAPCapture::SetFilter(const std::string sFilter)
{
	if (!m_fp)
	{
		Logger::Instance().LogWarning("Cannot apply filter; the capture device is not open");
		return false;
	}
	
	std::string sFormattedFilter = sFilter;

	// compile the filter		
	struct bpf_program fcode;
	std::string s = "Setting PCAP filter (";
	s += sFormattedFilter;
	s += ")...";
	Logger::Instance().LogDebug(s);
	bool bRes = true;
	for (int i = 0; i < m_nHandleCount; i++)
	{		
		if (m_nHandleCount > 1)
		{
			s = "Applying filter to source #";
			s += boost::lexical_cast<std::string>(i + 1);
			s += "...";
			Logger::Instance().LogDebug(s);
		}				
		if (pcap_compile(m_fp[i], &fcode, sFormattedFilter.c_str(), 1, 0) < 0)
		{
			Logger::Instance().LogError("Error compiling packet filter");
			bRes = false;
			continue;
		}
	    
		//set the filter
		int nRes = pcap_setfilter(m_fp[i], &fcode);
		pcap_freecode(&fcode);
		if (nRes < 0)
		{
			Logger::Instance().LogError("Error setting packet filter");
			bRes = false;
			continue;
		}		
		Logger::Instance().LogInfo("Filter applied successfully");
	}
	Logger::Instance().LogDebug("PCAP filter ready");
	return bRes;
}

void PCAPCapture::CloseCaptureFile()
{
	if (m_dumpFiles)
	{
		boost::mutex::scoped_lock lock(m_mutex);
		for (int i = 0; i < m_nHandleCount; i++)
			pcap_dump_close(m_dumpFiles[i]);
		delete m_dumpFiles;
		m_dumpFiles = 0;
	}
}

bool PCAPCapture::InitCaptureFile(const std::string sFile)
{
	if (!m_nHandleCount)
		return false;
	CloseCaptureFile();
	bool bRes = true;
	if (sFile.length() > 0)
	{
		boost::mutex::scoped_lock lock(m_mutex);
		m_dumpFiles = new pcap_dumper_t * [m_nHandleCount];
		memset(m_dumpFiles, 0, sizeof(pcap_dumper_t*) * m_nHandleCount);

		std::string sName = sFile;
		std::string sExt = "";
		if (m_nHandleCount > 1)
		{
			size_t nSize = sFile.find_last_of('.');
			if (nSize < sFile.length())
			{
				sName = sFile.substr(0, nSize);
				sExt = sFile.substr(nSize);
			}
		}
		
		for (int i = 0; i < m_nHandleCount; i++)
		{
			std::string sTmp = sFile;
			if (m_nHandleCount > 1)
			{
				sTmp = sName;
				sTmp += boost::lexical_cast<std::string>(i + 1);
				sTmp += sExt;
			}
			m_dumpFiles[i] = pcap_dump_open(m_fp[i], sTmp.c_str());
			if (!m_dumpFiles[i])
			{
				std::string s = "Failed to create dump file \"";
				s += sTmp;
				s += "\"";
				Logger::Instance().LogError(s);
				return false;
			}
		}
		Logger::Instance().LogDebug("PCAP dump file successfully initialized");
	}
	else
		Logger::Instance().LogDebug("Dump file disabled");
	return true;;
}