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

#include "authutils.h"

#include <zencore/crypto.h>
#include <zencore/filesystem.h>
#include <zencore/fmtutils.h>
#include <zencore/iobuffer.h>
#include <zencore/logging.h>

#include <zenhttp/auth/authmgr.h>
#include <zenhttp/httpclient.h>
#include <zenhttp/httpclientauth.h>
#include <zenutil/authutils.h>

ZEN_THIRD_PARTY_INCLUDES_START
#include <json11.hpp>
ZEN_THIRD_PARTY_INCLUDES_END

namespace zen {
using namespace std::literals;

std::string_view
GetDefaultAccessTokenEnvVariableName()
{
#if ZEN_PLATFORM_WINDOWS
	return "UE-CloudDataCacheAccessToken"sv;
#endif
#if ZEN_PLATFORM_LINUX || ZEN_PLATFORM_MAC
	return "UE_CloudDataCacheAccessToken"sv;
#endif
}

std::string
ReadAccessTokenFromJsonFile(const std::filesystem::path& Path)
{
	if (!IsFile(Path))
	{
		throw std::runtime_error(fmt::format("the file '{}' does not exist", Path));
	}
	IoBuffer	 Body = IoBufferBuilder::MakeFromFile(Path);
	std::string	 JsonText(reinterpret_cast<const char*>(Body.GetData()), Body.GetSize());
	std::string	 JsonError;
	json11::Json TokenInfo = json11::Json::parse(JsonText, JsonError);
	if (!JsonError.empty())
	{
		throw std::runtime_error(fmt::format("failed parsing json file '{}'. Reason: '{}'", Path, JsonError));
	}
	const std::string AuthToken = TokenInfo["Token"].string_value();
	if (AuthToken.empty())
	{
		throw std::runtime_error(fmt::format("the json file '{}' does not contain a value for \"Token\"", Path));
	}
	return AuthToken;
}

void
AuthCommandLineOptions::AddOptions(cxxopts::Options& Ops)
{
	// Direct access token (may expire)
	Ops.add_option("auth-token", "", "access-token", "Remote host access token", cxxopts::value(m_AccessToken), "<accesstoken>");
	Ops.add_option("auth-token",
				   "",
				   "access-token-env",
				   "Name of environment variable that holds the remote host access token",
				   cxxopts::value(m_AccessTokenEnv)->default_value(std::string(GetDefaultAccessTokenEnvVariableName())),
				   "<envvariable>");
	Ops.add_option("auth-token",
				   "",
				   "access-token-path",
				   "Path to json file that holds the remote host access token",
				   cxxopts::value(m_AccessTokenPath),
				   "<filepath>");

	// Auth manager token encryption
	Ops.add_option("security", "", "encryption-aes-key", "256 bit AES encryption key", cxxopts::value<std::string>(m_EncryptionKey), "");
	Ops.add_option("security",
				   "",
				   "encryption-aes-iv",
				   "128 bit AES encryption initialization vector",
				   cxxopts::value<std::string>(m_EncryptionIV),
				   "");

	// OpenId acccess token
	Ops.add_option("openid",
				   "",
				   "openid-provider-name",
				   "Open ID provider name",
				   cxxopts::value<std::string>(m_OpenIdProviderName),
				   "Default");
	Ops.add_option("openid", "", "openid-provider-url", "Open ID provider url", cxxopts::value<std::string>(m_OpenIdProviderUrl), "");
	Ops.add_option("openid", "", "openid-client-id", "Open ID client id", cxxopts::value<std::string>(m_OpenIdClientId), "");
	Ops.add_option("openid", "", "openid-refresh-token", "Open ID refresh token", cxxopts::value<std::string>(m_OpenIdRefreshToken), "");

	// OAuth acccess token
	Ops.add_option("oauth", "", "oauth-url", "OAuth provier url", cxxopts::value<std::string>(m_OAuthUrl)->default_value(""), "");
	Ops.add_option("oauth", "", "oauth-clientid", "OAuth client id", cxxopts::value<std::string>(m_OAuthClientId)->default_value(""), "");
	Ops.add_option("oauth",
				   "",
				   "oauth-clientsecret",
				   "OAuth client secret",
				   cxxopts::value<std::string>(m_OAuthClientSecret)->default_value(""),
				   "");
	Ops.add_option("auth",
				   "",
				   "oidctoken-exe-path",
				   "Path to OidcToken executable",
				   cxxopts::value<std::string>(m_OidcTokenAuthExecutablePath)->default_value(""),
				   "");
	Ops.add_option("auth",
				   "",
				   "oidctoken-exe-unattended",
				   "Set mode to unattended when launcing OidcToken executable",
				   cxxopts::value<bool>(m_OidcTokenUnattended),
				   "");
};

// Load or generate a per-install machine AES key+IV under AuthDir/machinekey.dat
// so the auth-state file is encrypted with bytes unique to this machine rather
// than a hardcoded constant.
//
// When per-user OS-protected storage is available (DPAPI on Windows) the key
// material is wrapped before it lands on disk, so a copy of the file off-machine
// or out of a backup cannot be unwrapped without also stealing the user's OS
// master key. On platforms without OS-level wrapping we fall back to persisting
// the raw bytes with restrictive file permissions (0600 on POSIX; user-only on
// Windows via inheritance from the profile dir).
//
// File format:
//   [4-byte magic 'Z','E','N','\x01'] [1-byte flags] [payload]
//   flags bit 0 set -> payload is OS-protected (DPAPI blob)
//   flags bit 0 clear -> payload is raw KeyBytes+IvBytes bytes
// Legacy files without the magic are interpreted as raw bytes.
void
AuthCommandLineOptions::LoadOrCreateMachineKey(const std::filesystem::path& AuthDir, bool Quiet)
{
	constexpr size_t						KeyBytes	  = AesKey256Bit::ByteCount;
	constexpr size_t						IvBytes		  = AesIV128Bit::ByteCount;
	static constexpr std::array<uint8_t, 4> FileMagic	  = {'Z', 'E', 'N', 0x01};
	static constexpr uint8_t				FlagProtected = 0x01;
	const std::filesystem::path				KeyFile		  = AuthDir / "machinekey.dat";
	std::array<uint8_t, KeyBytes + IvBytes> KeyMaterial{};
	bool									Loaded = false;

	auto ParseFile = [&](MemoryView FileBytes) -> bool {
		// Legacy: raw KeyBytes+IvBytes payload.
		if (FileBytes.GetSize() == KeyMaterial.size())
		{
			memcpy(KeyMaterial.data(), FileBytes.GetData(), KeyMaterial.size());
			return true;
		}
		if (FileBytes.GetSize() < FileMagic.size() + 1)
		{
			return false;
		}
		if (memcmp(FileBytes.GetData(), FileMagic.data(), FileMagic.size()) != 0)
		{
			return false;
		}
		const uint8_t	 Flags	 = static_cast<const uint8_t*>(FileBytes.GetData())[FileMagic.size()];
		const MemoryView Payload = FileBytes.Mid(FileMagic.size() + 1);
		if (Flags & FlagProtected)
		{
			std::vector<uint8_t> Plaintext;
			if (!TryUnprotectData(Payload, Plaintext))
			{
				if (!Quiet)
				{
					ZEN_CONSOLE_WARN("Auth: failed to unwrap OS-protected machine key at '{}', regenerating", KeyFile);
				}
				return false;
			}
			if (Plaintext.size() != KeyMaterial.size())
			{
				return false;
			}
			memcpy(KeyMaterial.data(), Plaintext.data(), KeyMaterial.size());
			return true;
		}
		if (Payload.GetSize() != KeyMaterial.size())
		{
			return false;
		}
		memcpy(KeyMaterial.data(), Payload.GetData(), KeyMaterial.size());
		return true;
	};

	std::error_code Ec;
	if (std::filesystem::exists(KeyFile, Ec))
	{
		IoBuffer Data = ReadFile(KeyFile).Flatten();
		if (ParseFile(Data.GetView()))
		{
			Loaded = true;
		}
		else if (!Quiet)
		{
			ZEN_CONSOLE_WARN("Auth: machine key file '{}' is unreadable (size {}), regenerating", KeyFile, Data.GetSize());
		}
	}

	if (!Loaded)
	{
		CreateDirectories(AuthDir);
		if (!SecureRandomBytes(MutableMemoryView(KeyMaterial.data(), KeyMaterial.size())))
		{
			throw std::runtime_error("failed to obtain secure random bytes for auth machine key");
		}

		std::vector<uint8_t> FileBytes;
		FileBytes.reserve(FileMagic.size() + 1 + KeyMaterial.size());
		FileBytes.insert(FileBytes.end(), FileMagic.begin(), FileMagic.end());

		std::vector<uint8_t> Wrapped;
		if (TryProtectData(MemoryView(KeyMaterial.data(), KeyMaterial.size()), Wrapped))
		{
			FileBytes.push_back(FlagProtected);
			FileBytes.insert(FileBytes.end(), Wrapped.begin(), Wrapped.end());
			if (!Quiet)
			{
				ZEN_CONSOLE_WARN("Auth: generated OS-protected machine-specific auth encryption key at '{}'", KeyFile);
			}
		}
		else
		{
			FileBytes.push_back(0);
			FileBytes.insert(FileBytes.end(), KeyMaterial.begin(), KeyMaterial.end());
			if (!Quiet)
			{
				ZEN_CONSOLE_WARN("Auth: generated machine-specific auth encryption key at '{}' (no OS wrapping available)", KeyFile);
			}
		}

		WriteFile(KeyFile, IoBufferBuilder::MakeCloneFromMemory(FileBytes.data(), FileBytes.size()));

		// Belt and suspenders: restrict access on POSIX. On Windows the
		// default DACL inherited from a per-user profile dir is already
		// user-only in the common case; an explicit tighten there would
		// require touching the DACL which is more code than it's worth
		// while DPAPI wrapping is the primary defense.
#if !ZEN_PLATFORM_WINDOWS
		std::error_code PermEc;
		std::filesystem::permissions(KeyFile,
									 std::filesystem::perms::owner_read | std::filesystem::perms::owner_write,
									 std::filesystem::perm_options::replace,
									 PermEc);
#endif
	}

	m_EncryptionKey.assign(reinterpret_cast<const char*>(KeyMaterial.data()), KeyBytes);
	m_EncryptionIV.assign(reinterpret_cast<const char*>(KeyMaterial.data() + KeyBytes), IvBytes);
}

void
AuthCommandLineOptions::CreateAuthMgr(cxxopts::Options&			   Ops,
									  const std::filesystem::path& SystemRootDir,
									  std::unique_ptr<AuthMgr>&	   InOutAuth,
									  bool						   Quiet,
									  bool						   Verbose)
{
	ZEN_ASSERT(!SystemRootDir.empty());
	if (InOutAuth)
	{
		return;
	}

	const std::filesystem::path AuthDir = SystemRootDir / "auth";

	if (m_EncryptionKey.empty() != m_EncryptionIV.empty())
	{
		throw OptionParseException(
			std::string("'--encryption-aes-key' and '--encryption-aes-iv' must be supplied together or both omitted"),
			Ops.help());
	}

	if (m_EncryptionKey.empty() && m_EncryptionIV.empty())
	{
		LoadOrCreateMachineKey(AuthDir, Quiet);
	}

	AuthConfig AuthMgrConfig = {.RootDirectory = AuthDir,
								.EncryptionKey = AesKey256Bit::FromString(m_EncryptionKey),
								.EncryptionIV  = AesIV128Bit::FromString(m_EncryptionIV)};
	if (!AuthMgrConfig.EncryptionKey.IsValid())
	{
		throw OptionParseException(fmt::format("'--encryption-aes-key' ('{}') is malformed", m_EncryptionKey), Ops.help());
	}
	if (!AuthMgrConfig.EncryptionIV.IsValid())
	{
		throw OptionParseException(fmt::format("'--encryption-aes-iv' ('{}') is malformed", m_EncryptionIV), Ops.help());
	}
	if (Verbose)
	{
		ExtendableStringBuilder<128> SB;
		SB << "\n    RootDirectory: " << AuthMgrConfig.RootDirectory.string();
		SB << "\n    EncryptionKey: " << HideSensitiveString(m_EncryptionKey);
		SB << "\n    EncryptionIV:  " << HideSensitiveString(m_EncryptionIV);
		ZEN_CONSOLE("Auth: Creating auth manager with:{}", SB.ToString());
	}
	InOutAuth = AuthMgr::Create(AuthMgrConfig);
}

void
AuthCommandLineOptions::ParseOptions(cxxopts::Options&			  Ops,
									 const std::filesystem::path& SystemRootDir,
									 HttpClientSettings&		  ClientSettings,
									 std::string_view			  HostUrl,
									 std::unique_ptr<AuthMgr>&	  Auth,
									 bool						  Quiet,
									 bool						  Hidden,
									 bool						  Verbose)
{
	if (!m_OpenIdProviderUrl.empty() && !m_OpenIdClientId.empty())
	{
		CreateAuthMgr(Ops, SystemRootDir, Auth, Quiet, Verbose);
		std::string ProviderName = m_OpenIdProviderName.empty() ? "Default" : m_OpenIdProviderName;
		if (Verbose)
		{
			ExtendableStringBuilder<128> SB;
			SB << "\n    Name:     " << ProviderName;
			SB << "\n    Url:      " << m_OpenIdProviderUrl;
			SB << "\n    ClientId: " << HideSensitiveString(m_OpenIdClientId);
			ZEN_CONSOLE("Auth: Adding Open ID auth provider:{}", SB.ToString());
		}
		Auth->AddOpenIdProvider({.Name = ProviderName, .Url = m_OpenIdProviderUrl, .ClientId = m_OpenIdClientId});
		if (!m_OpenIdRefreshToken.empty())
		{
			if (!Quiet)
			{
				ZEN_CONSOLE("Auth: Adding open id refresh token {} to provider {}",
							HideSensitiveString(m_OpenIdRefreshToken),
							ProviderName);
			}
			Auth->AddOpenIdToken({.ProviderName = ProviderName, .RefreshToken = m_OpenIdRefreshToken});
		}
	}

	auto GetEnvAccessToken = [](const std::string& AccessTokenEnv) -> std::string {
		if (!AccessTokenEnv.empty())
		{
			return GetEnvVariable(AccessTokenEnv).value_or("");
		}
		return {};
	};

	if (!m_AccessToken.empty())
	{
		if (!Quiet)
		{
			ZEN_CONSOLE("Auth: Using static auth token: {}", HideSensitiveString(m_AccessToken));
		}
		ClientSettings.AccessTokenProvider = httpclientauth::CreateFromStaticToken(m_AccessToken);
	}
	else if (!m_AccessTokenPath.empty())
	{
		MakeSafeAbsolutePathInPlace(m_AccessTokenPath);
		std::string ResolvedAccessToken = ReadAccessTokenFromJsonFile(m_AccessTokenPath);
		if (!ResolvedAccessToken.empty())
		{
			if (!Quiet)
			{
				ZEN_CONSOLE("Auth: Adding static auth token from {}: {}", m_AccessTokenPath, HideSensitiveString(ResolvedAccessToken));
			}
			ClientSettings.AccessTokenProvider = httpclientauth::CreateFromStaticToken(ResolvedAccessToken);
		}
	}
	else if (!m_OAuthUrl.empty())
	{
		if (Verbose)
		{
			ExtendableStringBuilder<128> SB;
			SB << "\n    Url:          " << m_OAuthUrl;
			SB << "\n    ClientId:     " << HideSensitiveString(m_OAuthClientId);
			SB << "\n    ClientSecret: " << HideSensitiveString(m_OAuthClientSecret);
			ZEN_CONSOLE("Auth: Adding oauth provider:{}", SB.ToString());
		}
		ClientSettings.AccessTokenProvider = httpclientauth::CreateFromOAuthClientCredentials(
			{.Url = m_OAuthUrl, .ClientId = m_OAuthClientId, .ClientSecret = m_OAuthClientSecret});
	}
	else if (!m_OpenIdProviderName.empty())
	{
		CreateAuthMgr(Ops, SystemRootDir, Auth, Quiet, Verbose);
		if (!Quiet)
		{
			ZEN_CONSOLE("Auth: Using OpenId provider: {}", m_OpenIdProviderName);
		}
		ClientSettings.AccessTokenProvider = httpclientauth::CreateFromOpenIdProvider(*Auth, m_OpenIdProviderName);
	}
	else if (std::string ResolvedAccessToken = GetEnvAccessToken(m_AccessTokenEnv); !ResolvedAccessToken.empty())
	{
		if (!Quiet)
		{
			ZEN_CONSOLE("Auth: Resolved environment variable '{}' to access token '{}'",
						m_AccessTokenEnv,
						HideSensitiveString(ResolvedAccessToken));
		}
		ClientSettings.AccessTokenProvider = httpclientauth::CreateFromStaticToken(ResolvedAccessToken);
	}
	else if (std::filesystem::path OidcTokenExePath = FindOidcTokenExePath(m_OidcTokenAuthExecutablePath); !OidcTokenExePath.empty())
	{
		if (!Quiet)
		{
			ZEN_CONSOLE("Auth: Using oidctoken exe from path '{}'", OidcTokenExePath);
		}
		ClientSettings.AccessTokenProvider =
			httpclientauth::CreateFromOidcTokenExecutable(OidcTokenExePath, HostUrl, Quiet, m_OidcTokenUnattended, Hidden);
	}
	else if (!m_OidcTokenAuthExecutablePath.empty())
	{
		throw OptionParseException(fmt::format("'--oidctoken-exe-path' ('{}') does not exist", m_OidcTokenAuthExecutablePath), Ops.help());
	}

	if (!ClientSettings.AccessTokenProvider)
	{
		CreateAuthMgr(Ops, SystemRootDir, Auth, Quiet, Verbose);
		if (!Quiet)
		{
			ZEN_CONSOLE("Auth: Using default Open ID provider");
		}
		ClientSettings.AccessTokenProvider = httpclientauth::CreateFromDefaultOpenIdProvider(*Auth);
	}
}
}  // namespace zen