//------------------------------------------------------------------------------
// <copyright file="VirtualRelayCryptoProcessor.cpp" company="LiquidRobotics">
//
// Copyright (c) 2014 Liquid Robotics, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.
//
// This software is associated with a release.  Liquid Robotics makes no
// representation that the design, interfaces or component distribution will
// remain unchanged in future releases.
//
// </copyright>
//
// <summary>
// 	Implementation of VirtualRelayCryptoProcessor class, which manages the
//  encryption and decryption of incoming and outgoing messages
// </summary>
//
// <owner>Arash Ushani</owner>
//------------------------------------------------------------------------------

#include "VirtualRelayCryptoProcessor.h"
#include "VirtualRelayMessage.h"
#include "CryptoProcessor.h"
#include "LqrSMACLogger.h"
#include "FCCMessage.h"
#include "nsextend.h"

using namespace std;
using namespace boost;
using namespace LiquidRobotics::CryptoLib;
using namespace LiquidRobotics::PayloadInterface;

namespace LiquidRobotics {
namespace VirtualRelay {

VirtualRelayCryptoProcessor::VirtualRelayCryptoProcessor(Logger&  _logger)
: logger(_logger)
{
	connectionSecure = false;
	pendingSecurityChange = false;
	pendingConnectionSecurityState = false;
	securityAckWGMSID = 0xFFFFFFFF;

	activeCrypto.reset(new CryptoProcessor());
	pendingCrypto.reset(new CryptoProcessor());

}

VirtualRelayCryptoProcessor::~VirtualRelayCryptoProcessor()
{
	Logger::Locked(logger)->Debug << "Crypto Processor Destructor" << endl;
}

VirtualRelayCryptoProcessor::UInt8Vector VirtualRelayCryptoProcessor::IncomingLevel2Message(UInt8Vector& packet)
{
	incomingMutex.lock();

	UInt8Vector result;

	if (connectionSecure)
	{
		result = DecryptPacket(packet);
	}
	else
	{
		result = packet;
	}

	// Is this command something special that we have to deal with,
	// like something that will change the security state of the channel?
	boost::uint16_t command = result[8] + (result[9]<<8);
	if (command == FCCMessage::CMD_SECURE_ON)
	{
		// Change to Secure On 2 (type 0xA51E)

		result[8] = 0x1E;
		// commandPacket[9] = 0xA5; // Don't need this because it already is 0xA5

		result.pop_back();
		result.pop_back();
		VirtualRelayMessage::AppendCRC(result);
	}
	if (command == FCCMessage::CMD_SECURE_OFF)
	{
		pendingSecurityChange = true;
		pendingConnectionSecurityState = false;

		// Store WGMS ID so we can track the message later
		securityAckWGMSID = (result[4]<<24) + (result[3]<<16) + (result[2]<<8) + result[1];
	}

	incomingMutex.unlock();
	return result;
}

VirtualRelayCryptoProcessor::UInt8Vector VirtualRelayCryptoProcessor::OutgoingLevel2Message(UInt8Vector& packet)
{
	outgoingMutex.lock();

	UInt8Vector result;

	// Will this message change the security state of this channel?
	bool toggleSecurityState = false;
	if (pendingSecurityChange)
	{
		unsigned int wgmsID = (packet[4]<<24) + (packet[3]<<16) + (packet[2]<<8) + packet[1];
		if (securityAckWGMSID == wgmsID)
			toggleSecurityState = true;
	}

	// Messages that toggle security state (ie, secure authenticate ack, secure off ack) are always sent in clear text

	if (connectionSecure && !toggleSecurityState)
	{
		result = EncryptPacket(packet);
	}
	else
	{
		result = packet;
	}

	if (toggleSecurityState)
	{
		connectionSecure = pendingConnectionSecurityState;
		activeCrypto.swap(pendingCrypto);
		Logger::Locked(logger)->Info << "Toggling Security State. Channel is now " << (connectionSecure ? "SECURE":"NOT SECURE") << endl;

		pendingSecurityChange = false;
	}

	outgoingMutex.unlock();
	return result;
}

VirtualRelayCryptoProcessor::UInt8Vector VirtualRelayCryptoProcessor::VerifySecureAuthenticate(UInt8Vector& packet)
{
	// Extract Client Channel ID, IV, and Pass Phrase (Cipher Text)

	// TODO: This chunk of code should probably be redone using a C-struct
	int headerSize = 10;
	int channelIDSize = 32;
	int ivSizeSize = 1;
	int cipherTextSizeSize = 2;

	string WGMSChannelID(packet.begin()+headerSize, packet.begin()+headerSize+channelIDSize);
	boost::uint8_t ivSize = packet[headerSize+channelIDSize];
	boost::uint16_t cipherTextSize = packet[headerSize+channelIDSize+ivSizeSize] + (packet[headerSize+channelIDSize+ivSizeSize+1]<<8);

	unsigned char* iv = &packet[headerSize+channelIDSize+cipherTextSizeSize+ivSizeSize];
	unsigned char* cipherText = &packet[headerSize+channelIDSize+cipherTextSizeSize+ivSizeSize+ivSize];

	int result_len=0;
	unsigned char decryptedCipherText[cipherTextSize+8];
	pendingCrypto->Decrypt(cipherText, cipherTextSize, iv, ivSize, decryptedCipherText, &result_len);
	string WGMSPassPhrase((const char *) decryptedCipherText, result_len);

	VirtualRelayMessage::VR_ACK_BASIC secureAuthResponse;
	secureAuthResponse.commandID = (packet[4]<<24) + (packet[3]<<16) + (packet[2]<<8) + packet[1];
	secureAuthResponse.responseSize = 0;

	// See if these match what we got from the Float during the Secure On Command
	if (WGMSPassPhrase.compare(PassPhrase)==0 && WGMSChannelID.compare(ChannelID)==0)
	{
		Logger::Locked(logger)->Info << "Secure Authenticate Matches!" << endl;

		pendingSecurityChange = true;
		pendingConnectionSecurityState = true;
		secureAuthResponse.type = FCCMessage::ACK_BASIC;
		securityAckWGMSID = secureAuthResponse.commandID;	// Store the commandID so we can track the secure auth ack message
	}
	else
	{
		Logger::Locked(logger)->Alert << "Secure Authenticate did not match." << endl;

		Logger::Locked(logger)->Alert << "WGMS Pass Phrase: " << WGMSPassPhrase << endl;
		Logger::Locked(logger)->Alert << "Stored Pass Phrase: " << PassPhrase << endl;

		Logger::Locked(logger)->Alert << "WGMS Channel ID: " << WGMSChannelID << endl;
		Logger::Locked(logger)->Alert << "Stored Channel ID: " << ChannelID << endl;

		pendingSecurityChange = false;
		secureAuthResponse.type = FCCMessage::NAK_SECURITY;
	}

	UInt8Vector outgoingMessage;
	outgoingMessage.insert(outgoingMessage.end(), (char*)&secureAuthResponse, ((char*)&secureAuthResponse) + sizeof(VirtualRelayMessage::VR_ACK_BASIC));
	VirtualRelayMessage::AppendCRC(outgoingMessage);

	return outgoingMessage;
}

void VirtualRelayCryptoProcessor::ProcessSecureOnAck2(UInt8Vector& packet)
{
	boost::uint16_t permPublicKeySize = packet[7] + (packet[8]<<8);
	boost::uint16_t ephmPublicKeySize = packet[9] + (packet[10]<<8);
	boost::uint16_t  signatureSize = packet[11] + (packet[12]<<8);

	int aesKeyIndex = 13 + permPublicKeySize + ephmPublicKeySize + signatureSize;

	// Extract AES Key, Pass Phrase, Channel ID
	int aesKeySize = 32;
	int passPhraseSize = 20;
	int channelIDSize = 32;

	unsigned char* aesKey = &(packet[aesKeyIndex]);
	pendingCrypto->SetKey(aesKey, aesKeySize);

	PassPhrase = string(packet.begin()+aesKeyIndex+aesKeySize, packet.begin()+aesKeyIndex+aesKeySize+passPhraseSize);
	ChannelID = string(packet.begin()+aesKeyIndex+aesKeySize+passPhraseSize, packet.begin()+aesKeyIndex+aesKeySize+passPhraseSize+channelIDSize);

	// Now format the message to be sent back to shore

	packet[0] = 0xA5;

	int oldSize = packet[5] + (packet[6]<<8);
	int newSize = oldSize - aesKeySize - passPhraseSize - channelIDSize;
	packet[5] = newSize & 0xFF;
	packet[6] = (newSize & 0xFF00)>>8;

	packet.erase(packet.begin()+aesKeyIndex, packet.end());

	VirtualRelayMessage::AppendCRC(packet);
}

bool VirtualRelayCryptoProcessor::IsEncrypted()
{
	return connectionSecure;
}

VirtualRelayCryptoProcessor::UInt8Vector VirtualRelayCryptoProcessor::EncryptPacket(UInt8Vector& clearTextPacket)
{
	UInt8Vector encryptedMsg;
	unsigned char* clearTextData = &clearTextPacket[0];

	int ivSize = 16;
	unsigned char theIV[ivSize];
	CryptoProcessor::GenerateRandomIV(ivSize, theIV);

	encryptedMsg.insert(encryptedMsg.end(), theIV, theIV + ivSize);

	int result_len = 0;
	unsigned char cipherData[clearTextPacket.size()+8];
	memset(cipherData, 0, clearTextPacket.size()+8);
	activeCrypto->Encrypt(clearTextData, clearTextPacket.size(), theIV, ivSize, cipherData, &result_len);

	VirtualRelayMessage::VR_SECURE_PACKET_HEADER securePacketHeader;
	securePacketHeader.type = VirtualRelayMessage::EncryptedShiptoShore;
	securePacketHeader.sizeIV = ivSize;
	securePacketHeader.sizeCipher = result_len;

	encryptedMsg.insert(encryptedMsg.begin(), (char*)&securePacketHeader, ((char*)&securePacketHeader) + sizeof(VirtualRelayMessage::VR_SECURE_PACKET_HEADER));
	encryptedMsg.insert(encryptedMsg.end(), cipherData, cipherData + result_len);

	VirtualRelayMessage::AppendCRC(encryptedMsg);

	return encryptedMsg;
}

VirtualRelayCryptoProcessor::UInt8Vector VirtualRelayCryptoProcessor::DecryptPacket(UInt8Vector& securePacket)
{
	boost::uint8_t sizeIV = securePacket[1];
	int sizeCipherData = securePacket[2] + (securePacket[3]<<8);

	unsigned char* iv = &securePacket[4];
	unsigned char* cipherData = &securePacket[4+sizeIV];

	int result_len = 0;
	unsigned char level2Msg[sizeCipherData+8];
	activeCrypto->Decrypt(cipherData, sizeCipherData, iv, sizeIV, level2Msg, &result_len);

	UInt8Vector decryptedPacket((char*) level2Msg, ((char*) level2Msg) + result_len);

	return decryptedPacket;
}

}
}
