#include "pch.h"
#include "Utils.h"
#include "tiny-AES-c/aes.h"

const uint8_t IV[] = { 0x72, 0xE0, 0x67, 0xFB, 0xDD, 0xCB, 0xCF, 0x77, 0xEB, 0xE8, 0xBC, 0x64, 0x3F, 0x63, 0x0D, 0x93 };

bool Utils::Detour32(char* src, char* dst, const intptr_t len)
{
	if (len < 5) return false;

	DWORD  curProtection;
	VirtualProtect(src, len, PAGE_EXECUTE_READWRITE, &curProtection);

	intptr_t  relativeAddress = (intptr_t)(dst - (intptr_t)src) - 5;

	*src = (char)'\xE9';
	*(intptr_t*)((intptr_t)src + 1) = relativeAddress;

	VirtualProtect(src, len, curProtection, &curProtection);
	return true;
}

char* Utils::TrampHook32(char* src, char* dst, const intptr_t len)
{
	// Make sure the length is greater than 5
	if (len < 5) return 0;

	// Create the gateway (len + 5 for the overwritten bytes + the jmp)
	void* gateway = VirtualAlloc(0, len + 5, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);

	if (gateway == NULL) return 0;

	// Write the stolen bytes into the gateway
	memcpy(gateway, src, len);

	// Get the gateway to destination addy
	intptr_t gatewayRelativeAddr = ((intptr_t)src - (intptr_t)gateway) - 5;

	// Add the jmp opcode to the end of the gateway
	*(char*)((intptr_t)gateway + len) = 0xE9;

	// Add the address to the jmp
	*(intptr_t*)((intptr_t)gateway + len + 1) = gatewayRelativeAddr;

	// Perform the detour
	Detour32(src, dst, len);

	return (char*)gateway;
}

int spotifyVer = -1;
int spotifyVerEnd = -1;
int Utils::GetSpotifyVersion()
{
	if (spotifyVer != -1)
		return spotifyVer;

	LPCWSTR lpszFilePath = L"Spotify.exe";
	DWORD dwDummy;
	DWORD dwFVISize = GetFileVersionInfoSize(lpszFilePath, &dwDummy);
	LPBYTE lpVersionInfo = new BYTE[dwFVISize];
	GetFileVersionInfo(lpszFilePath, 0, dwFVISize, lpVersionInfo);
	UINT uLen;
	VS_FIXEDFILEINFO* lpFfi;
	VerQueryValue(lpVersionInfo, _T("\\"), (LPVOID*)&lpFfi, &uLen);
	DWORD dwFileVersionMS = lpFfi->dwFileVersionMS;
	DWORD dwFileVersionLS = lpFfi->dwFileVersionLS;
	delete[] lpVersionInfo;

	DWORD dwLeftMost = HIWORD(dwFileVersionMS);
	DWORD dwSecondLeft = LOWORD(dwFileVersionMS);
	DWORD dwSecondRight = HIWORD(dwFileVersionLS);
	DWORD dwRightMost = LOWORD(dwFileVersionLS);

	spotifyVerEnd = dwRightMost;
	return spotifyVer = dwSecondRight;
}

std::string Utils::HexString(BYTE* data, int len)
{
	std::stringstream ss;
	ss << std::hex;

	for (int i(0); i < len; ++i)
		ss << std::setw(2) << std::setfill('0') << (int)data[i];

	return ss.str();
}

std::wstring Utils::FixPathStr(std::wstring str)
{
	// No forbidden characters
	std::wstring badCharsRegex = std::wstring(L"[<>:\"/\\|?*\x00-\x1F]", 14); // Use this constructor for \x00
	str = std::regex_replace(str, std::wregex(badCharsRegex), L"_");

	// No forbidden words or ending periods or spaces
	if (std::regex_match(str, std::wregex(L"^(CON|PRN|AUX|NUL|COM[0-9]|LPT[0-9])$| +$|\\.+$",
		std::regex_constants::ECMAScript | std::regex_constants::icase)))
		str.append(L"_");

	return str;
}

std::wstring Utils::Utf8ToUtf16(const std::string& str)
{
	std::wstring convertedString;
	int requiredSize = MultiByteToWideChar(CP_UTF8, 0, str.c_str(), -1, 0, 0);
	if (requiredSize > 0)
	{
		std::vector<wchar_t> buffer(requiredSize);
		MultiByteToWideChar(CP_UTF8, 0, str.c_str(), -1, &buffer[0], requiredSize);
		convertedString.assign(buffer.begin(), buffer.end() - 1);
	}

	return convertedString;
}

struct SongInfo
{
	std::string title, artist, album, cover;
} songInfo;

void ClearSongInfo()
{
	songInfo.title = "";
	songInfo.artist = "";
	songInfo.album = "";
	songInfo.cover = "";
}

static const std::string songRegex =
    "https?:\\/\\/(?:www\.)?([-a-zA-Z0-9@:%._\\+~#=]{1,256}\\.[a-zA-Z0-9()]{1,6})"
	"\\b([-a-zA-Z0-9()@:%_\\+.~#?&\\/\\/=]*)";
static const std::string albumSearchPattern = "\x68\x65\x69\x67\x68\x74\x22\x20\x3A\x20\x36\x34\x30";
static const std::wstring songDirRoot = L"Downloads";
static std::wstring songDir = songDirRoot;
void Utils::DownloadSong(std::string fileId, std::string uri, std::string key, std::string authToken)
{
	std::string downloadStr;
	std::wstring songExtension = L".ogg";

	if (fileId.empty() || uri.empty() || authToken.empty())
	{
		std::cout << "Could not download song or episode: missing fileId, trackUri, or authToken!" << std::endl;
		return;
	}

	// Is length check even needed for compare()?
	if (uri.length() > 13 && uri.compare(0, 13, "spotify:track") == 0)
	{
		std::cout << "Downloading song..." << std::endl;

		if (key.empty())
		{
			std::cout << "Could not download song: missing key!";
			return;
		}

		// Get storage resolve from Spotify
		std::string srStr = DownloadSpotifyUrl("spclient.wg.spotify.com",
			"/storage-resolve/files/audio/interactive_prefetch/" + fileId + "?product=0", authToken);

		if (srStr.length() <= 5)
		{
			std::cout << "Error: Couldn't fetch storage resolve!" << std::endl;
			return;
		}

		if (srStr.substr(0, 5).compare("Error") == 0)
		{
			std::cout << srStr << std::endl;
			return;
		}

		// Parse storage resolve response to get the encrypted song data's URL
		std::string songHost;
		std::string songPath;
		try
		{
			std::regex re(songRegex);
			std::smatch match;
			if (std::regex_search(srStr, match, re) && match.size() > 1)
			{
				songHost = match.str(1);
				songPath = match.str(2);
			}
			else
			{
				std::cout << "Error: Download URL not found" << std::endl;
				return;
			}
		}
		catch (std::regex_error& e)
		{
			// Syntax error in the regular expression
			std::cout << "Error: Invalid regex!" << std::endl;
			return;
		}

		// Download encrypted song data from Spotify
		downloadStr = DownloadSpotifyUrl(songHost, songPath, "");

		if (downloadStr.length() <= 6)
		{
			std::cout << "Error: Could not download audio!" << std::endl;
			return;
		}

		if (downloadStr.substr(0, 6).compare("<HTML>") == 0)
		{
			std::cout << "Error: " + downloadStr << std::endl;
			return;
		}

		// Decrypt encrypted song data with Tiny AES in C
		struct AES_ctx ctx;
		AES_init_ctx_iv(&ctx, reinterpret_cast<const uint8_t*>(&key[0]), IV);
		AES_CTR_xcrypt_buffer(&ctx, reinterpret_cast<uint8_t*>(&downloadStr[0]), downloadStr.size());

		// Check if decrypted song data has valid header
		if (downloadStr.length() < 4 || downloadStr.compare(0, 4, "OggS") != 0)
		{
			std::cout << "Error: Could not properly decrypt song (try downloading again)!" << std::endl;
			return;
		}

		// Remove custom Spotify Ogg page from beginning of file
		downloadStr = downloadStr.substr(downloadStr.find("\xFF\xFF\xFF\xFFOggS") + 4);

		// Download song metadata from Spotify API
		std::string metadata = DownloadSpotifyUrl("api.spotify.com", "/v1/tracks/"
			+ uri.substr(uri.find("spotify:track:") + 14), authToken);

		songInfo.title = strtok((char*)(metadata.substr(metadata.find("is_local") + 55)).c_str(), "\"");
		songInfo.artist = strtok((char*)(metadata.substr(metadata.find("name") + 9)).c_str(), "\"");
		songInfo.album = strtok((char*)(metadata.substr(metadata.find(albumSearchPattern) + 404)).c_str(), "\"");
		songInfo.cover = strtok((char*)(metadata.substr(metadata.find("height") + 30)).c_str(), "\"");

		songExtension = L".ogg";
	}
	else if (uri.length() > 15 && uri.compare(0, 15, "spotify:episode") == 0)
	{
		std::string episodeUrl = fileId;
		std::string songHost;
		std::string songPath;

		std::cout << "Downloading episode..." << std::endl;

		// Parse episode URL to separate host and path
		try
		{
			std::regex re(songRegex);
			std::smatch match;
			if (std::regex_search(episodeUrl, match, re) && match.size() > 1)
			{
				songHost = match.str(1);
				songPath = match.str(2);
			}
			else
			{
				std::cout << "Error: Download URL is not valid!" << std::endl;
				return;
			}
		}
		catch (std::regex_error& e)
		{
			// Syntax error in the regular expression
			std::cout << "Error: Invalid regex!" << std::endl;
			return;
		}

		// Download episode data from Spotify
		downloadStr = DownloadSpotifyUrl(songHost, songPath, "");

		// TODO: MP3 files appear to have a variety of headers, not sure this is 100% reliable
		// Check if downloaded episode data has valid header
		/*if (downloadStr.length() < 3 || downloadStr.compare(0, 3, "ID3") != 0)
		{
			std::cout << "Error: Could not properly download podcast episode (try downloading again)!" << std::endl;
			return;
		}*/

		// Download episode and show metadata from Spotify API
		std::string metadata = DownloadSpotifyUrl("api.spotify.com", "/v1/episodes/"
			+ uri.substr(uri.find("spotify:episode:") + 16), authToken);

		songInfo.title = strtok((char*)(metadata.substr(metadata.find("name") + 9)).c_str(), "\"");
		songInfo.artist = strtok((char*)(metadata.substr(metadata.find("publisher") + 14)).c_str(), "\"");
		songInfo.album = strtok((char*)(metadata.substr(metadata.find("media_type") + 37)).c_str(), "\"");
		songInfo.cover = strtok((char*)(metadata.substr(metadata.find("height") + 28)).c_str(), "\"");

		songExtension = L".mp3";
	}

	if (songInfo.title.empty() || songInfo.artist.empty() || songInfo.album.empty())
	{
		std::cout << "Error: Invalid title/artist/album name!" << std::endl;
		return;
	}

	std::wstring tempDirArtist = FixPathStr(Utf8ToUtf16(songInfo.artist));

	songDir = songDirRoot;
	if (!CreateDirectoryW(songDir.c_str(), NULL) && ERROR_ALREADY_EXISTS != GetLastError())
		std::cout << "Couldn't create main downloads directory!" << std::endl;

	if (CreateDirectoryW(std::wstring(songDir + L"\\" + tempDirArtist).c_str(), NULL)
		|| ERROR_ALREADY_EXISTS == GetLastError())
	{
		std::wstring tempDirAlbum = FixPathStr(Utf8ToUtf16(songInfo.album));

		if (CreateDirectoryW(std::wstring(songDir + L"\\" + tempDirArtist + std::wstring(L"\\")
			+ tempDirAlbum).c_str(), NULL) || ERROR_ALREADY_EXISTS == GetLastError())
		{
			songDir += L"\\" + tempDirArtist + std::wstring(L"\\") + tempDirAlbum;

			std::wstring tempDirSong = FixPathStr(Utf8ToUtf16(songInfo.title));

			std::ofstream songFileOut(songDir + L".\\" + tempDirArtist + L" - " + tempDirSong + songExtension,
				std::ios_base::binary);
			songFileOut.write(downloadStr.c_str(), downloadStr.size());
			songFileOut.close();

			std::cout << "Finished downloading: " << songInfo.artist << " - \"" << songInfo.title << "\"!" << std::endl;

			ClearSongInfo();
			return;
		}
		else
		{
			std::cout << "Couldn't create album directory!" << std::endl;
		}
	}
	else
	{
		std::cout << "Couldn't create artist directory!" << std::endl;
	}

	std::cout << "Could not finish downloading song!" << std::endl;

	ClearSongInfo();
}

std::string GetLastErrorAsString()
{
	//Get the error message, if any.
	DWORD errorMessageID = ::GetLastError();
	if (errorMessageID == 0)
		return std::string(); //No error message has been recorded

	LPSTR messageBuffer = nullptr;
	size_t size = FormatMessageA(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
		NULL, errorMessageID, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), (LPSTR)&messageBuffer, 0, NULL);

	std::string message(messageBuffer, size);

	//Free the buffer.
	LocalFree(messageBuffer);

	return message;
}

std::string Utils::DownloadSpotifyUrl(std::string host, std::string path, std::string authToken)
{
	std::string response;
	std::string authHeader = (authToken.empty()) ? "" : "Authorization: Bearer " + authToken;
	std::string userAgent = "Spotify/11" + std::to_string(spotifyVer) + std::string("00")
		+ std::to_string(spotifyVerEnd) + std::string(" Win32/Windows 10 (10.0.19042; x64)");
	HINTERNET hSession, hConnect, hRequest;
	BOOL bRequestSent;
	const int bufferSize = 1024;

	hSession = InternetOpenA(userAgent.c_str(), INTERNET_OPEN_TYPE_DIRECT, NULL, NULL, 0);

	if (hSession == NULL)
		return "Error: Could not initialize request!";

	hConnect = InternetConnectA(hSession, host.c_str(), 80, NULL, NULL, INTERNET_SERVICE_HTTP, 0,
		NULL);

	if (hConnect == NULL)
		return "Error: Could not create connect!";

	hRequest = HttpOpenRequestA(hConnect, "GET", path.c_str(), NULL, NULL, NULL, INTERNET_FLAG_NO_AUTH, 0);

	if (hRequest == NULL)
		return "Error: Could not create open request!";

	HttpAddRequestHeadersA(hRequest, authHeader.c_str(), -1, HTTP_ADDREQ_FLAG_ADD | HTTP_ADDREQ_FLAG_REPLACE);
	bRequestSent = HttpSendRequestA(hRequest, NULL, 0, NULL, 0);

	if (!bRequestSent)
		return "Error: Could not send request!";

	char tmpBuffer[bufferSize];
	BOOL canRead = true;
	DWORD bytesRead = -1;

	while (InternetReadFile(hRequest, tmpBuffer, bufferSize, &bytesRead) && bytesRead)
		response.append(tmpBuffer, bytesRead);

	InternetCloseHandle(hRequest);

	return response;
}

bool Utils::BadPtr(void* ptr)
{
	MEMORY_BASIC_INFORMATION mbi = { 0 };
	if (VirtualQuery(ptr, &mbi, sizeof(mbi)))
	{
		DWORD mask = (PAGE_READONLY | PAGE_READWRITE | PAGE_WRITECOPY | PAGE_EXECUTE_READ | PAGE_EXECUTE_READWRITE
			| PAGE_EXECUTE_WRITECOPY);
		bool b = !(mbi.Protect & mask);

		if (mbi.Protect & (PAGE_GUARD | PAGE_NOACCESS))
			b = true;

		return b;
	}

	return true;
}