/**
 * @file ConnectionController.cpp
 * @author Tomas Urban
 * @version 0.5
 * @date 23/07/2009
 */
#include "ConnectionController.h"
#include "PCAPLiveCapture.h"
#include "Logger/Logger.h"
#include "CaptureFactory.h"
#include "Messages/TrafficCaptureMessageFactory.h"
#include "Messages/CommonTrafficCaptureMessages.h"
#include <boost/filesystem.hpp>

ConnectionController::ConnectionController(SOCKET_TYPE hSocket) :
	m_hSocket(hSocket),
	m_pCapture(0)
{
}

ConnectionController::~ConnectionController(void)
{
	CaptureFactory::Instance().DisposeCaptureInstance(m_pCapture);
}

void ConnectionController::SendSimpleReply(CommonReplyMessage & reply, bool bRes)
{
	reply.SetResult(bRes);
	send(m_hSocket, reply.GetEncodedMessage(), reply.GetEncodedDataLength(), 0);
}

void ConnectionController::ProcessMessage(const struct TMessageHeader & header, const char * pPayload,
										  TcpipDispatch &td)
{
	std::string s = "Control message received (id: ";
	s += boost::lexical_cast<std::string>(header.nMsgType);
	s += ", len: ";
	s += boost::lexical_cast<std::string>(header.nMsgLen);
	s += ")";
	Logger::Instance().LogDebug(s);
	TrafficCaptureMessage * pMsg = TrafficCaptureMessageFactory::Instance().CreateMessage(
		header.nMsgType, pPayload, header.nMsgLen);
	if (!pMsg)
		return;
	const char * pszIgnoreWarn = "Warning: message ignored; the capturing device hasn't been created yet";
	switch (pMsg->GetId())
	{
	case OPEN_DEVICE_REQ:
		{
			Logger::Instance().LogInfo("OpenDeviceRequest message received");
			ECaptureInitResult res = ECaptureInit_Failed;
			if (m_pCapture == 0)
			{
				Logger::Instance().LogDebug("Creating capture instance...");
				OpenDeviceRequest * pReq = dynamic_cast<OpenDeviceRequest*>(pMsg);
				m_pCapture = CaptureFactory::Instance().CreateCaptureInstance(
					pReq->GetCaptureType(), pReq->GetCaptureMode());				
				if (m_pCapture != 0)
				{
					m_pCapture->InitDispatch(&td);
					m_pCapture->SetStartOffset(pReq->GetTimestamp());
					Logger::Instance().LogDebug("Initializing capture instance...");
					res = m_pCapture->OpenDevice(pReq->GetParameters());
					std::string sFile = pReq->GetCaptureFile();
					if (res != ECaptureInit_Failed && sFile.length() > 0)
					{
						if (!m_pCapture->SetCaptureFile(sFile))
							res = ECaptureInit_Failed;
					}
					if (res == ECaptureInit_Failed)
					{
						CaptureFactory::Instance().DisposeCaptureInstance(m_pCapture);
						m_pCapture = 0;
					}
				}				
			}
			else
				Logger::Instance().LogWarning("Capture device has been already created");
			OpenDeviceReply reply;
			reply.SetResult(res);
			send(m_hSocket, reply.GetEncodedMessage(), reply.GetEncodedDataLength(), 0);
			Logger::Instance().LogDebug("OpenDeviceReply message sent");
		}
		break;
	case SET_FILTER_REQ:
		Logger::Instance().LogInfo("SetFilterRequest message received");
		{
			bool bRes = (m_pCapture != 0);
			if (bRes)
			{
				SetFilterRequest * pReq = dynamic_cast<SetFilterRequest*>(pMsg);
				std::string sFilter = pReq->GetFilter();
				std::string s = "The received filter value is: ";
				s += sFilter;
				Logger::Instance().LogDebug(s);
				bRes = m_pCapture->SetFilter(sFilter);
			}
			else 
				Logger::Instance().LogWarning(pszIgnoreWarn);
			SetFilterReply reply;
			SendSimpleReply(reply, !bRes);
			Logger::Instance().LogDebug("SetFilterReply message sent");
		}
		break;
	case START_CAPTURE_REQ:
		{
			Logger::Instance().LogInfo("StartCaptureRequest message received");
			if (m_pCapture)
				m_pCapture->StartCapture();
			else 
				Logger::Instance().LogWarning(pszIgnoreWarn);
			StartCaptureReply reply;
			SendSimpleReply(reply, m_pCapture == 0);
		}
		Logger::Instance().LogDebug("StartCaptureReply message sent");
		break;
	case STOP_CAPTURE_REQ:
		{
			Logger::Instance().LogInfo("StopCaptureRequest message received");
			if (m_pCapture)
				m_pCapture->StopCapture();
			else 
				Logger::Instance().LogWarning(pszIgnoreWarn);
			StopCaptureReply reply;
			SendSimpleReply(reply, m_pCapture == 0);
			Logger::Instance().LogDebug("StopCaptureReply message sent");
		}
		break;
	case MERGE_FILES_REQ:
		{
			MergePcapFilesRequest * pReq = dynamic_cast<MergePcapFilesRequest*>(pMsg);
			Logger::Instance().LogInfo("MergePcapFilesRequest message received");
			MergePcapFilesReply reply;
			bool bRes;
			std::vector<std::string> vsFiles = pReq->GetFilesToMerge();
			if (bRes = vsFiles.size() > 0)
			{
				boost::filesystem::path targetFile = pReq->GetMergePath();
				targetFile /= "mergefile.pcap";				
				if (boost::filesystem::exists(targetFile))
					try
					{
						boost::filesystem::remove(targetFile);
						Logger::Instance().LogDebug(
							"Merge file from a previous session has been successfully removed");
					}
					catch (boost::filesystem::filesystem_error)					
					{
						Logger::Instance().LogError(
							"Merge file from a previous session could not be removed");
						bRes = false;
					}
				else
					Logger::Instance().LogDebug("No merge file from a previous session");
				if (bRes)
				{
					boost::filesystem::path mergecapPath = pReq->GetMergecapDirectory();
					mergecapPath /= "mergecap";
					std::string sCmd = "\"";
#ifdef WIN32
					sCmd += "\"";
#endif
					sCmd += mergecapPath.string();
					sCmd += "\" -w \"";
					sCmd += targetFile.string();
					sCmd += "\"";
					for (unsigned int i = 0; i < vsFiles.size(); ++i)
					{
						sCmd += " \"";
						sCmd += vsFiles[i];
						sCmd += "\"";
					}
#ifdef WIN32
					sCmd += "\"";
#endif
					Logger::Instance().LogDebug("Starting mergecap...");
					Logger::Instance().LogDebug(sCmd);
					system(sCmd.c_str());
					reply.SetMergeFile(targetFile.string());
					if (bRes = boost::filesystem::exists(targetFile))
						Logger::Instance().LogDebug("Files merged successfully");
					else
						Logger::Instance().LogError("File merging failed");
				}
					
			}
			SendSimpleReply(reply, bRes);
			Logger::Instance().LogDebug("MergePcapFilesReply message sent");
		}
		break;
	default:
		Logger::Instance().LogWarning("Warning: the received control message is not supported");
	}	
	TrafficCaptureMessageFactory::Instance().DisposeMessage(pMsg);
}

void ConnectionController::Run()
{
	TcpipDispatch td(m_hSocket);
	Logger::Instance().LogInfo("Waiting for control messages to inialize capturing...");

	// infinite receiving loop: interrupted when the connection is closed
	TMessageHeader header;
	int nReceived = 0, nLen, nRemaining;
	char * pBuf = NULL;
	while (-1)
	{
		if (nReceived < sizeof(header))
		{
			nRemaining = sizeof(header) - nReceived;
			if ((nLen = recv(m_hSocket, reinterpret_cast<char*>(&header) + nReceived, nRemaining, 0)) <= 0)
				break;
			nReceived += nLen;
		}
		else
		{
			if (header.nMsgLen > 0)
			{
				if (pBuf == NULL)
					pBuf = new char[header.nMsgLen];
				nRemaining = sizeof(header) + header.nMsgLen - nReceived;
				if ((nLen = recv(m_hSocket, pBuf + nReceived - sizeof(header), nRemaining , 0)) <= 0)
					break;
				nReceived += nLen;
			}
			if (nReceived == sizeof(header) + header.nMsgLen)
			{
				ProcessMessage(header, pBuf, td);
				delete[] pBuf;
				pBuf = NULL;
				nReceived = 0;
				if (!m_hSocket)
					break;
			}
		}
	}

	boost::mutex::scoped_lock cond(m_mutex);
	td.Stop();
	CaptureFactory::Instance().DisposeCaptureInstance(m_pCapture);
	m_pCapture = NULL;
	Logger::Instance().LogInfo("Client connection terminated");
	if (m_hSocket != SOCKET_CALL_ERROR)
	{
		CLOSE_SOCKET(m_hSocket);
		m_hSocket = SOCKET_CALL_ERROR;
	}
	m_captureClosed.notify_one();
	
}

void ConnectionController::Stop()
{
	boost::mutex::scoped_lock cond(m_mutex);
	if (m_hSocket != SOCKET_CALL_ERROR)
	{
		CLOSE_SOCKET(m_hSocket);
		m_hSocket = SOCKET_CALL_ERROR;
		m_captureClosed.wait(cond);		
	}
}