// Copyright 2017 Polymorph. All Rights Reserved.

#include "Dinkey.h"
#include "Interfaces/IPluginManager.h"
#include "Misc/Paths.h"

#include "dris.h"									// for DRIS definition + DDProtCheck declaration

DEFINE_LOG_CATEGORY(LogDinkeyPlugin);

#define DISABLE_PROTECTION 0

#define ASTERION_PS		13290
#define MY_PRODCODE     "PROSERVE"					// Dongle Product Code

void* UDinkey::DpWinDllHandle = nullptr;

void UDinkey::Initialize() {
	UE_LOG(LogDinkeyPlugin, Log, TEXT("UDinkey::Initialize()"));

	const TSharedPtr<IPlugin> Plugin = IPluginManager::Get().FindPlugin(TEXT("DinkeyPlugin"));
	if (!Plugin.IsValid()) {
		UE_LOG(LogDinkeyPlugin, Fatal, TEXT("DinkeyPlugin not found by IPluginManager"));
		return;
	}

	const FString DllPath = FPaths::Combine(
		Plugin->GetBaseDir(),
		TEXT("Binaries/ThirdParty/DinkeyPro/Win64/dpwin64.dll"));

	if (!FPaths::FileExists(DllPath)) {
		UE_LOG(LogDinkeyPlugin, Fatal, TEXT("Dinkey DLL not found at: %s"), *DllPath);
		return;
	}

	DpWinDllHandle = FPlatformProcess::GetDllHandle(*DllPath);
	if (DpWinDllHandle == nullptr) {
		UE_LOG(LogDinkeyPlugin, Fatal, TEXT("Failed to load Dinkey DLL: %s"), *DllPath);
	}
}

void UDinkey::Finalize() {
	UE_LOG(LogDinkeyPlugin, Log, TEXT("UDinkey::Finalize()"));

	if (DpWinDllHandle != nullptr) {
		FPlatformProcess::FreeDllHandle(DpWinDllHandle);
		DpWinDllHandle = nullptr;
	}
}

UDinkey::UDinkey(const FObjectInitializer& ObjectInitializer)
	: Super(ObjectInitializer) {
}

int UDinkey::ProtCheck() {
	DRIS dris;
	int ret_code = 0;

	// set the DRIS to have random values
	random_set(&dris, sizeof(DRIS));
	// then set the values we want to use
	memcpy(dris.header, "DRIS", 4);
	dris.size = sizeof(DRIS);
	dris.function = PROTECTION_CHECK;					// standard protection check
	dris.flags = 0;										// no extra flags, but you may want to specify some if you want to start a network user or decrement execs,...

	if (DpWinDllHandle == nullptr) {
		UE_LOG(LogDinkeyPlugin, Error, TEXT("Dinkey library not loaded!"));
#if DISABLE_PROTECTION || WITH_EDITOR
		return 0;
#else
		return -1;
#endif
	}

	ret_code = DDProtCheck(&dris, NULL);

	if (ret_code != 0) {
		DisplayError(ret_code, dris.ext_err);			// display messages for any errors that have occurred
#if DISABLE_PROTECTION || WITH_EDITOR
		ret_code = 0;
#endif
		return ret_code;
	}

	// later in your code you can check other values in the DRIS...

	if (dris.sdsn != ASTERION_PS ) {
		UE_LOG(LogDinkeyPlugin, Error, TEXT("Incorrect SDSN (%d) !"), dris.sdsn);
#if DISABLE_PROTECTION || WITH_EDITOR
		return 0;
#else
		return -1;
#endif
	}

	if (strcmp(dris.prodcode, MY_PRODCODE) != 0) {
		UE_LOG(LogDinkeyPlugin, Error, TEXT("Incorrect Product Code (%s) !"), ANSI_TO_TCHAR(dris.prodcode));
#if DISABLE_PROTECTION || WITH_EDITOR
		return 0;
#else
		return -1;
#endif
	}

	// later on in your program you can check the return code again
	if (dris.ret_code != 0) {
		UE_LOG(LogDinkeyPlugin, Error, TEXT("Dinkey Dongle protection error (%d) !"), dris.ret_code);
#if DISABLE_PROTECTION || WITH_EDITOR
		return 0;
#else
		return -1;
#endif
	}

	UE_LOG(LogDinkeyPlugin, Log, TEXT("License check suceeded."));

	return 0;
}

// *************  useful routines that are called from our functions  **********

// fills the buffer with length bytes of random data
// NB there are better ways of doing this (for example with the Windows Crypto API calls)
void UDinkey::random_set(void *buffer, int length) {
	int i;

//	srand(GetTickCount());

	for (i = 0; i<length; i++)
		((unsigned char *)buffer)[i] = (unsigned char)rand();

	return;
}

// displays messages for the most common errors. You will want to change this for your code
void UDinkey::DisplayError(int ret_code, int extended_error) {
	switch (ret_code) {
	case 401:
		UE_LOG(LogDinkeyPlugin, Error, TEXT("Error! No dongles detected!"));
		break;

	case 403:
		UE_LOG(LogDinkeyPlugin, Error, TEXT("Error! The dongle detected has a different type to the one specified in DinkeyAdd."));
		break;

	case 404:
		UE_LOG(LogDinkeyPlugin, Error, TEXT("Error! The dongle detected has a different model to those specified in DinkeyAdd."));
		break;

	case 409:
		UE_LOG(LogDinkeyPlugin, Error, TEXT("Error! The dongle detected has not been programmed by DinkeyAdd."));
		break;

	case 410:
		UE_LOG(LogDinkeyPlugin, Error, TEXT("Error! The dongle detected has a different Product Code to the one specified in DinkeyAdd."));
		break;

	case 411:
		UE_LOG(LogDinkeyPlugin, Error, TEXT("Error! The dongle detected does not contain the licence associated with this program."));
		break;

	case 413:
		UE_LOG(LogDinkeyPlugin, Error, TEXT("Error! This program has not been protected by DinkeyAdd. For guidance please read the DinkeyAdd chapter of the Dinkey manual."));
		break;

	case 417:
		UE_LOG(LogDinkeyPlugin, Error, TEXT("Error! One or more of the parameters set in the DRIS is incorrect. This could be caused if you are encrypting the DRIS in your code but did not specify DRIS encryption in DinkeyAdd - or vice versa."));
		break;

	case 423:
		UE_LOG(LogDinkeyPlugin, Error, TEXT("Error! The number of network users has been exceeded."));
		break;

	case 435:
		UE_LOG(LogDinkeyPlugin, Error, TEXT("Error! DinkeyServer has not been detected on the network."));
		break;

	case 922:
		UE_LOG(LogDinkeyPlugin, Error, TEXT("Error! The Software Key has expired."));
		break;

	default:
		UE_LOG(LogDinkeyPlugin, Error, TEXT("An error occurred checking the dongle.\nError: %d, Extended Error: %d"), ret_code, extended_error);
		break;
	}
	return;
}
