YubiKey support + stability fix & bug fix (#17)

.gitignore

@@ -364,3 +364,9 @@ MigrationBackup/
 # Fody - auto-generated XML schema
 FodyWeavers.xsd
 /Final Release
+/PROD-YUBI
+/PROD-YUBI-NOSTARTUP
+/ProtectionPayload/PROD-YUBI
+/ProtectionPayload/PROD-YUBI-NOSTARTUP
+/DiscordTokenProtector/PROD-YUBI-NOSTARTUP
+/DiscordTokenProtector/PROD-YUBI

DiscordTokenProtector.sln

@@ -15,6 +15,10 @@ Global
 		PROD|x86 = PROD|x86
 		PROD-NOSTARTUP|x64 = PROD-NOSTARTUP|x64
 		PROD-NOSTARTUP|x86 = PROD-NOSTARTUP|x86
+		PROD-YUBI|x64 = PROD-YUBI|x64
+		PROD-YUBI|x86 = PROD-YUBI|x86
+		PROD-YUBI-NOSTARTUP|x64 = PROD-YUBI-NOSTARTUP|x64
+		PROD-YUBI-NOSTARTUP|x86 = PROD-YUBI-NOSTARTUP|x86
 		Release|x64 = Release|x64
 		Release|x86 = Release|x86
 	EndGlobalSection
@@ -31,6 +35,14 @@ Global
 		{83DABD03-F840-4BF5-A090-859F73A3E9EF}.PROD-NOSTARTUP|x64.Build.0 = PROD-NOSTARTUP|x64
 		{83DABD03-F840-4BF5-A090-859F73A3E9EF}.PROD-NOSTARTUP|x86.ActiveCfg = PROD-NOSTARTUP|Win32
 		{83DABD03-F840-4BF5-A090-859F73A3E9EF}.PROD-NOSTARTUP|x86.Build.0 = PROD-NOSTARTUP|Win32
+		{83DABD03-F840-4BF5-A090-859F73A3E9EF}.PROD-YUBI|x64.ActiveCfg = PROD-YUBI|x64
+		{83DABD03-F840-4BF5-A090-859F73A3E9EF}.PROD-YUBI|x64.Build.0 = PROD-YUBI|x64
+		{83DABD03-F840-4BF5-A090-859F73A3E9EF}.PROD-YUBI|x86.ActiveCfg = PROD-YUBI|Win32
+		{83DABD03-F840-4BF5-A090-859F73A3E9EF}.PROD-YUBI|x86.Build.0 = PROD-YUBI|Win32
+		{83DABD03-F840-4BF5-A090-859F73A3E9EF}.PROD-YUBI-NOSTARTUP|x64.ActiveCfg = PROD-YUBI-NOSTARTUP|x64
+		{83DABD03-F840-4BF5-A090-859F73A3E9EF}.PROD-YUBI-NOSTARTUP|x64.Build.0 = PROD-YUBI-NOSTARTUP|x64
+		{83DABD03-F840-4BF5-A090-859F73A3E9EF}.PROD-YUBI-NOSTARTUP|x86.ActiveCfg = PROD-YUBI-NOSTARTUP|Win32
+		{83DABD03-F840-4BF5-A090-859F73A3E9EF}.PROD-YUBI-NOSTARTUP|x86.Build.0 = PROD-YUBI-NOSTARTUP|Win32
 		{83DABD03-F840-4BF5-A090-859F73A3E9EF}.Release|x64.ActiveCfg = Release|x64
 		{83DABD03-F840-4BF5-A090-859F73A3E9EF}.Release|x64.Build.0 = Release|x64
 		{83DABD03-F840-4BF5-A090-859F73A3E9EF}.Release|x86.ActiveCfg = Release|Win32
@@ -47,6 +59,14 @@ Global
 		{D0E45CD8-EEB4-4448-9E50-893266B0CED9}.PROD-NOSTARTUP|x64.Build.0 = PROD|x64
 		{D0E45CD8-EEB4-4448-9E50-893266B0CED9}.PROD-NOSTARTUP|x86.ActiveCfg = PROD|Win32
 		{D0E45CD8-EEB4-4448-9E50-893266B0CED9}.PROD-NOSTARTUP|x86.Build.0 = PROD|Win32
+		{D0E45CD8-EEB4-4448-9E50-893266B0CED9}.PROD-YUBI|x64.ActiveCfg = PROD-YUBI|x64
+		{D0E45CD8-EEB4-4448-9E50-893266B0CED9}.PROD-YUBI|x64.Build.0 = PROD-YUBI|x64
+		{D0E45CD8-EEB4-4448-9E50-893266B0CED9}.PROD-YUBI|x86.ActiveCfg = PROD-YUBI|Win32
+		{D0E45CD8-EEB4-4448-9E50-893266B0CED9}.PROD-YUBI|x86.Build.0 = PROD-YUBI|Win32
+		{D0E45CD8-EEB4-4448-9E50-893266B0CED9}.PROD-YUBI-NOSTARTUP|x64.ActiveCfg = PROD-YUBI-NOSTARTUP|x64
+		{D0E45CD8-EEB4-4448-9E50-893266B0CED9}.PROD-YUBI-NOSTARTUP|x64.Build.0 = PROD-YUBI-NOSTARTUP|x64
+		{D0E45CD8-EEB4-4448-9E50-893266B0CED9}.PROD-YUBI-NOSTARTUP|x86.ActiveCfg = PROD-YUBI-NOSTARTUP|Win32
+		{D0E45CD8-EEB4-4448-9E50-893266B0CED9}.PROD-YUBI-NOSTARTUP|x86.Build.0 = PROD-YUBI-NOSTARTUP|Win32
 		{D0E45CD8-EEB4-4448-9E50-893266B0CED9}.Release|x64.ActiveCfg = Release|x64
 		{D0E45CD8-EEB4-4448-9E50-893266B0CED9}.Release|x64.Build.0 = Release|x64
 		{D0E45CD8-EEB4-4448-9E50-893266B0CED9}.Release|x86.ActiveCfg = Release|Win32

DiscordTokenProtector/Context.h

@@ -63,8 +63,11 @@ class Context {
 			g_logger.warning("Tried to stop the protection when it\'s already stopped.");
 			return;
 		}
+
+		m_protectionState = ProtectionStates::Stop;
 		Discord::killDiscord();
 		m_running = false;
+
 		m_protectionThread.join();
 	}
 
@@ -97,7 +100,9 @@ class Context {
 			else {
 				kd = HWID_kd;
 				state = State::TokenSecure;
-				startProtection();
+
+				if (g_config->read<bool>("auto_start"))
+					startProtection();
 			}
 		}
 		else {
@@ -146,13 +151,27 @@ class Context {
 #endif
 	}
 
+	std::string getCurrentStateString() {
+		static const std::map<ProtectionStates, std::string> stateStrings = {
+			{ProtectionStates::Idle, "Waiting for Discord..."},
+			{ProtectionStates::Starting, "Starting..."},
+			{ProtectionStates::Checking, "Checking the integrity of Discord."},
+			{ProtectionStates::CheckIssues, "Found issues."},
+			{ProtectionStates::Injecting, "Injecting payload..."},
+			{ProtectionStates::Connected, "Connected!"},
+			{ProtectionStates::Stop, "Stopping..."}
+		};
+
+		return stateStrings.find(m_protectionState)->second;
+	}
+
 	//TODO getter setter ?
 
 	State state = State::None;
 	ProtectionStates m_protectionState = ProtectionStates::Idle;
 
 	EncryptionType encryptionType_cache = EncryptionType::Unknown;//TODO remove this
-	KeyData kd;//TODO THIS IS NOT SECURED! but we can't ask the users password all the time.
+	KeyData kd;
 
 	FolderRemover remover_LocalStorage;
 	FolderRemover remover_SessionStorage;
@@ -162,7 +181,7 @@ class Context {
 
 	IntegrityCheck integrityCheck;
 
-	bool m_running = false;
+	std::atomic_bool m_running = false;
 	std::mutex m_threadMutex;
 
 	bool m_isAutoStarting = false;
@@ -179,6 +198,8 @@ class Context {
 		while (m_protectionState == ProtectionStates::Connected) {
 			try {
 				std::string msg = m_networkManager.Recv();
+				if (msg == "KeepAlive") continue;
+
 				json jsonMsg = json::parse(msg);
 
 				if (jsonMsg["code"] == "HANDOFF") {
@@ -251,23 +272,28 @@ class Context {
 				}
 
 				m_protectionState = ProtectionStates::Starting;
+				hasStartedDiscord = true;
 
 				stopRemovers();
 
+				if (m_protectionState == ProtectionStates::Stop) continue;
+
 				Discord::killDiscord();
 				remover_LocalStorage.Remove();
 				remover_canary_LocalStorage.Remove();
 
 				//Check before launching!
 				if (g_config->read<bool>("integrity")) {
 					m_protectionState = ProtectionStates::Checking;
+					if (m_protectionState == ProtectionStates::Stop) continue;
 
 					integrityCheck.setCheckHash(g_config->read<bool>("integrity_checkhash"));
 					integrityCheck.setCheckExecutableSig(g_config->read<bool>("integrity_checkexecutable"));
 					integrityCheck.setCheckModule(g_config->read<bool>("integrity_checkmodule"));
 					integrityCheck.setCheckResources(g_config->read<bool>("integrity_checkresource"));
 					integrityCheck.setCheckScripts(g_config->read<bool>("integrity_checkscripts"));
 					integrityCheck.setAllowBetterDiscord(g_config->read<bool>("integrity_allowbetterdiscord"));
+					integrityCheck.setRedownloadHashes(g_config->read<bool>("integrity_redownloadhashes"));
 
 					integrityCheck.setDiscordVersion(g_discord->getDiscordVersion(discordType));
 
@@ -285,13 +311,17 @@ class Context {
 					}
 				}
 
+				if (m_protectionState == ProtectionStates::Stop) continue;
+
 				m_protectionState = ProtectionStates::Injecting;
 
 				PROCESS_INFORMATION discordProcess = g_discord->startSuspendedDiscord(discordType);//TODO CloseHandle?
 				//g_processprotection->ProtectProcess(discordProcess.hProcess);//TODO Fix
 
 				//TODO make this thing async
 				try {
+					if (m_protectionState == ProtectionStates::Stop) continue;
+
 					std::promise<USHORT> portPromise;
 
 					auto start = std::async(std::launch::async, &NetworkManager::Start, &m_networkManager, std::ref(portPromise));
@@ -303,11 +333,13 @@ class Context {
 
 					std::cout << "Injected!" << std::endl;
 
+					if (m_protectionState == ProtectionStates::Stop) continue;
+
 					if (ResumeThread(discordProcess.hThread) == -1)
 						throw std::runtime_error(sf() << "Failed ResumeThread : " << GetLastError());
 
 					//Wait for the payload client
-					for (int i = 0; i < 60 * 10; i++) {//Wait 60 seconds
+					for (int i = 0; (i < 60 * 10 && m_protectionState != ProtectionStates::Stop); i++) {//Wait 60 seconds
 						if (start.wait_for(std::chrono::milliseconds(100)) == std::future_status::ready)
 							break;
 					}
@@ -316,6 +348,8 @@ class Context {
 						throw std::exception(("Discord payload timeout!"));
 					}
 
+					if (m_protectionState == ProtectionStates::Stop) continue;
+
 					std::cout << "Finish!" << std::endl;
 
 					if (start.valid())
@@ -349,8 +383,6 @@ class Context {
 					if (m_networkHandlerThread.joinable()) m_networkHandlerThread.join();
 
 					m_networkHandlerThread = std::thread(&Context::networkHandler, this);
-
-					hasStartedDiscord = true;
 				}
 				catch (std::exception& e) {
 					g_logger.error(sf() << __FUNCSIG__ " : Failed to load payload : " << e.what());

DiscordTokenProtector/Crypto/Crypto.cpp

@@ -1,5 +1,6 @@
 #include "Crypto.h"
 #include "CryptoUtils.h"
+#include "../Storage/Config.h"
 
 namespace Crypto {
 	//Just some random bytes as salt
@@ -124,4 +125,133 @@ namespace Crypto {
 
 		return out;
 	}
+
+	void encryptSBB(CryptoPP::SecByteBlock& in) {
+		if (auto mod = in.size() % CRYPTPROTECTMEMORY_BLOCK_SIZE; mod != 0)
+			in.CleanGrow(in.size() + (CRYPTPROTECTMEMORY_BLOCK_SIZE - mod));
+
+		if (!CryptProtectMemory(in.data(), in.size(), CRYPTPROTECTMEMORY_SAME_PROCESS))
+			throw std::runtime_error(sf() << "CryptProtectMemory failed : " << GetLastError());
+	}
+
+	void decryptSBB(CryptoPP::SecByteBlock& in) {
+		if (!CryptUnprotectMemory(in.data(), in.size(), CRYPTPROTECTMEMORY_SAME_PROCESS))
+			throw std::runtime_error(sf() << "CryptProtectMemory failed : " << GetLastError());
+	}
+
+#ifdef YUBIKEYSUPPORT
+	Yubi::Yubi() {
+		throwOnError(ykpiv_init(&m_state, true), "ykpiv_init");
+		throwOnError(ykpiv_connect(m_state, NULL), "ykpiv_connect");
+
+		m_model = ykpiv_util_devicemodel(m_state);
+		if (m_model == NULL)
+			throw std::runtime_error(sf() << "Failed ykpiv_util_devicemodel : model was NULL");
+		if (m_model == DEVTYPE_NEOr3)
+			throw std::runtime_error(sf() << "NEO YubiKey are not supported");
+	}
+
+	Yubi::~Yubi() {
+		if (m_state) {
+			ykpiv_disconnect(m_state);
+			ykpiv_done(m_state);
+		}
+	}
+
+	int Yubi::getRetryCount() {
+		int retries = 0;
+		throwOnError(ykpiv_get_pin_retries(m_state, &retries), "ykpiv_get_pin_retries");
+		return retries;
+	}
+
+	int Yubi::authentificate(const secure_string& pin) {
+		int tries = 0;
+		m_err = ykpiv_verify(m_state, pin.c_str(), &tries);
+		if (m_err == YKPIV_OK || m_err == YKPIV_WRONG_PIN)
+			return tries;
+		throwOnError(m_err, "ykpiv_verify");
+		return 0;//Should not hit this
+	}
+
+	secure_string Yubi::signData(const CryptoPP::SecByteBlock& in) {
+		secure_string signature;
+		signature.resize(1024);
+
+		size_t siglen = signature.size();
+
+		throwOnError(
+			ykpiv_sign_data(
+				m_state, in.data(), in.size(),
+				reinterpret_cast<byte*>(signature.data()), &siglen, YKPIV_ALGO_RSA2048, YKPIV_KEY_CARDAUTH
+			), "ykpiv_sign_data");
+
+		signature.resize(siglen);
+		return signature;
+	}
+
+	std::string Yubi::getModelName() const {
+		switch (m_model){
+			case DEVTYPE_NEO: return "YubiKey NEO";
+			case DEVTYPE_YK: return "YubiKey";
+			case DEVTYPE_NEOr3: return "YubiKey NEO R3";
+			case DEVTYPE_YK4: return "YubiKey 4";
+			case DEVTYPE_YK5: return "YubiKey 5";
+		}
+		return "Unknown";
+	}
+
+	CryptoPP::SecByteBlock Yubi::generateKeyFile() {
+		CryptoPP::SecByteBlock keydata = CryptoUtils::randomSBB(YUBIKEY_DATA_LEN);
+		std::ofstream file(Config::getConfigPath() + YUBIKEY_KEY_FILE, std::ios::binary);
+		if (!file.is_open())
+			throw std::runtime_error("generateKeyFile : Failed to open YubiKey data file");
+
+		file.write(reinterpret_cast<const char*>(keydata.data()), keydata.size());
+		return keydata;
+	}
+
+	CryptoPP::SecByteBlock Yubi::readKeyFile() {
+		std::ifstream file(Config::getConfigPath() + YUBIKEY_KEY_FILE, std::ios::binary);
+		if (!file.is_open())
+			throw std::runtime_error("readKeyFile : Failed to open YubiKey data file");
+
+		CryptoPP::SecByteBlock keydata(YUBIKEY_DATA_LEN);
+		file.read(reinterpret_cast<char*>(keydata.data()), YUBIKEY_DATA_LEN);
+
+		return keydata;
+	}
+
+	void Yubi::throwOnError(ykpiv_rc err, const std::string& action) {
+		m_err = err;
+		if (err == YKPIV_OK) return;
+
+		std::string errorString = "UNKNOWN_ERROR";
+
+#define ADDCASE(ERR)\
+	case ERR: errorString = #ERR; break;
+
+		switch (err) {
+			ADDCASE(YKPIV_MEMORY_ERROR);
+			ADDCASE(YKPIV_PCSC_ERROR);
+			ADDCASE(YKPIV_SIZE_ERROR);
+			ADDCASE(YKPIV_APPLET_ERROR);
+			ADDCASE(YKPIV_AUTHENTICATION_ERROR);
+			ADDCASE(YKPIV_RANDOMNESS_ERROR);
+			ADDCASE(YKPIV_GENERIC_ERROR);
+			ADDCASE(YKPIV_KEY_ERROR);
+			ADDCASE(YKPIV_PARSE_ERROR);
+			ADDCASE(YKPIV_WRONG_PIN);
+			ADDCASE(YKPIV_INVALID_OBJECT);
+			ADDCASE(YKPIV_ALGORITHM_ERROR);
+			ADDCASE(YKPIV_PIN_LOCKED);
+			ADDCASE(YKPIV_ARGUMENT_ERROR);
+			ADDCASE(YKPIV_RANGE_ERROR);
+			ADDCASE(YKPIV_NOT_SUPPORTED);
+		}
+
+#undef ADDCASE
+
+		throw std::runtime_error(sf() << "Failed " << action << " : " << errorString);
+	}
+#endif
 }

DiscordTokenProtector/Crypto/Crypto.h

@@ -11,6 +11,12 @@
 #include <cryptopp/eax.h>
 #include <cryptopp/hex.h>
 
+#ifdef YUBIKEYSUPPORT
+#include <ykpiv/ykpiv.h>
+#pragma comment(lib, "ykpiv_static.lib")
+#pragma comment(lib, "Winscard.lib")
+#endif
+
 //https://stackoverflow.com/a/56888301/13544464
 using secure_string = std::basic_string<char, std::char_traits<char>, CryptoPP::AllocatorWithCleanup<char>>;
 
@@ -22,4 +28,39 @@ namespace Crypto {
 
 	secure_string encryptHWID(const secure_string& content);
 	secure_string decryptHWID(const secure_string& content);
+
+	void encryptSBB(CryptoPP::SecByteBlock& in);
+	void decryptSBB(CryptoPP::SecByteBlock& in);
+
+#ifdef YUBIKEYSUPPORT
+	constexpr auto YUBIKEY_KEY_FILE = L"yk.dat";
+	constexpr auto YUBIKEY_DATA_LEN = 256;
+
+	class Yubi {
+	public:
+		Yubi();
+		~Yubi();
+
+		int getRetryCount();
+
+		//Returns -1 if authentificated, else it returns the amount of retries
+		int authentificate(const secure_string& pin);
+		//in.size() == 256
+		secure_string signData(const CryptoPP::SecByteBlock& in);
+
+		ykpiv_rc getLastError() const { return m_err; }
+
+		std::string getModelName() const;
+
+		static CryptoPP::SecByteBlock generateKeyFile();
+		static CryptoPP::SecByteBlock readKeyFile();
+
+	private:
+		void throwOnError(ykpiv_rc err, const std::string& action);
+
+		ykpiv_state* m_state;
+		ykpiv_rc m_err;
+		ykpiv_devmodel m_model = NULL;
+	};
+#endif
 }