// Copyright Epic Games, Inc. All Rights Reserved.

#include <zenutil/cloud/sigv4.h>

#include <zencore/string.h>
#include <zencore/testing.h>

#include <algorithm>
#include <chrono>
#include <cstring>
#include <ctime>

// Platform-specific crypto backends
#if ZEN_PLATFORM_WINDOWS
#	define ZEN_S3_USE_BCRYPT 1
#else
#	define ZEN_S3_USE_BCRYPT 0
#endif

#ifndef ZEN_S3_USE_OPENSSL
#	if ZEN_S3_USE_BCRYPT
#		define ZEN_S3_USE_OPENSSL 0
#	else
#		define ZEN_S3_USE_OPENSSL 1
#	endif
#endif

ZEN_THIRD_PARTY_INCLUDES_START
#include <fmt/format.h>

#if ZEN_S3_USE_OPENSSL
#	include <openssl/evp.h>
#elif ZEN_S3_USE_BCRYPT
#	include <zencore/windows.h>
#	include <bcrypt.h>
#endif
ZEN_THIRD_PARTY_INCLUDES_END

namespace zen {

//////////////////////////////////////////////////////////////////////////
// SHA-256

#if ZEN_S3_USE_OPENSSL

Sha256Digest
ComputeSha256(const void* Data, size_t Size)
{
	Sha256Digest Result;
	unsigned int Len = 0;
	EVP_Digest(Data, Size, Result.data(), &Len, EVP_sha256(), nullptr);
	ZEN_ASSERT(Len == 32);
	return Result;
}

Sha256Digest
ComputeHmacSha256(const void* Key, size_t KeySize, const void* Data, size_t DataSize)
{
	Sha256Digest Result;

	EVP_MAC* Mac = EVP_MAC_fetch(nullptr, "HMAC", nullptr);
	ZEN_ASSERT(Mac != nullptr);

	EVP_MAC_CTX* Ctx = EVP_MAC_CTX_new(Mac);
	ZEN_ASSERT(Ctx != nullptr);

	OSSL_PARAM Params[] = {
		OSSL_PARAM_construct_utf8_string("digest", const_cast<char*>("SHA256"), 0),
		OSSL_PARAM_construct_end(),
	};

	int Rc = EVP_MAC_init(Ctx, reinterpret_cast<const unsigned char*>(Key), KeySize, Params);
	ZEN_ASSERT(Rc == 1);

	Rc = EVP_MAC_update(Ctx, reinterpret_cast<const unsigned char*>(Data), DataSize);
	ZEN_ASSERT(Rc == 1);

	size_t OutLen = 0;
	Rc			  = EVP_MAC_final(Ctx, Result.data(), &OutLen, Result.size());
	ZEN_ASSERT(Rc == 1);
	ZEN_ASSERT(OutLen == 32);

	EVP_MAC_CTX_free(Ctx);
	EVP_MAC_free(Mac);

	return Result;
}

#elif ZEN_S3_USE_BCRYPT

namespace {

#	define NT_SUCCESS(Status) (((NTSTATUS)(Status)) >= 0)

	Sha256Digest BcryptHash(BCRYPT_ALG_HANDLE Algorithm, const void* Data, size_t DataSize)
	{
		Sha256Digest	   Result;
		BCRYPT_HASH_HANDLE HashHandle = nullptr;
		NTSTATUS		   Status;

		Status = BCryptCreateHash(Algorithm, &HashHandle, nullptr, 0, nullptr, 0, 0);
		ZEN_ASSERT(NT_SUCCESS(Status));

		Status = BCryptHashData(HashHandle, (PUCHAR)Data, (ULONG)DataSize, 0);
		ZEN_ASSERT(NT_SUCCESS(Status));

		Status = BCryptFinishHash(HashHandle, Result.data(), (ULONG)Result.size(), 0);
		ZEN_ASSERT(NT_SUCCESS(Status));

		BCryptDestroyHash(HashHandle);
		return Result;
	}

	Sha256Digest BcryptHmac(BCRYPT_ALG_HANDLE Algorithm, const void* Key, size_t KeySize, const void* Data, size_t DataSize)
	{
		Sha256Digest	   Result;
		BCRYPT_HASH_HANDLE HashHandle = nullptr;
		NTSTATUS		   Status;

		Status = BCryptCreateHash(Algorithm, &HashHandle, nullptr, 0, (PUCHAR)Key, (ULONG)KeySize, 0);
		ZEN_ASSERT(NT_SUCCESS(Status));

		Status = BCryptHashData(HashHandle, (PUCHAR)Data, (ULONG)DataSize, 0);
		ZEN_ASSERT(NT_SUCCESS(Status));

		Status = BCryptFinishHash(HashHandle, Result.data(), (ULONG)Result.size(), 0);
		ZEN_ASSERT(NT_SUCCESS(Status));

		BCryptDestroyHash(HashHandle);
		return Result;
	}

	struct BcryptAlgorithmHandles
	{
		BCRYPT_ALG_HANDLE Sha256	 = nullptr;
		BCRYPT_ALG_HANDLE HmacSha256 = nullptr;

		BcryptAlgorithmHandles()
		{
			NTSTATUS Status;
			Status = BCryptOpenAlgorithmProvider(&Sha256, BCRYPT_SHA256_ALGORITHM, nullptr, 0);
			ZEN_ASSERT(NT_SUCCESS(Status));
			Status = BCryptOpenAlgorithmProvider(&HmacSha256, BCRYPT_SHA256_ALGORITHM, nullptr, BCRYPT_ALG_HANDLE_HMAC_FLAG);
			ZEN_ASSERT(NT_SUCCESS(Status));
		}

		~BcryptAlgorithmHandles()
		{
			if (Sha256)
			{
				BCryptCloseAlgorithmProvider(Sha256, 0);
			}
			if (HmacSha256)
			{
				BCryptCloseAlgorithmProvider(HmacSha256, 0);
			}
		}
	};

	BcryptAlgorithmHandles& GetBcryptHandles()
	{
		static BcryptAlgorithmHandles s_Handles;
		return s_Handles;
	}

}  // namespace

Sha256Digest
ComputeSha256(const void* Data, size_t Size)
{
	return BcryptHash(GetBcryptHandles().Sha256, Data, Size);
}

Sha256Digest
ComputeHmacSha256(const void* Key, size_t KeySize, const void* Data, size_t DataSize)
{
	return BcryptHmac(GetBcryptHandles().HmacSha256, Key, KeySize, Data, DataSize);
}

#endif

Sha256Digest
ComputeSha256(std::string_view Data)
{
	return ComputeSha256(Data.data(), Data.size());
}

Sha256Digest
ComputeHmacSha256(const Sha256Digest& Key, std::string_view Data)
{
	return ComputeHmacSha256(Key.data(), Key.size(), Data.data(), Data.size());
}

std::string
Sha256ToHex(const Sha256Digest& Digest)
{
	std::string Result;
	Result.reserve(64);
	for (uint8_t Byte : Digest)
	{
		fmt::format_to(std::back_inserter(Result), "{:02x}", Byte);
	}
	return Result;
}

void
SecureZeroSecret(void* Data, size_t Size)
{
#if ZEN_PLATFORM_WINDOWS
	SecureZeroMemory(Data, Size);
#elif ZEN_PLATFORM_LINUX
	explicit_bzero(Data, Size);
#else
	// Portable fallback: volatile pointer prevents the compiler from optimizing away the memset
	static void* (*const volatile VolatileMemset)(void*, int, size_t) = memset;
	VolatileMemset(Data, 0, Size);
#endif
}

//////////////////////////////////////////////////////////////////////////
// SigV4 signing

namespace {

	std::string GetDateStamp(std::string_view AmzDate)
	{
		// AmzDate is "YYYYMMDDTHHMMSSZ", date stamp is first 8 chars
		return std::string(AmzDate.substr(0, 8));
	}

}  // namespace

std::string
GetAmzTimestamp()
{
	auto		Now		= std::chrono::system_clock::now();
	std::time_t NowTime = std::chrono::system_clock::to_time_t(Now);

	struct tm Tm;
#if ZEN_PLATFORM_WINDOWS
	gmtime_s(&Tm, &NowTime);
#else
	gmtime_r(&NowTime, &Tm);
#endif

	char Buf[32];
	std::strftime(Buf, sizeof(Buf), "%Y%m%dT%H%M%SZ", &Tm);
	return std::string(Buf);
}

std::string
AwsUriEncode(std::string_view Input, bool EncodeSlash)
{
	ExtendableStringBuilder<256> Result;
	for (char C : Input)
	{
		if ((C >= 'A' && C <= 'Z') || (C >= 'a' && C <= 'z') || (C >= '0' && C <= '9') || C == '_' || C == '-' || C == '~' || C == '.')
		{
			Result.Append(C);
		}
		else if (C == '/' && !EncodeSlash)
		{
			Result.Append(C);
		}
		else
		{
			Result.Append(fmt::format("%{:02X}", static_cast<unsigned char>(C)));
		}
	}
	return std::string(Result.ToView());
}

std::string
BuildCanonicalQueryString(std::vector<std::pair<std::string, std::string>> Parameters)
{
	if (Parameters.empty())
	{
		return {};
	}

	// Sort by key name, then by value (as required by SigV4)
	std::sort(Parameters.begin(), Parameters.end());

	ExtendableStringBuilder<512> Result;
	for (size_t i = 0; i < Parameters.size(); ++i)
	{
		if (i > 0)
		{
			Result.Append('&');
		}
		Result.Append(AwsUriEncode(Parameters[i].first));
		Result.Append('=');
		Result.Append(AwsUriEncode(Parameters[i].second));
	}
	return std::string(Result.ToView());
}

SigV4SignedHeaders
SignRequestV4(const SigV4Credentials&								  Credentials,
			  std::string_view										  Method,
			  std::string_view										  Url,
			  std::string_view										  CanonicalQueryString,
			  std::string_view										  Region,
			  std::string_view										  Service,
			  std::string_view										  AmzDate,
			  const std::vector<std::pair<std::string, std::string>>& Headers,
			  std::string_view										  PayloadHash,
			  const Sha256Digest*									  SigningKeyPtr)
{
	SigV4SignedHeaders Result;
	Result.AmzDate	   = std::string(AmzDate);
	Result.PayloadHash = std::string(PayloadHash);

	std::string DateStamp = GetDateStamp(Result.AmzDate);

	// Step 1: Create canonical request
	// CanonicalRequest =
	//   HTTPRequestMethod + '\n' +
	//   CanonicalURI + '\n' +
	//   CanonicalQueryString + '\n' +
	//   CanonicalHeaders + '\n' +
	//   SignedHeaders + '\n' +
	//   HexEncode(Hash(RequestPayload))

	std::string CanonicalUri = AwsUriEncode(Url, false);

	// Build canonical headers and signed headers (headers must be sorted by lowercase name)
	ExtendableStringBuilder<512> CanonicalHeadersSb;
	ExtendableStringBuilder<256> SignedHeadersSb;

	for (size_t i = 0; i < Headers.size(); ++i)
	{
		CanonicalHeadersSb.Append(Headers[i].first);
		CanonicalHeadersSb.Append(':');
		CanonicalHeadersSb.Append(Headers[i].second);
		CanonicalHeadersSb.Append('\n');

		if (i > 0)
		{
			SignedHeadersSb.Append(';');
		}
		SignedHeadersSb.Append(Headers[i].first);
	}

	std::string SignedHeaders = std::string(SignedHeadersSb.ToView());

	std::string CanonicalRequest = fmt::format("{}\n{}\n{}\n{}\n{}\n{}",
											   Method,
											   CanonicalUri,
											   CanonicalQueryString,
											   CanonicalHeadersSb.ToView(),
											   SignedHeaders,
											   PayloadHash);

	// Step 2: Create the string to sign
	std::string CredentialScope = fmt::format("{}/{}/{}/aws4_request", DateStamp, Region, Service);

	Sha256Digest CanonicalRequestHash = ComputeSha256(CanonicalRequest);
	std::string	 CanonicalRequestHex  = Sha256ToHex(CanonicalRequestHash);

	std::string StringToSign = fmt::format("AWS4-HMAC-SHA256\n{}\n{}\n{}", Result.AmzDate, CredentialScope, CanonicalRequestHex);

	// Step 3: Calculate the signing key
	// kDate    = HMAC("AWS4" + SecretKey, DateStamp)
	// kRegion  = HMAC(kDate, Region)
	// kService = HMAC(kRegion, Service)
	// kSigning = HMAC(kService, "aws4_request")

	Sha256Digest DerivedSigningKey;
	if (!SigningKeyPtr)
	{
		std::string SecretPrefix = fmt::format("AWS4{}", Credentials.SecretAccessKey);

		Sha256Digest DateKey = ComputeHmacSha256(SecretPrefix.data(), SecretPrefix.size(), DateStamp.data(), DateStamp.size());
		SecureZeroSecret(SecretPrefix.data(), SecretPrefix.size());

		Sha256Digest RegionKey	= ComputeHmacSha256(DateKey, Region);
		Sha256Digest ServiceKey = ComputeHmacSha256(RegionKey, Service);
		DerivedSigningKey		= ComputeHmacSha256(ServiceKey, "aws4_request");
		SigningKeyPtr			= &DerivedSigningKey;
	}

	// Step 4: Calculate the signature
	Sha256Digest Signature	  = ComputeHmacSha256(*SigningKeyPtr, StringToSign);
	std::string	 SignatureHex = Sha256ToHex(Signature);

	// Step 5: Build the Authorization header
	Result.Authorization = fmt::format("AWS4-HMAC-SHA256 Credential={}/{}, SignedHeaders={}, Signature={}",
									   Credentials.AccessKeyId,
									   CredentialScope,
									   SignedHeaders,
									   SignatureHex);

	return Result;
}

std::string
GeneratePresignedUrl(const SigV4Credentials&								 Credentials,
					 std::string_view										 Method,
					 std::string_view										 Scheme,
					 std::string_view										 Host,
					 std::string_view										 Path,
					 std::string_view										 Region,
					 std::string_view										 Service,
					 std::chrono::seconds									 ExpiresIn,
					 const std::vector<std::pair<std::string, std::string>>& ExtraQueryParams)
{
	// Pre-signed URLs use query string authentication:
	// https://docs.aws.amazon.com/AmazonS3/latest/API/sigv4-query-string-auth.html

	std::string AmzDate	  = GetAmzTimestamp();
	std::string DateStamp = GetDateStamp(AmzDate);

	std::string CredentialScope = fmt::format("{}/{}/{}/aws4_request", DateStamp, Region, Service);
	std::string Credential		= fmt::format("{}/{}", Credentials.AccessKeyId, CredentialScope);

	// The only signed header for pre-signed URLs is "host"
	constexpr std::string_view SignedHeaders = "host";

	// Build query parameters that will be part of the canonical request.
	// These are the auth params (minus X-Amz-Signature which is added after signing).
	std::vector<std::pair<std::string, std::string>> QueryParams = ExtraQueryParams;
	QueryParams.emplace_back("X-Amz-Algorithm", "AWS4-HMAC-SHA256");
	QueryParams.emplace_back("X-Amz-Credential", Credential);
	QueryParams.emplace_back("X-Amz-Date", AmzDate);
	QueryParams.emplace_back("X-Amz-Expires", fmt::format("{}", ExpiresIn.count()));
	if (!Credentials.SessionToken.empty())
	{
		QueryParams.emplace_back("X-Amz-Security-Token", Credentials.SessionToken);
	}
	QueryParams.emplace_back("X-Amz-SignedHeaders", std::string(SignedHeaders));

	std::string CanonicalQueryString = BuildCanonicalQueryString(QueryParams);
	std::string CanonicalUri		 = AwsUriEncode(Path, false);

	// For pre-signed URLs, the payload is always UNSIGNED-PAYLOAD
	constexpr std::string_view PayloadHash = "UNSIGNED-PAYLOAD";

	// Build the canonical request
	// Only "host" is in the canonical headers for pre-signed URLs
	std::string CanonicalHeaders = fmt::format("host:{}\n", Host);

	std::string CanonicalRequest =
		fmt::format("{}\n{}\n{}\n{}\n{}\n{}", Method, CanonicalUri, CanonicalQueryString, CanonicalHeaders, SignedHeaders, PayloadHash);

	// Create the string to sign
	Sha256Digest CanonicalRequestHash = ComputeSha256(CanonicalRequest);
	std::string	 CanonicalRequestHex  = Sha256ToHex(CanonicalRequestHash);

	std::string StringToSign = fmt::format("AWS4-HMAC-SHA256\n{}\n{}\n{}", AmzDate, CredentialScope, CanonicalRequestHex);

	// Calculate the signing key
	std::string	 SecretPrefix = fmt::format("AWS4{}", Credentials.SecretAccessKey);
	Sha256Digest DateKey	  = ComputeHmacSha256(SecretPrefix.data(), SecretPrefix.size(), DateStamp.data(), DateStamp.size());
	SecureZeroSecret(SecretPrefix.data(), SecretPrefix.size());
	Sha256Digest RegionKey	= ComputeHmacSha256(DateKey, Region);
	Sha256Digest ServiceKey = ComputeHmacSha256(RegionKey, Service);
	Sha256Digest SigningKey = ComputeHmacSha256(ServiceKey, "aws4_request");

	// Calculate the signature
	std::string SignatureHex = Sha256ToHex(ComputeHmacSha256(SigningKey, StringToSign));

	// Build the final URL (use the URI-encoded path so special characters are properly escaped)
	return fmt::format("{}://{}{}?{}&X-Amz-Signature={}", Scheme, Host, CanonicalUri, CanonicalQueryString, SignatureHex);
}

//////////////////////////////////////////////////////////////////////////
// Tests

#if ZEN_WITH_TESTS

void
sigv4_forcelink()
{
}

TEST_SUITE_BEGIN("util.cloud.sigv4");

TEST_CASE("sigv4.sha256")
{
	// Test with known test vector (empty string)
	Sha256Digest Empty = ComputeSha256("", 0);
	std::string	 Hex   = Sha256ToHex(Empty);
	CHECK(Hex == "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855");

	// Test with "hello"
	Sha256Digest Hello	  = ComputeSha256("hello");
	std::string	 HelloHex = Sha256ToHex(Hello);
	CHECK(HelloHex == "2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824");
}

TEST_CASE("sigv4.hmac_sha256")
{
	// RFC 4231 Test Case 2
	std::string_view Key  = "Jefe";
	std::string_view Data = "what do ya want for nothing?";

	Sha256Digest Result = ComputeHmacSha256(Key.data(), Key.size(), Data.data(), Data.size());
	std::string	 Hex	= Sha256ToHex(Result);
	CHECK(Hex == "5bdcc146bf60754e6a042426089575c75a003f089d2739839dec58b964ec3843");
}

TEST_CASE("sigv4.signing")
{
	// Based on the AWS SigV4 test suite example
	// https://docs.aws.amazon.com/general/latest/gr/sigv4-calculate-signature.html

	SigV4Credentials Creds;
	Creds.AccessKeyId	  = "AKIDEXAMPLE";
	Creds.SecretAccessKey = "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY";

	// We can't test with a fixed timestamp since SignRequestV4 uses current time,
	// but we can verify the crypto primitives produce correct results by testing
	// the signing key derivation manually.

	// Test signing key derivation: HMAC chain for "20150830" / "us-east-1" / "iam"
	std::string	 SecretPrefix = "AWS4wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY";
	Sha256Digest DateKey	  = ComputeHmacSha256(SecretPrefix.data(), SecretPrefix.size(), "20150830", 8);
	Sha256Digest RegionKey	  = ComputeHmacSha256(DateKey, "us-east-1");
	Sha256Digest ServiceKey	  = ComputeHmacSha256(RegionKey, "iam");
	Sha256Digest SigningKey	  = ComputeHmacSha256(ServiceKey, "aws4_request");

	std::string SigningKeyHex = Sha256ToHex(SigningKey);
	CHECK(SigningKeyHex == "c4afb1cc5771d871763a393e44b703571b55cc28424d1a5e86da6ed3c154a4b9");
}

TEST_SUITE_END();

#endif

}  // namespace zen