/*---------------------------------------------------------------------------*/
/* ntlm.cpp                                                                  */
/* copyright (c) innovaphone 2016                                            */
/*                                                                           */
/*---------------------------------------------------------------------------*/

// NOTE
//
// A lot of the NTLMv2 stuff was implemented with the help of the ntlm source code found at https://github.com/hamstergene/ntlm
// by Eugene Hmyakov (https://hamstergene.github.io/) and released under the MIT license.

#include <platform.h>

#include <ctype.h>
#include <stdexcept>
#include <algorithm>
#include <string>

#include "common/ilib/base64.h"
#include "common/ilib/cipher.h"
#include "common/ilib/hash.h"
#include "common/ilib/str.h"

#include "sdk/common/interface/random.h"
#include "sdk/common/interface/time.h"

#include "ntlm.h"

inline void WriteQWORD(void * addr, long64 data)
{
#if defined(BIGENDIAN)
	WRITE_QUAD_LE(addr, data);
#else
	*(static_cast<long64 *>(addr)) = data;
#endif
}


inline void WriteDWORD(void * addr, dword data)
{
#if defined(BIGENDIAN)
	WRITE_DWORD_LE(addr, data);
#else
	*(static_cast<dword *>(addr)) = data;
#endif
}


inline void WriteWORD(void * addr, word data)
{
#if defined(BIGENDIAN)
	WRITE_WORD_LE(addr, data);
#else
	*(static_cast<word *>(addr)) = data;
#endif
}


inline dword ReadDWORD(const void * addr)
{
#if defined(BIGENDIAN)
	return READ_DWORD_LE(addr);
#else
	return *(static_cast<const dword *>(addr));
#endif
}


inline word ReadWORD(const void * addr)
{
#if defined(BIGENDIAN)
	return READ_WORD_LE(addr);
#else
	return *(static_cast<const word *>(addr));
#endif
}


size_t WriteNTLMString(std::vector<byte> & buffer, const char * str, size_t strLen, int startIdx, size_t offset)
{
	WriteWORD(&buffer[startIdx], static_cast<word>(strLen));
	WriteWORD(&buffer[startIdx + sizeof(word)], static_cast<word>(strLen));
	WriteDWORD(&buffer[startIdx + sizeof(word) * 2], static_cast<dword>(offset));
	
	if (strLen > 0) {
		std::copy(str, str + strLen, &buffer[offset]);
	}
	
	return strLen;
}


size_t WriteNTLMString(std::vector<byte> & buffer, const std::string & str, int startIdx, size_t offset)
{
	return WriteNTLMString(buffer, str.c_str(), str.length(), startIdx, offset);
}


size_t WriteNTLMString(std::vector<byte> & buffer, const std::vector<char> & str, int startIdx, size_t offset)
{
	return WriteNTLMString(buffer, str.data(), str.size(), startIdx, offset);
}



/*---------------------------------------------------------------------------*/
/* class NTLMType1Message                                                    */
/*---------------------------------------------------------------------------*/

NTLMType1Message::NTLMType1Message(const char * domain, const char * workstation)
	: domain(domain == nullptr ? "" : domain),
	  workstation(workstation == nullptr ? "" : workstation),
	  flags(0)
{
	this->needBuild = true;
}


NTLMType1Message::~NTLMType1Message()
{
}


void NTLMType1Message::Build()
{
	size_t domLen = this->domain.length();
	size_t wsLen = this->workstation.length();
	this->buffer.clear();
	this->buffer.resize(32 + domLen + wsLen);
    
	memcpy(&buffer[IDX_NTLM_HEADER], "NTLMSSP", 7);
	WriteDWORD(&buffer[IDX_NTLM_MESSAGE_TYPE], NTLM_TYPE_1_MESSAGE);
	//this->SetFlag(NTLM_NEGOTIATE_UNICODE);
	
	size_t offset = IDX_TYPE_1_DATA_START_NO_OS_VERSION;
	if (domLen > 0) {
		this->SetFlag(NTLM_NEGOTIATE_DOMAIN);
		offset += WriteNTLMString(this->buffer, this->domain, IDX_TYPE_1_MESSAGE_DOMAIN_LEN, offset);
	}
	
	if (wsLen > 0) {
		this->SetFlag(NTLM_NEGOTIATE_WORKSTATION);
		WriteNTLMString(this->buffer, this->workstation, IDX_TYPE_1_MESSAGE_WORKSTATION_LEN, offset);
	}
	
	WriteDWORD(&buffer[IDX_TYPE_1_MESSAGE_FLAGS], this->flags);
	this->needBuild = false;
}


void NTLMType1Message::SetFlag(ntlm_flags flag)
{
	this->flags |= flag;
	this->needBuild = true;
}


dword NTLMType1Message::Flags()
{
	return this->flags;
}


byte * NTLMType1Message::GetByteStream()
{
	if (this->needBuild) this->Build();
	return this->buffer.data();
}


size_t NTLMType1Message::GetByteStreamLength()
{
	if (this->needBuild) this->Build();
	return this->buffer.size();
}

/*---------------------------------------------------------------------------*/
/* class NTLMType2Message                                                    */
/*---------------------------------------------------------------------------*/

NTLMType2Message::NTLMType2Message(const byte * b64Data)
{
	size_t len = strlen((char *)b64Data);
	// Make sure that the basic data length is available.
	if (len < 40) throw std::runtime_error("NTLMType2Message::NTLMType2Message(): Message byte stream to short.");
	
	this->buffer.resize(len);
    decode_base64((char *)b64Data, this->buffer.data(), len);
	
	if (!this->CheckHeader()) throw std::runtime_error("NTLMType2Message::NTLMType2Message(): NTLM response is not a valid type 2 message.");
	this->flags = ReadDWORD(&buffer[IDX_TYPE_2_FLAGS]);
	if (!this->FlagSet(NTLM_NEGOTIATE_NTLM) && !this->FlagSet(NTLM_NEGOTIATE_NTLM2)) throw std::runtime_error("NTLMType2Message::NTLMType2Message(): The response message doesn neither support NTLMv1 nor NLTMv2");
	this->ReadTargetName();
}


NTLMType2Message::~NTLMType2Message()
{
}


std::string NTLMType2Message::ReadUNICODEString(const byte * ucstrBuffer, size_t len)
{
	std::string result = "";
	if (ucstrBuffer != nullptr && len > 0) {
		std::vector<char> tmp(len * 2);
		str::from_ucs2(reinterpret_cast<const word *>(ucstrBuffer), tmp.data(), len / 2 + 1); // + 1 for the closing zero.
		result = tmp.data();
	}
	
	return result;
}


bool NTLMType2Message::CheckHeader()
{
	return (memcmp(&this->buffer[IDX_NTLM_HEADER], "NTLMSSP", 7) == 0 && ReadDWORD(&this->buffer[IDX_NTLM_MESSAGE_TYPE]) == NTLM_TYPE_2_MESSAGE);
}


void NTLMType2Message::ReadTargetName()
{
	size_t tlen = ReadWORD(&this->buffer[IDX_TYPE_2_TARGET_NAME_LEN]);
	size_t toffset = ReadDWORD(&this->buffer[IDX_TYPE_2_TARGET_NAME_OFFSET]);
	
	if (tlen > 0) {
		if (buffer.size() < toffset + tlen) throw std::runtime_error("NTLMType2Message::ReadTargetName(): indicated target name not available in data stream.");
		if (this->FlagSet(NTLM_NEGOTIATE_UNICODE)) this->targetName = this->ReadUNICODEString(&this->buffer[toffset], tlen);
		else this->targetName.assign(reinterpret_cast<const char *>(&this->buffer[toffset]), tlen);
	}
}


const byte * NTLMType2Message::Challenge()
{
	return &this->buffer[IDX_TYPE_2_CHALLENGE];
}


size_t NTLMType2Message::TargetInformationSize()
{
	return ReadWORD(&this->buffer[IDX_TYPE_2_TARGET_INFORMATION_LEN]);
}


const byte * NTLMType2Message::TargetInformationData()
{
	dword of = ReadDWORD(&this->buffer[IDX_TYPE_2_TARGET_INFORMATION_OFFSET]);
	return &this->buffer[of];
}


dword NTLMType2Message::Flags()
{
    return this->flags;
}


/*---------------------------------------------------------------------------*/
/* class NTLMType3Message                                                    */
/*---------------------------------------------------------------------------*/

NTLMType3Message::NTLMType3Message(const char * user, 
                                   const char * password,
                                   const char * domainName,
                                   const char * workstationName,
                                   NTLMType2Message * const type2Msg)
    : type2Message(type2Msg),
      user(user == nullptr ? "" : user),
      password(password == nullptr ? "" : password),
	  domainName(domainName == nullptr ? "" : domainName),
      workstationName(workstationName == nullptr ? "" : workstationName),
	  flags(0)
{
	if (this->user.empty() || this->user == "") throw std::runtime_error("No user name given to build type 3 response message");
	this->needBuild = true;
	this->useUnicode = this->type2Message->FlagSet(NTLM_NEGOTIATE_UNICODE);
}


NTLMType3Message::~NTLMType3Message()
{
}


void NTLMType3Message::AllocBuffer()
{
	// The buffer size to calculate is easy for NLTM1:
	//
	//  64 byte for the message base
	//  + length of the domain name (multiplied by 2 in case of unicode)
	//  + length of user name (multiplied by 2 in case of unicode)
	//  + lenght of local workstation name (multiplied by.... you got it)
	//  + LN response (always 24 bytes)
	//  + NT response (in case of NTLMv1 always 24 byte; for NTLMv2 the size must be calculated)
	if (this->type2Message->FlagSet(NTLM_NEGOTIATE_NTLM2)) {
		// The size of the NTv2 response part is 28 for the blob size plus 16 for the NTLMv2 hash.
		this->ntResponseSize = this->type2Message->TargetInformationSize() + NTLM_TYPE_3_MESSAGE_NT2_BLOB_BASE_SIZE + NTLM_TYPE_3_MESSAGE_NT2_HASH_SIZE;
	}
	else {
		this->ntResponseSize = NTLM_TYPE_3_MESSAGE_NT1_RESPONSE_SIZE;
	}
	
	size_t strSizes = this->domainName.length() + this->user.length() + this->workstationName.length();
	if (this->useUnicode) strSizes *= 2;
	
	this->buffer.resize(NTLM_TYPE_3_MESSAGE_BASE_SIZE + NTLM_TYPE_3_MESSAGE_LM_RESPONSE_SIZE + strSizes + this->ntResponseSize);
}


void NTLMType3Message::CreateHashHmacMD5(const byte * data, size_t dataLen, const byte * key, size_t keyLen, byte * out)
{
	hmac hmacMD5;
	hmacMD5.init(HASH_MD5, key, keyLen);
	hmacMD5.update(data, dataLen);
	hmacMD5.final(out);
}


std::vector<char> NTLMType3Message::StringToUnicode(const std::string & str)
{
	std::vector<char> result;
	if (!str.empty()) {
		result.resize(str.length() * 2);
		str::to_ucs2_n(str.c_str(), reinterpret_cast<word *>(result.data()), str.length());
#if defined(BIGENDIAN)
		str::ucs2_to_le_n(reinterpret_cast<word *>(result.data()), reinterpret_cast<byte *>(result.data()), str.length());
#endif
	}
	return result;
}


void NTLMType3Message::Build()
{
	// AllocBuffer() also converts the strings into unicode strings, if neccessary...
	this->AllocBuffer();
	this->SetFlag(NTLM_NEGOTIATE_ALWAYS_SIGN); 
	
	memcpy(&buffer[IDX_NTLM_HEADER], "NTLMSSP", 7);
	WriteDWORD(&buffer[IDX_NTLM_MESSAGE_TYPE], NTLM_TYPE_3_MESSAGE);

	// The security key is optional, so it's not supported for now.
	WriteWORD(&this->buffer[IDX_SESSION_KEY_LENGTH], 0);
	WriteWORD(&this->buffer[IDX_SESSION_KEY_LENGTH_MAX], 0);
	WriteWORD(&this->buffer[IDX_SESSION_KEY_OFFSET], 0);

	size_t offset = NTLM_TYPE_3_MESSAGE_BASE_SIZE;
	if (this->useUnicode) {
		this->SetFlag(NTLM_NEGOTIATE_UNICODE);
		offset += WriteNTLMString(this->buffer, this->StringToUnicode(this->domainName), IDX_DOMAIN_LENGTH, offset);
		offset += WriteNTLMString(this->buffer, this->StringToUnicode(this->user), IDX_USER_NAME_LENGTH, offset);
		offset += WriteNTLMString(this->buffer, this->StringToUnicode(this->workstationName), IDX_WORKSTATION_LENGTH, offset);
	}
	else {
		this->SetFlag(NTLM_NEGOTIATE_OEM);
		offset += WriteNTLMString(this->buffer, this->domainName, IDX_DOMAIN_LENGTH, offset);
		offset += WriteNTLMString(this->buffer, this->user, IDX_USER_NAME_LENGTH, offset);
		offset += WriteNTLMString(this->buffer, this->workstationName, IDX_WORKSTATION_LENGTH, offset);
	}
	
	// Size and offset for the LM/LMv2 data block
	WriteWORD(&buffer[IDX_LM_RESPONSE_LENGTH], NTLM_TYPE_3_MESSAGE_LM_RESPONSE_SIZE);
	WriteWORD(&buffer[IDX_LM_RESPONSE_LENGTH_MAX], NTLM_TYPE_3_MESSAGE_LM_RESPONSE_SIZE);
	WriteDWORD(&buffer[IDX_LM_RESPONSE_OFFSET], offset);

	// Size and offset for the NT/NTv2 data block...
	WriteWORD(&buffer[IDX_NT_RESPONSE_LENGTH], static_cast<word>(this->ntResponseSize));
	WriteWORD(&buffer[IDX_NT_RESPONSE_LENGTH_MAX], static_cast<word>(this->ntResponseSize));
	WriteDWORD(&buffer[IDX_NT_RESPONSE_OFFSET], offset + NTLM_TYPE_3_MESSAGE_LM_RESPONSE_SIZE);
	
	if (this->type2Message->FlagSet(NTLM_NEGOTIATE_NTLM2)) {
		byte v1PwdHash[HASH_SIZE_MD4];
		byte hash[HASH_SIZE_MD5];
		IRandom::GetChars(reinterpret_cast<char *>(this->v2ClientNonce), 8);

		this->SetFlag(NTLM_NEGOTIATE_NTLM2);
		this->CreateV2Hash(hash, v1PwdHash);
		
		offset += this->WriteLMv2Response(offset, hash, v1PwdHash);
		this->WriteNT2Response(offset, hash);
	}
	else {
		this->SetFlag(NTLM_NEGOTIATE_NTLM);
		offset += this->WriteLMResponse(offset);
		this->WriteNTResponse(offset);
	}

	WriteWORD(&this->buffer[IDX_TYPE_3_MESSAGE_FLAGS], this->flags);
	this->needBuild = false;
}


size_t NTLMType3Message::WriteLMResponse(size_t offset)
{
    size_t pwdLen = this->password.length();
    if (pwdLen > 14) pwdLen = 14;
    byte lm_pwd[14];

	for (size_t i = 0; i < pwdLen; i++) {
		lm_pwd[i] = toupper(password[i]);
	}

	if (pwdLen < 14) {
		memset(&lm_pwd[pwdLen], 0, 14 - pwdLen);
	}

    const byte LMR_MAGIC[] = "KGS!@#$%";
    byte lm_resp[21];
    Cipher cipher;
    cipher.Init(CIPHER_DES, lm_pwd, 7, true);
    cipher.Block(LMR_MAGIC, lm_resp);
    cipher.Init(CIPHER_DES, lm_pwd + 7, 7, true);
    cipher.Block(LMR_MAGIC, lm_resp + 8);
    memset(lm_resp + 16, 0, 5);

	byte * lmr_dest = &buffer[offset];
	const byte * nonce = this->type2Message->Challenge();
    cipher.Init(CIPHER_DES, lm_resp, 7, true);
    cipher.Block(nonce, lmr_dest);
    cipher.Init(CIPHER_DES, lm_resp + 7, 7, true);
    cipher.Block(nonce, lmr_dest + 8);
    cipher.Init(CIPHER_DES, lm_resp + 14, 7, true);
    cipher.Block(nonce, lmr_dest + 16);
	
	return 24;	
}


size_t NTLMType3Message::WriteLMv2Response(size_t offset, const byte * v2Hash, const byte * v1PwdHash)
{
	byte challenge[16];
	std::copy(this->type2Message->Challenge(), this->type2Message->Challenge() + 8, &challenge[0]);
	std::copy(std::begin(this->v2ClientNonce), std::end(this->v2ClientNonce), &challenge[8]);
	
	this->CreateHashHmacMD5(challenge, sizeof(challenge), v2Hash, HASH_SIZE_MD5, &this->buffer[offset]);
	std::copy(this->v2ClientNonce, &this->v2ClientNonce[8], &this->buffer[offset + HASH_SIZE_MD5]);	
	return 24; 
}


void NTLMType3Message::WriteNTResponse(size_t offset)
{
	byte nt_resp[21];
	this->CreateNTLMPasswordHash(nt_resp);
	std::fill(std::begin(nt_resp) + HASH_SIZE_MD4, std::end(nt_resp), 0);
	
    byte * ntr_dest = &this->buffer[offset];
    Cipher cipher;
	const byte * nonce = this->type2Message->Challenge();
    cipher.Init(CIPHER_DES, nt_resp, 7, true);
    cipher.Block(nonce, ntr_dest);
    cipher.Init(CIPHER_DES, nt_resp + 7, 7, true);
    cipher.Block(nonce, ntr_dest + 8);
    cipher.Init(CIPHER_DES, nt_resp + 14, 7, true);
    cipher.Block(nonce, ntr_dest + 16);
}


void NTLMType3Message::WriteNT2Response(size_t offset, const byte * v2Hash)
{
	NTv2Blob blob(this->type2Message, this->v2ClientNonce);
	size_t bs = blob.Size();
	std::vector<byte> challengeAndBlob(NTLM_TYPE_2_MESSAGE_CHALLENGE_SIZE + bs);
	std::copy(this->type2Message->Challenge(), this->type2Message->Challenge() + NTLM_TYPE_2_MESSAGE_CHALLENGE_SIZE, &challengeAndBlob[0]);
	std::copy(blob.Begin(), blob.End(), &challengeAndBlob[NTLM_TYPE_2_MESSAGE_CHALLENGE_SIZE]);
	this->CreateHashHmacMD5(challengeAndBlob.data(), challengeAndBlob.size(), v2Hash, HASH_SIZE_MD5, &this->buffer[offset]);
	std::copy(blob.Begin(), blob.End(), &buffer[offset + HASH_SIZE_MD5]);
}


void NTLMType3Message::CreateNTLMPasswordHash(byte * dest)
{
	if (!this->password.empty()) {
		size_t pwdLen = this->password.length();
		std::vector<word> ntPwd(pwdLen);
		str::to_ucs2_n(this->password.c_str(), ntPwd.data(), pwdLen);

		hash md4;
		md4.init(HASH_MD4);
		md4.update(ntPwd.data(), pwdLen * 2);
		md4.final(dest);
	}
}


void NTLMType3Message::CreateV2Hash(byte * hash, byte * v1PasswordHash)
{
	this->CreateNTLMPasswordHash(v1PasswordHash);
	
	std::string userUpper = this->user;
	for (char & s : userUpper) s = toupper(s);
	
	std::vector<byte> userAndDomain(userUpper.length() * 2 + this->domainName.length() * 2);
	str::to_ucs2_n(userUpper.c_str(), reinterpret_cast<word *>(userAndDomain.data()), userUpper.length());
	str::to_ucs2_n(this->domainName.c_str(), reinterpret_cast<word *>(userAndDomain.data() + userUpper.length() * 2), this->domainName.length());
	
	this->CreateHashHmacMD5(userAndDomain.data(), userAndDomain.size(), v1PasswordHash, HASH_SIZE_MD4, hash);
}


void NTLMType3Message::SetFlag(ntlm_flags flag)
{
	this->flags |= flag;
	this->needBuild = true;
}


dword NTLMType3Message::GetFlags()
{
	return this->flags;
}


const byte * NTLMType3Message::GetByteStream()
{
	if (this->needBuild) this->Build();
	return this->buffer.data();
}


size_t NTLMType3Message::GetByteStreamLength()
{
	if (this->needBuild) this->Build();
	return this->buffer.size();
}


/*---------------------------------------------------------------------------*/
/* class NTv2Blob                                                            */
/*---------------------------------------------------------------------------*/

NTv2Blob::NTv2Blob(NTLMType2Message * type2Message, const byte (&clientNonce)[8])
{
	size_t tiSize = type2Message->TargetInformationSize();
	const byte * tiBegin = type2Message->TargetInformationData();
	const byte * tiEnd = tiBegin + tiSize;
	this->buffer.resize(tiSize + NTLM_TYPE_3_MESSAGE_NT2_BLOB_BASE_SIZE, 0);

	// Since the vector is initialized with zeros, the unkown and reserved values of the blob don't neet to be written (see
	// description of the blob in the header file). The timestamp will be converted to windows timestamps, which start
	// at January 1st, 1601. And they have to be in tenth of milliseconds.
	WriteWORD(&this->buffer[IDX_V2_BLOB_SIGNATURE], 0x0101);
	WriteQWORD(&buffer[IDX_V2_BLOB_TIMESTAMPT], (ITime::TimeStampMilliseconds() + 11644473600000) * 10000);


	std::copy(std::begin(clientNonce), std::end(clientNonce), &this->buffer[IDX_V2_BLOB_CLIENT_CONCE]);
	std::copy(tiBegin, tiEnd, &this->buffer[IDX_V2_BLOB_TARGET_INFORMATION]);
}