/**
 * WinPR: Windows Portable Runtime
 * NTLM Security Package
 *
 * Copyright 2011-2014 Marc-Andre Moreau <marcandre.moreau@gmail.com>
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <winpr/config.h>

#include <winpr/crt.h>
#include <winpr/assert.h>
#include <winpr/sspi.h>
#include <winpr/print.h>
#include <winpr/string.h>
#include <winpr/tchar.h>
#include <winpr/sysinfo.h>
#include <winpr/registry.h>
#include <winpr/endian.h>
#include <winpr/build-config.h>

#include "ntlm.h"
#include "ntlm_export.h"
#include "../sspi.h"

#include "ntlm_message.h"

#include "../../utils.h"

#include "../../log.h"
#define TAG WINPR_TAG("sspi.NTLM")

#define WINPR_KEY "Software\\%s\\WinPR\\NTLM"

static char* NTLM_PACKAGE_NAME = "NTLM";

#define check_context(ctx) check_context_((ctx), __FILE__, __func__, __LINE__)
static BOOL check_context_(NTLM_CONTEXT* context, const char* file, const char* fkt, size_t line)
{
	BOOL rc = TRUE;
	wLog* log = WLog_Get(TAG);
	const DWORD log_level = WLOG_ERROR;

	if (!context)
	{
		if (WLog_IsLevelActive(log, log_level))
			WLog_PrintTextMessage(log, log_level, line, file, fkt, "invalid context");

		return FALSE;
	}

	if (!context->RecvRc4Seal)
	{
		if (WLog_IsLevelActive(log, log_level))
			WLog_PrintTextMessage(log, log_level, line, file, fkt, "invalid context->RecvRc4Seal");
		rc = FALSE;
	}
	if (!context->SendRc4Seal)
	{
		if (WLog_IsLevelActive(log, log_level))
			WLog_PrintTextMessage(log, log_level, line, file, fkt, "invalid context->SendRc4Seal");
		rc = FALSE;
	}

	if (!context->SendSigningKey)
	{
		if (WLog_IsLevelActive(log, log_level))
			WLog_PrintTextMessage(log, log_level, line, file, fkt,
			                      "invalid context->SendSigningKey");
		rc = FALSE;
	}
	if (!context->RecvSigningKey)
	{
		if (WLog_IsLevelActive(log, log_level))
			WLog_PrintTextMessage(log, log_level, line, file, fkt,
			                      "invalid context->RecvSigningKey");
		rc = FALSE;
	}
	if (!context->SendSealingKey)
	{
		if (WLog_IsLevelActive(log, log_level))
			WLog_PrintTextMessage(log, log_level, line, file, fkt,
			                      "invalid context->SendSealingKey");
		rc = FALSE;
	}
	if (!context->RecvSealingKey)
	{
		if (WLog_IsLevelActive(log, log_level))
			WLog_PrintTextMessage(log, log_level, line, file, fkt,
			                      "invalid context->RecvSealingKey");
		rc = FALSE;
	}
	return rc;
}

static char* get_name(COMPUTER_NAME_FORMAT type)
{
	DWORD nSize = 0;

	if (GetComputerNameExA(type, nullptr, &nSize))
		return nullptr;

	if (GetLastError() != ERROR_MORE_DATA)
		return nullptr;

	char* computerName = calloc(1, nSize);

	if (!computerName)
		return nullptr;

	if (!GetComputerNameExA(type, computerName, &nSize))
	{
		free(computerName);
		return nullptr;
	}

	return computerName;
}

static int ntlm_SetContextWorkstation(NTLM_CONTEXT* context, char* Workstation)
{
	char* ws = Workstation;
	CHAR* computerName = nullptr;

	WINPR_ASSERT(context);

	if (!Workstation)
	{
		computerName = get_name(ComputerNameNetBIOS);
		if (!computerName)
			return -1;
		ws = computerName;
	}

	size_t len = 0;
	context->Workstation.Buffer = ConvertUtf8ToWCharAlloc(ws, &len);

	free(computerName);

	if (!context->Workstation.Buffer || (len > UINT16_MAX / sizeof(WCHAR)))
		return -1;

	context->Workstation.Length = (USHORT)(len * sizeof(WCHAR));
	return 1;
}

static int ntlm_SetContextServicePrincipalNameW(NTLM_CONTEXT* context, LPWSTR ServicePrincipalName)
{
	WINPR_ASSERT(context);

	if (!ServicePrincipalName)
	{
		context->ServicePrincipalName.Buffer = nullptr;
		context->ServicePrincipalName.Length = 0;
		return 1;
	}

	context->ServicePrincipalName.Length = (USHORT)(_wcslen(ServicePrincipalName) * 2);
	context->ServicePrincipalName.Buffer = (PWSTR)malloc(context->ServicePrincipalName.Length + 2);

	if (!context->ServicePrincipalName.Buffer)
		return -1;

	memcpy(context->ServicePrincipalName.Buffer, ServicePrincipalName,
	       context->ServicePrincipalName.Length + 2);
	return 1;
}

static int ntlm_SetContextTargetName(NTLM_CONTEXT* context, char* TargetName)
{
	char* name = TargetName;
	DWORD nSize = 0;
	CHAR* computerName = nullptr;

	WINPR_ASSERT(context);

	if (!name)
	{
		if (GetComputerNameExA(ComputerNameNetBIOS, nullptr, &nSize) ||
		    GetLastError() != ERROR_MORE_DATA)
			return -1;

		computerName = calloc(nSize, sizeof(CHAR));

		if (!computerName)
			return -1;

		if (!GetComputerNameExA(ComputerNameNetBIOS, computerName, &nSize))
		{
			free(computerName);
			return -1;
		}

		if (nSize > MAX_COMPUTERNAME_LENGTH)
			computerName[MAX_COMPUTERNAME_LENGTH] = '\0';

		name = computerName;

		if (!name)
			return -1;

		CharUpperA(name);
	}

	size_t len = 0;
	context->TargetName.pvBuffer = ConvertUtf8ToWCharAlloc(name, &len);

	if (!context->TargetName.pvBuffer || (len > UINT16_MAX / sizeof(WCHAR)))
	{
		free(context->TargetName.pvBuffer);
		context->TargetName.pvBuffer = nullptr;

		if (!TargetName)
			free(name);

		return -1;
	}

	context->TargetName.cbBuffer = (USHORT)(len * sizeof(WCHAR));

	if (!TargetName)
		free(name);

	return 1;
}

static NTLM_CONTEXT* ntlm_ContextNew(void)
{
	HKEY hKey = nullptr;
	DWORD dwType = 0;
	DWORD dwSize = 0;
	DWORD dwValue = 0;
	NTLM_CONTEXT* context = (NTLM_CONTEXT*)calloc(1, sizeof(NTLM_CONTEXT));

	if (!context)
		return nullptr;

	context->NTLMv2 = TRUE;
	context->UseMIC = FALSE;
	context->SendVersionInfo = TRUE;
	context->SendSingleHostData = FALSE;
	context->SendWorkstationName = TRUE;
	context->NegotiateKeyExchange = TRUE;
	context->UseSamFileDatabase = TRUE;

	{
		char* key = winpr_getApplicatonDetailsRegKey(WINPR_KEY);
		if (key)
		{
			const LONG status =
			    RegOpenKeyExA(HKEY_LOCAL_MACHINE, key, 0, KEY_READ | KEY_WOW64_64KEY, &hKey);
			free(key);

			if (status == ERROR_SUCCESS)
			{
				if (RegQueryValueEx(hKey, _T("NTLMv2"), nullptr, &dwType, (BYTE*)&dwValue,
				                    &dwSize) == ERROR_SUCCESS)
					context->NTLMv2 = dwValue ? 1 : 0;

				if (RegQueryValueEx(hKey, _T("UseMIC"), nullptr, &dwType, (BYTE*)&dwValue,
				                    &dwSize) == ERROR_SUCCESS)
					context->UseMIC = dwValue ? 1 : 0;

				if (RegQueryValueEx(hKey, _T("SendVersionInfo"), nullptr, &dwType, (BYTE*)&dwValue,
				                    &dwSize) == ERROR_SUCCESS)
					context->SendVersionInfo = dwValue ? 1 : 0;

				if (RegQueryValueEx(hKey, _T("SendSingleHostData"), nullptr, &dwType,
				                    (BYTE*)&dwValue, &dwSize) == ERROR_SUCCESS)
					context->SendSingleHostData = dwValue ? 1 : 0;

				if (RegQueryValueEx(hKey, _T("SendWorkstationName"), nullptr, &dwType,
				                    (BYTE*)&dwValue, &dwSize) == ERROR_SUCCESS)
					context->SendWorkstationName = dwValue ? 1 : 0;

				if (RegQueryValueEx(hKey, _T("WorkstationName"), nullptr, &dwType, nullptr,
				                    &dwSize) == ERROR_SUCCESS)
				{
					char* workstation = (char*)malloc(dwSize + 1);

					if (!workstation)
					{
						free(context);
						return nullptr;
					}

					const LONG rc = RegQueryValueExA(hKey, "WorkstationName", nullptr, &dwType,
					                                 (BYTE*)workstation, &dwSize);
					if (rc != ERROR_SUCCESS)
						WLog_WARN(TAG, "Key ''WorkstationName' not found");
					workstation[dwSize] = '\0';

					if (ntlm_SetContextWorkstation(context, workstation) < 0)
					{
						free(workstation);
						free(context);
						return nullptr;
					}

					free(workstation);
				}

				RegCloseKey(hKey);
			}
		}
	}

	/*
	 * Extended Protection is enabled by default in Windows 7,
	 * but enabling it in WinPR breaks TS Gateway at this point
	 */
	context->SuppressExtendedProtection = FALSE;
	const LONG status =
	    RegOpenKeyEx(HKEY_LOCAL_MACHINE, _T("System\\CurrentControlSet\\Control\\LSA"), 0,
	                 KEY_READ | KEY_WOW64_64KEY, &hKey);

	if (status == ERROR_SUCCESS)
	{
		if (RegQueryValueEx(hKey, _T("SuppressExtendedProtection"), nullptr, &dwType,
		                    (BYTE*)&dwValue, &dwSize) == ERROR_SUCCESS)
			context->SuppressExtendedProtection = dwValue ? 1 : 0;

		RegCloseKey(hKey);
	}

	context->NegotiateFlags = 0;
	context->LmCompatibilityLevel = 3;
	ntlm_change_state(context, NTLM_STATE_INITIAL);
	FillMemory(context->MachineID, sizeof(context->MachineID), 0xAA);

	if (context->NTLMv2)
		context->UseMIC = TRUE;

	return context;
}

static void ntlm_ContextFree(NTLM_CONTEXT* context)
{
	if (!context)
		return;

	winpr_RC4_Free(context->SendRc4Seal);
	winpr_RC4_Free(context->RecvRc4Seal);
	sspi_SecBufferFree(&context->NegotiateMessage);
	sspi_SecBufferFree(&context->ChallengeMessage);
	sspi_SecBufferFree(&context->AuthenticateMessage);
	sspi_SecBufferFree(&context->ChallengeTargetInfo);
	sspi_SecBufferFree(&context->TargetName);
	sspi_SecBufferFree(&context->NtChallengeResponse);
	sspi_SecBufferFree(&context->LmChallengeResponse);
	free(context->ServicePrincipalName.Buffer);
	free(context->Workstation.Buffer);
	free(context);
}

static SECURITY_STATUS SEC_ENTRY ntlm_AcquireCredentialsHandleW(
    WINPR_ATTR_UNUSED SEC_WCHAR* pszPrincipal, WINPR_ATTR_UNUSED SEC_WCHAR* pszPackage,
    ULONG fCredentialUse, WINPR_ATTR_UNUSED void* pvLogonID, void* pAuthData,
    SEC_GET_KEY_FN pGetKeyFn, void* pvGetKeyArgument, PCredHandle phCredential,
    WINPR_ATTR_UNUSED PTimeStamp ptsExpiry)
{
	SEC_WINPR_NTLM_SETTINGS* settings = nullptr;

	if ((fCredentialUse != SECPKG_CRED_OUTBOUND) && (fCredentialUse != SECPKG_CRED_INBOUND) &&
	    (fCredentialUse != SECPKG_CRED_BOTH))
	{
		return SEC_E_INVALID_PARAMETER;
	}

	SSPI_CREDENTIALS* credentials = sspi_CredentialsNew();

	if (!credentials)
		return SEC_E_INTERNAL_ERROR;

	credentials->fCredentialUse = fCredentialUse;
	credentials->pGetKeyFn = pGetKeyFn;
	credentials->pvGetKeyArgument = pvGetKeyArgument;

	if (pAuthData)
	{
		UINT32 identityFlags = sspi_GetAuthIdentityFlags(pAuthData);

		if (sspi_CopyAuthIdentity(&(credentials->identity),
		                          (const SEC_WINNT_AUTH_IDENTITY_INFO*)pAuthData) < 0)
		{
			sspi_CredentialsFree(credentials);
			return SEC_E_INVALID_PARAMETER;
		}

		if (identityFlags & SEC_WINNT_AUTH_IDENTITY_EXTENDED)
			settings = (((SEC_WINNT_AUTH_IDENTITY_WINPR*)pAuthData)->ntlmSettings);
	}

	if (settings)
	{
		if (settings->samFile)
		{
			credentials->ntlmSettings.samFile = _strdup(settings->samFile);
			if (!credentials->ntlmSettings.samFile)
			{
				sspi_CredentialsFree(credentials);
				return SEC_E_INSUFFICIENT_MEMORY;
			}
		}
		credentials->ntlmSettings.hashCallback = settings->hashCallback;
		credentials->ntlmSettings.hashCallbackArg = settings->hashCallbackArg;
	}

	sspi_SecureHandleSetLowerPointer(phCredential, (void*)credentials);
	sspi_SecureHandleSetUpperPointer(phCredential, (void*)NTLM_PACKAGE_NAME);
	return SEC_E_OK;
}

static SECURITY_STATUS SEC_ENTRY ntlm_AcquireCredentialsHandleA(
    SEC_CHAR* pszPrincipal, SEC_CHAR* pszPackage, ULONG fCredentialUse, void* pvLogonID,
    void* pAuthData, SEC_GET_KEY_FN pGetKeyFn, void* pvGetKeyArgument, PCredHandle phCredential,
    PTimeStamp ptsExpiry)
{
	SECURITY_STATUS status = SEC_E_INSUFFICIENT_MEMORY;
	SEC_WCHAR* principal = nullptr;
	SEC_WCHAR* package = nullptr;

	if (pszPrincipal)
	{
		principal = ConvertUtf8ToWCharAlloc(pszPrincipal, nullptr);
		if (!principal)
			goto fail;
	}
	if (pszPackage)
	{
		package = ConvertUtf8ToWCharAlloc(pszPackage, nullptr);
		if (!package)
			goto fail;
	}

	status =
	    ntlm_AcquireCredentialsHandleW(principal, package, fCredentialUse, pvLogonID, pAuthData,
	                                   pGetKeyFn, pvGetKeyArgument, phCredential, ptsExpiry);

fail:
	free(principal);
	free(package);

	return status;
}

static SECURITY_STATUS SEC_ENTRY ntlm_FreeCredentialsHandle(PCredHandle phCredential)
{
	if (!phCredential)
		return SEC_E_INVALID_HANDLE;

	SSPI_CREDENTIALS* credentials =
	    (SSPI_CREDENTIALS*)sspi_SecureHandleGetLowerPointer(phCredential);

	if (!credentials)
		return SEC_E_INVALID_HANDLE;

	sspi_CredentialsFree(credentials);
	return SEC_E_OK;
}

static SECURITY_STATUS SEC_ENTRY ntlm_QueryCredentialsAttributesW(
    WINPR_ATTR_UNUSED PCredHandle phCredential, WINPR_ATTR_UNUSED ULONG ulAttribute,
    WINPR_ATTR_UNUSED void* pBuffer)
{
	if (ulAttribute == SECPKG_CRED_ATTR_NAMES)
	{
		return SEC_E_OK;
	}

	WLog_ERR(TAG, "TODO: Implement");
	return SEC_E_UNSUPPORTED_FUNCTION;
}

static SECURITY_STATUS SEC_ENTRY ntlm_QueryCredentialsAttributesA(PCredHandle phCredential,
                                                                  ULONG ulAttribute, void* pBuffer)
{
	return ntlm_QueryCredentialsAttributesW(phCredential, ulAttribute, pBuffer);
}

/**
 * @see http://msdn.microsoft.com/en-us/library/windows/desktop/aa374707
 */
static SECURITY_STATUS SEC_ENTRY ntlm_AcceptSecurityContext(
    PCredHandle phCredential, PCtxtHandle phContext, PSecBufferDesc pInput, ULONG fContextReq,
    WINPR_ATTR_UNUSED ULONG TargetDataRep, PCtxtHandle phNewContext, PSecBufferDesc pOutput,
    WINPR_ATTR_UNUSED PULONG pfContextAttr, WINPR_ATTR_UNUSED PTimeStamp ptsTimeStamp)
{
	SECURITY_STATUS status = 0;
	SSPI_CREDENTIALS* credentials = nullptr;
	PSecBuffer input_buffer = nullptr;
	PSecBuffer output_buffer = nullptr;

	/* behave like windows SSPIs that don't want empty context */
	if (phContext && !phContext->dwLower && !phContext->dwUpper)
		return SEC_E_INVALID_HANDLE;

	NTLM_CONTEXT* context = (NTLM_CONTEXT*)sspi_SecureHandleGetLowerPointer(phContext);

	if (!context)
	{
		context = ntlm_ContextNew();

		if (!context)
			return SEC_E_INSUFFICIENT_MEMORY;

		context->server = TRUE;

		if (fContextReq & ASC_REQ_CONFIDENTIALITY)
			context->confidentiality = TRUE;

		credentials = (SSPI_CREDENTIALS*)sspi_SecureHandleGetLowerPointer(phCredential);
		context->credentials = credentials;
		context->SamFile = credentials->ntlmSettings.samFile;
		context->HashCallback = credentials->ntlmSettings.hashCallback;
		context->HashCallbackArg = credentials->ntlmSettings.hashCallbackArg;

		ntlm_SetContextTargetName(context, nullptr);
		sspi_SecureHandleSetLowerPointer(phNewContext, context);
		sspi_SecureHandleSetUpperPointer(phNewContext, (void*)NTLM_PACKAGE_NAME);
	}

	switch (ntlm_get_state(context))
	{
		case NTLM_STATE_INITIAL:
		{
			ntlm_change_state(context, NTLM_STATE_NEGOTIATE);

			if (!pInput)
				return SEC_E_INVALID_TOKEN;

			if (pInput->cBuffers < 1)
				return SEC_E_INVALID_TOKEN;

			input_buffer = sspi_FindSecBuffer(pInput, SECBUFFER_TOKEN);

			if (!input_buffer)
				return SEC_E_INVALID_TOKEN;

			if (input_buffer->cbBuffer < 1)
				return SEC_E_INVALID_TOKEN;

			status = ntlm_read_NegotiateMessage(context, input_buffer);
			if (status != SEC_I_CONTINUE_NEEDED)
				return status;

			if (ntlm_get_state(context) == NTLM_STATE_CHALLENGE)
			{
				if (!pOutput)
					return SEC_E_INVALID_TOKEN;

				if (pOutput->cBuffers < 1)
					return SEC_E_INVALID_TOKEN;

				output_buffer = sspi_FindSecBuffer(pOutput, SECBUFFER_TOKEN);

				if (!output_buffer->BufferType)
					return SEC_E_INVALID_TOKEN;

				if (output_buffer->cbBuffer < 1)
					return SEC_E_INSUFFICIENT_MEMORY;

				return ntlm_write_ChallengeMessage(context, output_buffer);
			}

			return SEC_E_OUT_OF_SEQUENCE;
		}

		case NTLM_STATE_AUTHENTICATE:
		{
			if (!pInput)
				return SEC_E_INVALID_TOKEN;

			if (pInput->cBuffers < 1)
				return SEC_E_INVALID_TOKEN;

			input_buffer = sspi_FindSecBuffer(pInput, SECBUFFER_TOKEN);

			if (!input_buffer)
				return SEC_E_INVALID_TOKEN;

			if (input_buffer->cbBuffer < 1)
				return SEC_E_INVALID_TOKEN;

			status = ntlm_read_AuthenticateMessage(context, input_buffer);

			if (pOutput)
			{
				for (ULONG i = 0; i < pOutput->cBuffers; i++)
				{
					pOutput->pBuffers[i].cbBuffer = 0;
					pOutput->pBuffers[i].BufferType = SECBUFFER_TOKEN;
				}
			}

			return status;
		}

		default:
			return SEC_E_OUT_OF_SEQUENCE;
	}
}

static SECURITY_STATUS SEC_ENTRY
ntlm_ImpersonateSecurityContext(WINPR_ATTR_UNUSED PCtxtHandle phContext)
{
	return SEC_E_OK;
}

static SECURITY_STATUS SEC_ENTRY ntlm_InitializeSecurityContextW(
    PCredHandle phCredential, PCtxtHandle phContext, SEC_WCHAR* pszTargetName, ULONG fContextReq,
    WINPR_ATTR_UNUSED ULONG Reserved1, WINPR_ATTR_UNUSED ULONG TargetDataRep, PSecBufferDesc pInput,
    WINPR_ATTR_UNUSED ULONG Reserved2, PCtxtHandle phNewContext, PSecBufferDesc pOutput,
    WINPR_ATTR_UNUSED PULONG pfContextAttr, WINPR_ATTR_UNUSED PTimeStamp ptsExpiry)
{
	SECURITY_STATUS status = 0;
	SSPI_CREDENTIALS* credentials = nullptr;
	PSecBuffer input_buffer = nullptr;
	PSecBuffer output_buffer = nullptr;

	/* behave like windows SSPIs that don't want empty context */
	if (phContext && !phContext->dwLower && !phContext->dwUpper)
		return SEC_E_INVALID_HANDLE;

	NTLM_CONTEXT* context = (NTLM_CONTEXT*)sspi_SecureHandleGetLowerPointer(phContext);

	if (pInput)
	{
		input_buffer = sspi_FindSecBuffer(pInput, SECBUFFER_TOKEN);
	}

	if (!context)
	{
		context = ntlm_ContextNew();

		if (!context)
			return SEC_E_INSUFFICIENT_MEMORY;

		if (fContextReq & ISC_REQ_CONFIDENTIALITY)
			context->confidentiality = TRUE;

		credentials = (SSPI_CREDENTIALS*)sspi_SecureHandleGetLowerPointer(phCredential);
		context->credentials = credentials;

		if (context->Workstation.Length < 1)
		{
			if (ntlm_SetContextWorkstation(context, nullptr) < 0)
			{
				ntlm_ContextFree(context);
				return SEC_E_INTERNAL_ERROR;
			}
		}

		if (ntlm_SetContextServicePrincipalNameW(context, pszTargetName) < 0)
		{
			ntlm_ContextFree(context);
			return SEC_E_INTERNAL_ERROR;
		}

		sspi_SecureHandleSetLowerPointer(phNewContext, context);
		sspi_SecureHandleSetUpperPointer(phNewContext, NTLM_SSP_NAME);
	}

	if ((!input_buffer) || (ntlm_get_state(context) == NTLM_STATE_AUTHENTICATE))
	{
		if (!pOutput)
			return SEC_E_INVALID_TOKEN;

		if (pOutput->cBuffers < 1)
			return SEC_E_INVALID_TOKEN;

		output_buffer = sspi_FindSecBuffer(pOutput, SECBUFFER_TOKEN);

		if (!output_buffer)
			return SEC_E_INVALID_TOKEN;

		if (output_buffer->cbBuffer < 1)
			return SEC_E_INVALID_TOKEN;

		if (ntlm_get_state(context) == NTLM_STATE_INITIAL)
			ntlm_change_state(context, NTLM_STATE_NEGOTIATE);

		if (ntlm_get_state(context) == NTLM_STATE_NEGOTIATE)
			return ntlm_write_NegotiateMessage(context, output_buffer);

		return SEC_E_OUT_OF_SEQUENCE;
	}
	else
	{
		if (!input_buffer)
			return SEC_E_INVALID_TOKEN;

		if (input_buffer->cbBuffer < 1)
			return SEC_E_INVALID_TOKEN;

		PSecBuffer channel_bindings = sspi_FindSecBuffer(pInput, SECBUFFER_CHANNEL_BINDINGS);

		if (channel_bindings)
		{
			context->Bindings.BindingsLength = channel_bindings->cbBuffer;
			context->Bindings.Bindings = (SEC_CHANNEL_BINDINGS*)channel_bindings->pvBuffer;
		}

		if (ntlm_get_state(context) == NTLM_STATE_CHALLENGE)
		{
			status = ntlm_read_ChallengeMessage(context, input_buffer);

			if (status != SEC_I_CONTINUE_NEEDED)
				return status;

			if (!pOutput)
				return SEC_E_INVALID_TOKEN;

			if (pOutput->cBuffers < 1)
				return SEC_E_INVALID_TOKEN;

			output_buffer = sspi_FindSecBuffer(pOutput, SECBUFFER_TOKEN);

			if (!output_buffer)
				return SEC_E_INVALID_TOKEN;

			if (output_buffer->cbBuffer < 1)
				return SEC_E_INSUFFICIENT_MEMORY;

			if (ntlm_get_state(context) == NTLM_STATE_AUTHENTICATE)
				return ntlm_write_AuthenticateMessage(context, output_buffer);
		}

		return SEC_E_OUT_OF_SEQUENCE;
	}

	return SEC_E_OUT_OF_SEQUENCE;
}

/**
 * @see http://msdn.microsoft.com/en-us/library/windows/desktop/aa375512%28v=vs.85%29.aspx
 */
static SECURITY_STATUS SEC_ENTRY ntlm_InitializeSecurityContextA(
    PCredHandle phCredential, PCtxtHandle phContext, SEC_CHAR* pszTargetName, ULONG fContextReq,
    ULONG Reserved1, ULONG TargetDataRep, PSecBufferDesc pInput, ULONG Reserved2,
    PCtxtHandle phNewContext, PSecBufferDesc pOutput, PULONG pfContextAttr, PTimeStamp ptsExpiry)
{
	SECURITY_STATUS status = 0;
	SEC_WCHAR* pszTargetNameW = nullptr;

	if (pszTargetName)
	{
		pszTargetNameW = ConvertUtf8ToWCharAlloc(pszTargetName, nullptr);
		if (!pszTargetNameW)
			return SEC_E_INTERNAL_ERROR;
	}

	status = ntlm_InitializeSecurityContextW(phCredential, phContext, pszTargetNameW, fContextReq,
	                                         Reserved1, TargetDataRep, pInput, Reserved2,
	                                         phNewContext, pOutput, pfContextAttr, ptsExpiry);
	free(pszTargetNameW);
	return status;
}

/* http://msdn.microsoft.com/en-us/library/windows/desktop/aa375354 */

static SECURITY_STATUS SEC_ENTRY ntlm_DeleteSecurityContext(PCtxtHandle phContext)
{
	NTLM_CONTEXT* context = (NTLM_CONTEXT*)sspi_SecureHandleGetLowerPointer(phContext);
	ntlm_ContextFree(context);
	return SEC_E_OK;
}

SECURITY_STATUS ntlm_computeProofValue(NTLM_CONTEXT* ntlm, SecBuffer* ntproof)
{
	BYTE* blob = nullptr;
	SecBuffer* target = nullptr;

	WINPR_ASSERT(ntlm);
	WINPR_ASSERT(ntproof);

	target = &ntlm->ChallengeTargetInfo;

	if (!sspi_SecBufferAlloc(ntproof, 36 + target->cbBuffer))
		return SEC_E_INSUFFICIENT_MEMORY;

	blob = (BYTE*)ntproof->pvBuffer;
	CopyMemory(blob, ntlm->ServerChallenge, 8); /* Server challenge. */
	blob[8] = 1;                                /* Response version. */
	blob[9] = 1; /* Highest response version understood by the client. */
	/* Reserved 6B. */
	CopyMemory(&blob[16], ntlm->Timestamp, 8);       /* Time. */
	CopyMemory(&blob[24], ntlm->ClientChallenge, 8); /* Client challenge. */
	/* Reserved 4B. */
	/* Server name. */
	CopyMemory(&blob[36], target->pvBuffer, target->cbBuffer);
	return SEC_E_OK;
}

SECURITY_STATUS ntlm_computeMicValue(NTLM_CONTEXT* ntlm, SecBuffer* micvalue)
{
	BYTE* blob = nullptr;
	ULONG msgSize = 0;

	WINPR_ASSERT(ntlm);
	WINPR_ASSERT(micvalue);

	msgSize = ntlm->NegotiateMessage.cbBuffer + ntlm->ChallengeMessage.cbBuffer +
	          ntlm->AuthenticateMessage.cbBuffer;

	if (!sspi_SecBufferAlloc(micvalue, msgSize))
		return SEC_E_INSUFFICIENT_MEMORY;

	blob = (BYTE*)micvalue->pvBuffer;
	CopyMemory(blob, ntlm->NegotiateMessage.pvBuffer, ntlm->NegotiateMessage.cbBuffer);
	blob += ntlm->NegotiateMessage.cbBuffer;
	CopyMemory(blob, ntlm->ChallengeMessage.pvBuffer, ntlm->ChallengeMessage.cbBuffer);
	blob += ntlm->ChallengeMessage.cbBuffer;
	CopyMemory(blob, ntlm->AuthenticateMessage.pvBuffer, ntlm->AuthenticateMessage.cbBuffer);
	blob += ntlm->MessageIntegrityCheckOffset;
	ZeroMemory(blob, 16);
	return SEC_E_OK;
}

/* http://msdn.microsoft.com/en-us/library/windows/desktop/aa379337/ */

static SECURITY_STATUS SEC_ENTRY ntlm_QueryContextAttributesW(PCtxtHandle phContext,
                                                              ULONG ulAttribute, void* pBuffer)
{
	if (!phContext)
		return SEC_E_INVALID_HANDLE;

	if (!pBuffer)
		return SEC_E_INSUFFICIENT_MEMORY;

	NTLM_CONTEXT* context = (NTLM_CONTEXT*)sspi_SecureHandleGetLowerPointer(phContext);
	if (!check_context(context))
		return SEC_E_INVALID_HANDLE;

	if (ulAttribute == SECPKG_ATTR_SIZES)
	{
		SecPkgContext_Sizes* ContextSizes = (SecPkgContext_Sizes*)pBuffer;
		ContextSizes->cbMaxToken = 2010;
		ContextSizes->cbMaxSignature = 16;    /* the size of expected signature is 16 bytes */
		ContextSizes->cbBlockSize = 0;        /* no padding */
		ContextSizes->cbSecurityTrailer = 16; /* no security trailer appended in NTLM
		                            contrary to Kerberos */
		return SEC_E_OK;
	}
	else if (ulAttribute == SECPKG_ATTR_AUTH_IDENTITY)
	{
		SSPI_CREDENTIALS* credentials = nullptr;
		const SecPkgContext_AuthIdentity empty = WINPR_C_ARRAY_INIT;
		SecPkgContext_AuthIdentity* AuthIdentity = (SecPkgContext_AuthIdentity*)pBuffer;

		WINPR_ASSERT(AuthIdentity);
		*AuthIdentity = empty;

		context->UseSamFileDatabase = FALSE;
		credentials = context->credentials;

		if (credentials->identity.UserLength > 0)
		{
			if (ConvertWCharNToUtf8(credentials->identity.User, credentials->identity.UserLength,
			                        AuthIdentity->User, ARRAYSIZE(AuthIdentity->User)) <= 0)
				return SEC_E_INTERNAL_ERROR;
		}

		if (credentials->identity.DomainLength > 0)
		{
			if (ConvertWCharNToUtf8(credentials->identity.Domain,
			                        credentials->identity.DomainLength, AuthIdentity->Domain,
			                        ARRAYSIZE(AuthIdentity->Domain)) <= 0)
				return SEC_E_INTERNAL_ERROR;
		}

		return SEC_E_OK;
	}
	else if (ulAttribute == SECPKG_ATTR_AUTH_NTLM_NTPROOF_VALUE)
	{
		return ntlm_computeProofValue(context, (SecBuffer*)pBuffer);
	}
	else if (ulAttribute == SECPKG_ATTR_AUTH_NTLM_RANDKEY)
	{
		SecBuffer* randkey = nullptr;
		randkey = (SecBuffer*)pBuffer;

		if (!sspi_SecBufferAlloc(randkey, 16))
			return (SEC_E_INSUFFICIENT_MEMORY);

		CopyMemory(randkey->pvBuffer, context->EncryptedRandomSessionKey, 16);
		return (SEC_E_OK);
	}
	else if (ulAttribute == SECPKG_ATTR_AUTH_NTLM_MIC)
	{
		SecBuffer* mic = (SecBuffer*)pBuffer;
		NTLM_AUTHENTICATE_MESSAGE* message = &context->AUTHENTICATE_MESSAGE;

		if (!sspi_SecBufferAlloc(mic, 16))
			return (SEC_E_INSUFFICIENT_MEMORY);

		CopyMemory(mic->pvBuffer, message->MessageIntegrityCheck, 16);
		return (SEC_E_OK);
	}
	else if (ulAttribute == SECPKG_ATTR_AUTH_NTLM_MIC_VALUE)
	{
		return ntlm_computeMicValue(context, (SecBuffer*)pBuffer);
	}
	else if (ulAttribute == SECPKG_ATTR_PACKAGE_INFO)
	{
		SecPkgContext_PackageInfo* PackageInfo = (SecPkgContext_PackageInfo*)pBuffer;
		size_t size = sizeof(SecPkgInfoA);
		SecPkgInfoA* pPackageInfo =
		    (SecPkgInfoA*)sspi_ContextBufferAlloc(QuerySecurityPackageInfoIndex, size);

		if (!pPackageInfo)
			return SEC_E_INSUFFICIENT_MEMORY;

		pPackageInfo->fCapabilities = NTLM_SecPkgInfoA.fCapabilities;
		pPackageInfo->wVersion = NTLM_SecPkgInfoA.wVersion;
		pPackageInfo->wRPCID = NTLM_SecPkgInfoA.wRPCID;
		pPackageInfo->cbMaxToken = NTLM_SecPkgInfoA.cbMaxToken;
		pPackageInfo->Name = _strdup(NTLM_SecPkgInfoA.Name);
		pPackageInfo->Comment = _strdup(NTLM_SecPkgInfoA.Comment);

		if (!pPackageInfo->Name || !pPackageInfo->Comment)
		{
			sspi_ContextBufferFree(pPackageInfo);
			return SEC_E_INSUFFICIENT_MEMORY;
		}
		PackageInfo->PackageInfo = pPackageInfo;
		return SEC_E_OK;
	}

	WLog_ERR(TAG, "TODO: Implement ulAttribute=0x%08" PRIx32, ulAttribute);
	return SEC_E_UNSUPPORTED_FUNCTION;
}

static SECURITY_STATUS SEC_ENTRY ntlm_QueryContextAttributesA(PCtxtHandle phContext,
                                                              ULONG ulAttribute, void* pBuffer)
{
	return ntlm_QueryContextAttributesW(phContext, ulAttribute, pBuffer);
}

static SECURITY_STATUS SEC_ENTRY ntlm_SetContextAttributesW(PCtxtHandle phContext,
                                                            ULONG ulAttribute, void* pBuffer,
                                                            ULONG cbBuffer)
{
	if (!phContext)
		return SEC_E_INVALID_HANDLE;

	if (!pBuffer)
		return SEC_E_INVALID_PARAMETER;

	NTLM_CONTEXT* context = (NTLM_CONTEXT*)sspi_SecureHandleGetLowerPointer(phContext);
	if (!context)
		return SEC_E_INVALID_HANDLE;

	if (ulAttribute == SECPKG_ATTR_AUTH_NTLM_HASH)
	{
		SecPkgContext_AuthNtlmHash* AuthNtlmHash = (SecPkgContext_AuthNtlmHash*)pBuffer;

		if (cbBuffer < sizeof(SecPkgContext_AuthNtlmHash))
			return SEC_E_INVALID_PARAMETER;

		if (AuthNtlmHash->Version == 1)
			CopyMemory(context->NtlmHash, AuthNtlmHash->NtlmHash, 16);
		else if (AuthNtlmHash->Version == 2)
			CopyMemory(context->NtlmV2Hash, AuthNtlmHash->NtlmHash, 16);

		return SEC_E_OK;
	}
	else if (ulAttribute == SECPKG_ATTR_AUTH_NTLM_MESSAGE)
	{
		SecPkgContext_AuthNtlmMessage* AuthNtlmMessage = (SecPkgContext_AuthNtlmMessage*)pBuffer;

		if (cbBuffer < sizeof(SecPkgContext_AuthNtlmMessage))
			return SEC_E_INVALID_PARAMETER;

		if (AuthNtlmMessage->type == 1)
		{
			sspi_SecBufferFree(&context->NegotiateMessage);

			if (!sspi_SecBufferAlloc(&context->NegotiateMessage, AuthNtlmMessage->length))
				return SEC_E_INSUFFICIENT_MEMORY;

			CopyMemory(context->NegotiateMessage.pvBuffer, AuthNtlmMessage->buffer,
			           AuthNtlmMessage->length);
		}
		else if (AuthNtlmMessage->type == 2)
		{
			sspi_SecBufferFree(&context->ChallengeMessage);

			if (!sspi_SecBufferAlloc(&context->ChallengeMessage, AuthNtlmMessage->length))
				return SEC_E_INSUFFICIENT_MEMORY;

			CopyMemory(context->ChallengeMessage.pvBuffer, AuthNtlmMessage->buffer,
			           AuthNtlmMessage->length);
		}
		else if (AuthNtlmMessage->type == 3)
		{
			sspi_SecBufferFree(&context->AuthenticateMessage);

			if (!sspi_SecBufferAlloc(&context->AuthenticateMessage, AuthNtlmMessage->length))
				return SEC_E_INSUFFICIENT_MEMORY;

			CopyMemory(context->AuthenticateMessage.pvBuffer, AuthNtlmMessage->buffer,
			           AuthNtlmMessage->length);
		}

		return SEC_E_OK;
	}
	else if (ulAttribute == SECPKG_ATTR_AUTH_NTLM_TIMESTAMP)
	{
		SecPkgContext_AuthNtlmTimestamp* AuthNtlmTimestamp =
		    (SecPkgContext_AuthNtlmTimestamp*)pBuffer;

		if (cbBuffer < sizeof(SecPkgContext_AuthNtlmTimestamp))
			return SEC_E_INVALID_PARAMETER;

		if (AuthNtlmTimestamp->ChallengeOrResponse)
			CopyMemory(context->ChallengeTimestamp, AuthNtlmTimestamp->Timestamp, 8);
		else
			CopyMemory(context->Timestamp, AuthNtlmTimestamp->Timestamp, 8);

		return SEC_E_OK;
	}
	else if (ulAttribute == SECPKG_ATTR_AUTH_NTLM_CLIENT_CHALLENGE)
	{
		SecPkgContext_AuthNtlmClientChallenge* AuthNtlmClientChallenge =
		    (SecPkgContext_AuthNtlmClientChallenge*)pBuffer;

		if (cbBuffer < sizeof(SecPkgContext_AuthNtlmClientChallenge))
			return SEC_E_INVALID_PARAMETER;

		CopyMemory(context->ClientChallenge, AuthNtlmClientChallenge->ClientChallenge, 8);
		return SEC_E_OK;
	}
	else if (ulAttribute == SECPKG_ATTR_AUTH_NTLM_SERVER_CHALLENGE)
	{
		SecPkgContext_AuthNtlmServerChallenge* AuthNtlmServerChallenge =
		    (SecPkgContext_AuthNtlmServerChallenge*)pBuffer;

		if (cbBuffer < sizeof(SecPkgContext_AuthNtlmServerChallenge))
			return SEC_E_INVALID_PARAMETER;

		CopyMemory(context->ServerChallenge, AuthNtlmServerChallenge->ServerChallenge, 8);
		return SEC_E_OK;
	}

	WLog_ERR(TAG, "TODO: Implement ulAttribute=%08" PRIx32, ulAttribute);
	return SEC_E_UNSUPPORTED_FUNCTION;
}

static SECURITY_STATUS SEC_ENTRY ntlm_SetContextAttributesA(PCtxtHandle phContext,
                                                            ULONG ulAttribute, void* pBuffer,
                                                            ULONG cbBuffer)
{
	return ntlm_SetContextAttributesW(phContext, ulAttribute, pBuffer, cbBuffer);
}

static SECURITY_STATUS SEC_ENTRY ntlm_SetCredentialsAttributesW(
    WINPR_ATTR_UNUSED PCredHandle phCredential, WINPR_ATTR_UNUSED ULONG ulAttribute,
    WINPR_ATTR_UNUSED void* pBuffer, WINPR_ATTR_UNUSED ULONG cbBuffer)
{
	return SEC_E_UNSUPPORTED_FUNCTION;
}

static SECURITY_STATUS SEC_ENTRY ntlm_SetCredentialsAttributesA(
    WINPR_ATTR_UNUSED PCredHandle phCredential, WINPR_ATTR_UNUSED ULONG ulAttribute,
    WINPR_ATTR_UNUSED void* pBuffer, WINPR_ATTR_UNUSED ULONG cbBuffer)
{
	return SEC_E_UNSUPPORTED_FUNCTION;
}

static SECURITY_STATUS SEC_ENTRY ntlm_RevertSecurityContext(WINPR_ATTR_UNUSED PCtxtHandle phContext)
{
	return SEC_E_OK;
}

static SECURITY_STATUS SEC_ENTRY ntlm_EncryptMessage(PCtxtHandle phContext,
                                                     WINPR_ATTR_UNUSED ULONG fQOP,
                                                     PSecBufferDesc pMessage, ULONG MessageSeqNo)
{
	const UINT32 SeqNo = MessageSeqNo;
	UINT32 value = 0;
	BYTE digest[WINPR_MD5_DIGEST_LENGTH] = WINPR_C_ARRAY_INIT;
	BYTE checksum[8] = WINPR_C_ARRAY_INIT;
	ULONG version = 1;
	PSecBuffer data_buffer = nullptr;
	PSecBuffer signature_buffer = nullptr;
	NTLM_CONTEXT* context = (NTLM_CONTEXT*)sspi_SecureHandleGetLowerPointer(phContext);
	if (!check_context(context))
		return SEC_E_INVALID_HANDLE;

	for (ULONG index = 0; index < pMessage->cBuffers; index++)
	{
		SecBuffer* cur = &pMessage->pBuffers[index];

		if (cur->BufferType & SECBUFFER_DATA)
			data_buffer = cur;
		else if (cur->BufferType & SECBUFFER_TOKEN)
			signature_buffer = cur;
	}

	if (!data_buffer)
		return SEC_E_INVALID_TOKEN;

	if (!signature_buffer)
		return SEC_E_INVALID_TOKEN;

	/* Copy original data buffer */
	ULONG length = data_buffer->cbBuffer;
	void* data = malloc(length);

	if (!data)
		return SEC_E_INSUFFICIENT_MEMORY;

	CopyMemory(data, data_buffer->pvBuffer, length);
	/* Compute the HMAC-MD5 hash of ConcatenationOf(seq_num,data) using the client signing key */
	WINPR_HMAC_CTX* hmac = winpr_HMAC_New();

	BOOL success = FALSE;
	{
		if (!hmac)
			goto hmac_fail;
		if (!winpr_HMAC_Init(hmac, WINPR_MD_MD5, context->SendSigningKey, WINPR_MD5_DIGEST_LENGTH))
			goto hmac_fail;

		winpr_Data_Write_UINT32(&value, SeqNo);

		if (!winpr_HMAC_Update(hmac, (void*)&value, 4))
			goto hmac_fail;
		if (!winpr_HMAC_Update(hmac, data, length))
			goto hmac_fail;
		if (!winpr_HMAC_Final(hmac, digest, WINPR_MD5_DIGEST_LENGTH))
			goto hmac_fail;
	}

	success = TRUE;

hmac_fail:
	winpr_HMAC_Free(hmac);
	if (!success)
	{
		free(data);
		return SEC_E_INSUFFICIENT_MEMORY;
	}

	/* Encrypt message using with RC4, result overwrites original buffer */
	if ((data_buffer->BufferType & SECBUFFER_READONLY) == 0)
	{
		if (context->confidentiality)
		{
			if (!winpr_RC4_Update(context->SendRc4Seal, length, (BYTE*)data,
			                      (BYTE*)data_buffer->pvBuffer))
			{
				free(data);
				return SEC_E_INSUFFICIENT_MEMORY;
			}
		}
		else
			CopyMemory(data_buffer->pvBuffer, data, length);
	}

#ifdef WITH_DEBUG_NTLM
	WLog_DBG(TAG, "Data Buffer (length = %" PRIu32 ")", length);
	winpr_HexDump(TAG, WLOG_DEBUG, data, length);
	WLog_DBG(TAG, "Encrypted Data Buffer (length = %" PRIu32 ")", data_buffer->cbBuffer);
	winpr_HexDump(TAG, WLOG_DEBUG, data_buffer->pvBuffer, data_buffer->cbBuffer);
#endif
	free(data);
	/* RC4-encrypt first 8 bytes of digest */
	if (!winpr_RC4_Update(context->SendRc4Seal, 8, digest, checksum))
		return SEC_E_INSUFFICIENT_MEMORY;
	if ((signature_buffer->BufferType & SECBUFFER_READONLY) == 0)
	{
		BYTE* signature = signature_buffer->pvBuffer;
		/* Concatenate version, ciphertext and sequence number to build signature */
		winpr_Data_Write_UINT32(signature, version);
		CopyMemory(&signature[4], (void*)checksum, 8);
		winpr_Data_Write_UINT32(&signature[12], SeqNo);
	}
	context->SendSeqNum++;
#ifdef WITH_DEBUG_NTLM
	WLog_DBG(TAG, "Signature (length = %" PRIu32 ")", signature_buffer->cbBuffer);
	winpr_HexDump(TAG, WLOG_DEBUG, signature_buffer->pvBuffer, signature_buffer->cbBuffer);
#endif
	return SEC_E_OK;
}

static SECURITY_STATUS SEC_ENTRY ntlm_DecryptMessage(PCtxtHandle phContext, PSecBufferDesc pMessage,
                                                     ULONG MessageSeqNo,
                                                     WINPR_ATTR_UNUSED PULONG pfQOP)
{
	const UINT32 SeqNo = (UINT32)MessageSeqNo;
	UINT32 value = 0;
	BYTE digest[WINPR_MD5_DIGEST_LENGTH] = WINPR_C_ARRAY_INIT;
	BYTE checksum[8] = WINPR_C_ARRAY_INIT;
	UINT32 version = 1;
	BYTE expected_signature[WINPR_MD5_DIGEST_LENGTH] = WINPR_C_ARRAY_INIT;
	PSecBuffer data_buffer = nullptr;
	PSecBuffer signature_buffer = nullptr;
	NTLM_CONTEXT* context = (NTLM_CONTEXT*)sspi_SecureHandleGetLowerPointer(phContext);
	if (!check_context(context))
		return SEC_E_INVALID_HANDLE;

	for (ULONG index = 0; index < pMessage->cBuffers; index++)
	{
		if (pMessage->pBuffers[index].BufferType == SECBUFFER_DATA)
			data_buffer = &pMessage->pBuffers[index];
		else if (pMessage->pBuffers[index].BufferType == SECBUFFER_TOKEN)
			signature_buffer = &pMessage->pBuffers[index];
	}

	if (!data_buffer)
		return SEC_E_INVALID_TOKEN;

	if (!signature_buffer)
		return SEC_E_INVALID_TOKEN;

	/* Copy original data buffer */
	const ULONG length = data_buffer->cbBuffer;
	void* data = malloc(length);

	if (!data)
		return SEC_E_INSUFFICIENT_MEMORY;

	CopyMemory(data, data_buffer->pvBuffer, length);

	/* Decrypt message using with RC4, result overwrites original buffer */

	if (context->confidentiality)
	{
		if (!winpr_RC4_Update(context->RecvRc4Seal, length, (BYTE*)data,
		                      (BYTE*)data_buffer->pvBuffer))
		{
			free(data);
			return SEC_E_INSUFFICIENT_MEMORY;
		}
	}
	else
		CopyMemory(data_buffer->pvBuffer, data, length);

	/* Compute the HMAC-MD5 hash of ConcatenationOf(seq_num,data) using the client signing key */
	WINPR_HMAC_CTX* hmac = winpr_HMAC_New();

	BOOL success = FALSE;
	{
		if (!hmac)
			goto hmac_fail;

		if (!winpr_HMAC_Init(hmac, WINPR_MD_MD5, context->RecvSigningKey, WINPR_MD5_DIGEST_LENGTH))
			goto hmac_fail;

		winpr_Data_Write_UINT32(&value, SeqNo);

		if (!winpr_HMAC_Update(hmac, (void*)&value, 4))
			goto hmac_fail;
		if (!winpr_HMAC_Update(hmac, data_buffer->pvBuffer, data_buffer->cbBuffer))
			goto hmac_fail;
		if (!winpr_HMAC_Final(hmac, digest, WINPR_MD5_DIGEST_LENGTH))
			goto hmac_fail;

		success = TRUE;
	}
hmac_fail:
	winpr_HMAC_Free(hmac);
	if (!success)
	{
		free(data);
		return SEC_E_INSUFFICIENT_MEMORY;
	}

#ifdef WITH_DEBUG_NTLM
	WLog_DBG(TAG, "Encrypted Data Buffer (length = %" PRIu32 ")", length);
	winpr_HexDump(TAG, WLOG_DEBUG, data, length);
	WLog_DBG(TAG, "Data Buffer (length = %" PRIu32 ")", data_buffer->cbBuffer);
	winpr_HexDump(TAG, WLOG_DEBUG, data_buffer->pvBuffer, data_buffer->cbBuffer);
#endif
	free(data);
	/* RC4-encrypt first 8 bytes of digest */
	if (!winpr_RC4_Update(context->RecvRc4Seal, 8, digest, checksum))
		return SEC_E_MESSAGE_ALTERED;

	/* Concatenate version, ciphertext and sequence number to build signature */
	winpr_Data_Write_UINT32(expected_signature, version);
	CopyMemory(&expected_signature[4], (void*)checksum, 8);
	winpr_Data_Write_UINT32(&expected_signature[12], SeqNo);
	context->RecvSeqNum++;

	if (memcmp(signature_buffer->pvBuffer, expected_signature, 16) != 0)
	{
		/* signature verification failed! */
		WLog_ERR(TAG, "signature verification failed, something nasty is going on!");
#ifdef WITH_DEBUG_NTLM
		WLog_ERR(TAG, "Expected Signature:");
		winpr_HexDump(TAG, WLOG_ERROR, expected_signature, 16);
		WLog_ERR(TAG, "Actual Signature:");
		winpr_HexDump(TAG, WLOG_ERROR, (BYTE*)signature_buffer->pvBuffer, 16);
#endif
		return SEC_E_MESSAGE_ALTERED;
	}

	return SEC_E_OK;
}

static SECURITY_STATUS SEC_ENTRY ntlm_MakeSignature(PCtxtHandle phContext,
                                                    WINPR_ATTR_UNUSED ULONG fQOP,
                                                    PSecBufferDesc pMessage, ULONG MessageSeqNo)
{
	SECURITY_STATUS status = SEC_E_INTERNAL_ERROR;
	PSecBuffer data_buffer = nullptr;
	PSecBuffer sig_buffer = nullptr;
	UINT32 seq_no = 0;
	BYTE digest[WINPR_MD5_DIGEST_LENGTH] = WINPR_C_ARRAY_INIT;
	BYTE checksum[8] = WINPR_C_ARRAY_INIT;

	NTLM_CONTEXT* context = sspi_SecureHandleGetLowerPointer(phContext);
	if (!check_context(context))
		return SEC_E_INVALID_HANDLE;

	for (ULONG i = 0; i < pMessage->cBuffers; i++)
	{
		if (pMessage->pBuffers[i].BufferType == SECBUFFER_DATA)
			data_buffer = &pMessage->pBuffers[i];
		else if (pMessage->pBuffers[i].BufferType == SECBUFFER_TOKEN)
			sig_buffer = &pMessage->pBuffers[i];
	}

	if (!data_buffer || !sig_buffer)
		return SEC_E_INVALID_TOKEN;

	WINPR_HMAC_CTX* hmac = winpr_HMAC_New();

	if (!winpr_HMAC_Init(hmac, WINPR_MD_MD5, context->SendSigningKey, WINPR_MD5_DIGEST_LENGTH))
		goto fail;

	winpr_Data_Write_UINT32(&seq_no, MessageSeqNo);
	if (!winpr_HMAC_Update(hmac, (BYTE*)&seq_no, 4))
		goto fail;
	if (!winpr_HMAC_Update(hmac, data_buffer->pvBuffer, data_buffer->cbBuffer))
		goto fail;
	if (!winpr_HMAC_Final(hmac, digest, WINPR_MD5_DIGEST_LENGTH))
		goto fail;

	if (!winpr_RC4_Update(context->SendRc4Seal, 8, digest, checksum))
		goto fail;

	BYTE* signature = sig_buffer->pvBuffer;
	winpr_Data_Write_UINT32(signature, 1L);
	CopyMemory(&signature[4], checksum, 8);
	winpr_Data_Write_UINT32(&signature[12], seq_no);
	sig_buffer->cbBuffer = 16;

	status = SEC_E_OK;

fail:
	winpr_HMAC_Free(hmac);
	return status;
}

static SECURITY_STATUS SEC_ENTRY ntlm_VerifySignature(PCtxtHandle phContext,
                                                      PSecBufferDesc pMessage, ULONG MessageSeqNo,
                                                      WINPR_ATTR_UNUSED PULONG pfQOP)
{
	SECURITY_STATUS status = SEC_E_INTERNAL_ERROR;
	PSecBuffer data_buffer = nullptr;
	PSecBuffer sig_buffer = nullptr;
	UINT32 seq_no = 0;
	BYTE digest[WINPR_MD5_DIGEST_LENGTH] = WINPR_C_ARRAY_INIT;
	BYTE checksum[8] = WINPR_C_ARRAY_INIT;
	BYTE signature[16] = WINPR_C_ARRAY_INIT;

	NTLM_CONTEXT* context = sspi_SecureHandleGetLowerPointer(phContext);
	if (!check_context(context))
		return SEC_E_INVALID_HANDLE;

	for (ULONG i = 0; i < pMessage->cBuffers; i++)
	{
		if (pMessage->pBuffers[i].BufferType == SECBUFFER_DATA)
			data_buffer = &pMessage->pBuffers[i];
		else if (pMessage->pBuffers[i].BufferType == SECBUFFER_TOKEN)
			sig_buffer = &pMessage->pBuffers[i];
	}

	if (!data_buffer || !sig_buffer)
		return SEC_E_INVALID_TOKEN;

	WINPR_HMAC_CTX* hmac = winpr_HMAC_New();

	if (!winpr_HMAC_Init(hmac, WINPR_MD_MD5, context->RecvSigningKey, WINPR_MD5_DIGEST_LENGTH))
		goto fail;

	winpr_Data_Write_UINT32(&seq_no, MessageSeqNo);
	if (!winpr_HMAC_Update(hmac, (BYTE*)&seq_no, 4))
		goto fail;
	if (!winpr_HMAC_Update(hmac, data_buffer->pvBuffer, data_buffer->cbBuffer))
		goto fail;
	if (!winpr_HMAC_Final(hmac, digest, WINPR_MD5_DIGEST_LENGTH))
		goto fail;

	if (!winpr_RC4_Update(context->RecvRc4Seal, 8, digest, checksum))
		goto fail;

	winpr_Data_Write_UINT32(signature, 1L);
	CopyMemory(&signature[4], checksum, 8);
	winpr_Data_Write_UINT32(&signature[12], seq_no);

	status = SEC_E_OK;
	if (memcmp(sig_buffer->pvBuffer, signature, 16) != 0)
		status = SEC_E_MESSAGE_ALTERED;

fail:
	winpr_HMAC_Free(hmac);
	return status;
}

const SecurityFunctionTableA NTLM_SecurityFunctionTableA = {
	3,                                /* dwVersion */
	nullptr,                          /* EnumerateSecurityPackages */
	ntlm_QueryCredentialsAttributesA, /* QueryCredentialsAttributes */
	ntlm_AcquireCredentialsHandleA,   /* AcquireCredentialsHandle */
	ntlm_FreeCredentialsHandle,       /* FreeCredentialsHandle */
	nullptr,                          /* Reserved2 */
	ntlm_InitializeSecurityContextA,  /* InitializeSecurityContext */
	ntlm_AcceptSecurityContext,       /* AcceptSecurityContext */
	nullptr,                          /* CompleteAuthToken */
	ntlm_DeleteSecurityContext,       /* DeleteSecurityContext */
	nullptr,                          /* ApplyControlToken */
	ntlm_QueryContextAttributesA,     /* QueryContextAttributes */
	ntlm_ImpersonateSecurityContext,  /* ImpersonateSecurityContext */
	ntlm_RevertSecurityContext,       /* RevertSecurityContext */
	ntlm_MakeSignature,               /* MakeSignature */
	ntlm_VerifySignature,             /* VerifySignature */
	nullptr,                          /* FreeContextBuffer */
	nullptr,                          /* QuerySecurityPackageInfo */
	nullptr,                          /* Reserved3 */
	nullptr,                          /* Reserved4 */
	nullptr,                          /* ExportSecurityContext */
	nullptr,                          /* ImportSecurityContext */
	nullptr,                          /* AddCredentials */
	nullptr,                          /* Reserved8 */
	nullptr,                          /* QuerySecurityContextToken */
	ntlm_EncryptMessage,              /* EncryptMessage */
	ntlm_DecryptMessage,              /* DecryptMessage */
	ntlm_SetContextAttributesA,       /* SetContextAttributes */
	ntlm_SetCredentialsAttributesA,   /* SetCredentialsAttributes */
};

const SecurityFunctionTableW NTLM_SecurityFunctionTableW = {
	3,                                /* dwVersion */
	nullptr,                          /* EnumerateSecurityPackages */
	ntlm_QueryCredentialsAttributesW, /* QueryCredentialsAttributes */
	ntlm_AcquireCredentialsHandleW,   /* AcquireCredentialsHandle */
	ntlm_FreeCredentialsHandle,       /* FreeCredentialsHandle */
	nullptr,                          /* Reserved2 */
	ntlm_InitializeSecurityContextW,  /* InitializeSecurityContext */
	ntlm_AcceptSecurityContext,       /* AcceptSecurityContext */
	nullptr,                          /* CompleteAuthToken */
	ntlm_DeleteSecurityContext,       /* DeleteSecurityContext */
	nullptr,                          /* ApplyControlToken */
	ntlm_QueryContextAttributesW,     /* QueryContextAttributes */
	ntlm_ImpersonateSecurityContext,  /* ImpersonateSecurityContext */
	ntlm_RevertSecurityContext,       /* RevertSecurityContext */
	ntlm_MakeSignature,               /* MakeSignature */
	ntlm_VerifySignature,             /* VerifySignature */
	nullptr,                          /* FreeContextBuffer */
	nullptr,                          /* QuerySecurityPackageInfo */
	nullptr,                          /* Reserved3 */
	nullptr,                          /* Reserved4 */
	nullptr,                          /* ExportSecurityContext */
	nullptr,                          /* ImportSecurityContext */
	nullptr,                          /* AddCredentials */
	nullptr,                          /* Reserved8 */
	nullptr,                          /* QuerySecurityContextToken */
	ntlm_EncryptMessage,              /* EncryptMessage */
	ntlm_DecryptMessage,              /* DecryptMessage */
	ntlm_SetContextAttributesW,       /* SetContextAttributes */
	ntlm_SetCredentialsAttributesW,   /* SetCredentialsAttributes */
};

const SecPkgInfoA NTLM_SecPkgInfoA = {
	0x00082B37,             /* fCapabilities */
	1,                      /* wVersion */
	0x000A,                 /* wRPCID */
	0x00000B48,             /* cbMaxToken */
	"NTLM",                 /* Name */
	"NTLM Security Package" /* Comment */
};

static WCHAR NTLM_SecPkgInfoW_NameBuffer[32] = WINPR_C_ARRAY_INIT;
static WCHAR NTLM_SecPkgInfoW_CommentBuffer[32] = WINPR_C_ARRAY_INIT;

const SecPkgInfoW NTLM_SecPkgInfoW = {
	0x00082B37,                    /* fCapabilities */
	1,                             /* wVersion */
	0x000A,                        /* wRPCID */
	0x00000B48,                    /* cbMaxToken */
	NTLM_SecPkgInfoW_NameBuffer,   /* Name */
	NTLM_SecPkgInfoW_CommentBuffer /* Comment */
};

char* ntlm_negotiate_flags_string(char* buffer, size_t size, UINT32 flags)
{
	if (!buffer || (size == 0))
		return buffer;

	(void)_snprintf(buffer, size, "[0x%08" PRIx32 "] ", flags);

	for (int x = 0; x < 31; x++)
	{
		const UINT32 mask = 1u << x;
		size_t len = strnlen(buffer, size);
		if (flags & mask)
		{
			const char* str = ntlm_get_negotiate_string(mask);
			const size_t flen = strlen(str);

			if ((len > 0) && (buffer[len - 1] != ' '))
			{
				if (size - len < 1)
					break;
				winpr_str_append("|", buffer, size, nullptr);
				len++;
			}

			if (size - len < flen)
				break;
			winpr_str_append(str, buffer, size, nullptr);
		}
	}

	return buffer;
}

const char* ntlm_message_type_string(UINT32 messageType)
{
	switch (messageType)
	{
		case MESSAGE_TYPE_NEGOTIATE:
			return "MESSAGE_TYPE_NEGOTIATE";
		case MESSAGE_TYPE_CHALLENGE:
			return "MESSAGE_TYPE_CHALLENGE";
		case MESSAGE_TYPE_AUTHENTICATE:
			return "MESSAGE_TYPE_AUTHENTICATE";
		default:
			return "MESSAGE_TYPE_UNKNOWN";
	}
}

const char* ntlm_state_string(NTLM_STATE state)
{
	switch (state)
	{
		case NTLM_STATE_INITIAL:
			return "NTLM_STATE_INITIAL";
		case NTLM_STATE_NEGOTIATE:
			return "NTLM_STATE_NEGOTIATE";
		case NTLM_STATE_CHALLENGE:
			return "NTLM_STATE_CHALLENGE";
		case NTLM_STATE_AUTHENTICATE:
			return "NTLM_STATE_AUTHENTICATE";
		case NTLM_STATE_FINAL:
			return "NTLM_STATE_FINAL";
		default:
			return "NTLM_STATE_UNKNOWN";
	}
}
void ntlm_change_state(NTLM_CONTEXT* ntlm, NTLM_STATE state)
{
	WINPR_ASSERT(ntlm);
	WLog_DBG(TAG, "change state from %s to %s", ntlm_state_string(ntlm->state),
	         ntlm_state_string(state));
	ntlm->state = state;
}

NTLM_STATE ntlm_get_state(NTLM_CONTEXT* ntlm)
{
	WINPR_ASSERT(ntlm);
	return ntlm->state;
}

BOOL ntlm_reset_cipher_state(PSecHandle phContext)
{
	NTLM_CONTEXT* context = sspi_SecureHandleGetLowerPointer(phContext);

	if (context)
	{
		check_context(context);
		winpr_RC4_Free(context->SendRc4Seal);
		winpr_RC4_Free(context->RecvRc4Seal);
		context->SendRc4Seal = winpr_RC4_New(context->RecvSealingKey, 16);
		context->RecvRc4Seal = winpr_RC4_New(context->SendSealingKey, 16);

		if (!context->SendRc4Seal)
		{
			WLog_ERR(TAG, "Failed to allocate context->SendRc4Seal");
			return FALSE;
		}
		if (!context->RecvRc4Seal)
		{
			WLog_ERR(TAG, "Failed to allocate context->RecvRc4Seal");
			return FALSE;
		}
	}

	return TRUE;
}

BOOL NTLM_init(void)
{
	InitializeConstWCharFromUtf8(NTLM_SecPkgInfoA.Name, NTLM_SecPkgInfoW_NameBuffer,
	                             ARRAYSIZE(NTLM_SecPkgInfoW_NameBuffer));
	InitializeConstWCharFromUtf8(NTLM_SecPkgInfoA.Comment, NTLM_SecPkgInfoW_CommentBuffer,
	                             ARRAYSIZE(NTLM_SecPkgInfoW_CommentBuffer));

	return TRUE;
}
