Main.cpp

/*
 * This file is part of the Rhocore-extrem distribution.
 * Copyright (c) 2024 Thomas Baumann.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, version 3.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program. If not, see <http://www.gnu.org/licenses/>.
 */

#include "Timer.h"
#include "KeyHunt.h"
#include "Base58.h"
#include "ArgParse.h"
#include "Int.h"
#include <fstream>
#include <string>
#include <string.h>
#include <stdexcept>
#include <cassert>
#ifndef WIN64
#include <signal.h>
#include <unistd.h>
#endif

#ifdef WITHGPU
#include <cuda.h>
#include <cuda_runtime.h>
#endif

#define RELEASE "1.00"

using namespace std;
using namespace argparse;
bool should_exit = false;

// ----------------------------------------------------------------------------
// ----------------------------------------------------------------------------

const char* vstr = "Print version                                                                                   ";
const char* cstr = "Check the working of the codes                                                                  ";
const char* ustr = "Search uncompressed addresses                                                                   ";
const char* bstr = "Search both uncompressed or compressed addresses                                                ";
const char* gstr = "Enable GPU calculation                                                                          ";
const char* istr = "GPU ids: 0,1...: List of GPU(s) to use, default is 0                                            ";
const char* xstr = "GPU gridsize: g0x,g0y,g1x,g1y, ...: Specify GPU(s) kernel gridsize, default is 16*(Device MP count),128";
const char* ostr = "Outputfile: Output results to the specified file, default: Found.txt                            ";
const char* mstr = "Specify maximun number of addresses found by each kernel call                                   ";
//const char* sstr = "Seed: Specify a seed for the base key, default is random                                        ";
const char* tstr = "threadNumber: Specify number of CPU thread, default is number of core                           ";
//const char* estr = "Disable SSE hash function                                                                       ";
const char* lstr = "List cuda enabled devices                                                                       ";
//const char* rstr = "Rkey: Rekey interval in MegaKey, default is disabled                                            ";
//const char* nstr = "Number of base key random bits                                                                  ";
const char* fstr = "Ripemd160 binary hash file path                                                                 ";
const char* astr = "P2PKH Address (single address mode)                                                             ";

const char* pstr = "Range start in hex                                                                              ";
const char* qstr = "Range end in hex, if not provided then, endRange would be set to the maximum value for secp256k1 curve";

// ----------------------------------------------------------------------------
// ----------------------------------------------------------------------------

void getInts(string name, vector<int>& tokens, const string& text, char sep)
{

	size_t start = 0, end = 0;
	tokens.clear();
	int item;

	try {

		while ((end = text.find(sep, start)) != string::npos) {
			item = std::stoi(text.substr(start, end - start));
			tokens.push_back(item);
			start = end + 1;
		}

		item = std::stoi(text.substr(start));
		tokens.push_back(item);

	}
	catch (std::invalid_argument&) {

		printf("Invalid %s argument, number expected\n", name.c_str());
		exit(-1);

	}

}

#ifdef WIN64
BOOL WINAPI CtrlHandler(DWORD fdwCtrlType)
{
	switch (fdwCtrlType) {
	case CTRL_C_EVENT:
		//printf("\n\nCtrl-C event\n\n");
		should_exit = true;
		return TRUE;

	default:
		return TRUE;
	}
}
#else
void CtrlHandler(int signum) {
	printf("\n\nBYE\n");
	exit(signum);
}
#endif

int main(int argc, const char* argv[])
{
	// Global Init
	Timer::Init();
	rseed(Timer::getSeed32());

	bool gpuEnable = false;
	bool gpuAutoGrid = true;
	int searchMode = SEARCH_COMPRESSED;
	vector<int> gpuId = {0,1,2,3,4,5,6,7};
	vector<int> gridSize;
	//string seed = "";
	string outputFile = "Found.txt";
	string addressFile = "";
	string address = "";
	//string addressHash = "";
	std::vector<unsigned char> addressHash;
	bool singleAddress = false;
	int nbCPUThread = Timer::getCoreNumber();
	//int nbit = 0;
	bool tSpecified = false;
	bool sse = true;
	uint32_t maxFound = 1024 * 256;  // Erhöht von 64K auf 256K
	//uint64_t rekey = 0;
	//bool paranoiacSeed = false;
	string rangeStart = "";
	string rangeEnd = "";
	addressHash.clear();

	ArgumentParser parser("Rhocore-extrem", "Hunt for Bitcoin private keys with extreme performance.");
	parser.add_argument("-v", "--version", vstr, false);
	parser.add_argument("-c", "--check", cstr, false);
	parser.add_argument("-u", "--uncomp", ustr, false);
	parser.add_argument("-b", "--both", bstr, false);
	parser.add_argument("-g", "--gpu", gstr, false);
	parser.add_argument("-i", "--gpui", istr, false);
	parser.add_argument("-x", "--gpux", xstr, false);
	parser.add_argument("-o", "--out", ostr, false);
	parser.add_argument("-m", "--max", mstr, false);
	parser.add_argument("-t", "--thread", tstr, false);
	//parser.add_argument("-e", "--nosse", estr, false);
	parser.add_argument("-l", "--list", lstr, false);
	//parser.add_argument("-r", "--rkey", rstr, false);
	//parser.add_argument("-n", "--nbit", nstr, false);
	parser.add_argument("-f", "--file", fstr, false);
	parser.add_argument("-a", "--addr", astr, false);

	parser.add_argument("-s", "--start", pstr, false);
	parser.add_argument("-e", "--end", qstr, false);

	parser.enable_help();

	auto err = parser.parse(argc, argv);
	if (err) {
		std::cout << err << std::endl;
		parser.print_help();
		return -1;
	}

	if (parser.exists("help")) {
		parser.print_help();
		return 0;
	}

	if (parser.exists("version")) {
		printf("Rhocore-extrem v" RELEASE "\n");
		return 0;
	}

	if (parser.exists("check")) {
		printf("Rhocore-extrem v" RELEASE "\n\n");

		printf("\nChecking... Secp256K1\n\n");
		Secp256K1 sec;
		sec.Init();
		sec.Check();

		printf("\n\nChecking... Int\n\n");
		Int K;
		K.SetBase16("3EF7CEF65557B61DC4FF2313D0049C584017659A32B002C105D04A19DA52CB47");
		K.Check();

		printf("\n\nCheck completed successfully\n\n");
		return 0;
	}

	if (parser.exists("uncomp")) {
		searchMode = SEARCH_UNCOMPRESSED;
	}
	if (parser.exists("both")) {
		searchMode = SEARCH_BOTH;
	}

	if (parser.exists("gpu")) {
		gpuEnable = true;
		nbCPUThread = 0;
	}

	if (parser.exists("gpui")) {
		string ids = parser.get<string>("i");
		getInts("gpui", gpuId, ids, ',');
	}

	if (parser.exists("gpux")) {
		string grids = parser.get<string>("x");
		getInts("gpux", gridSize, grids, ',');
		gpuAutoGrid = false;
	}

	if (parser.exists("out")) {
		outputFile = parser.get<string>("o");
	}

	if (parser.exists("max")) {
		maxFound = parser.get<uint32_t>("m");
	}

	//if (parser.exists("seed")) {
	//	seed = parser.get<string>("s");
	//	paranoiacSeed = true;
	//}

	if (parser.exists("thread")) {
		nbCPUThread = parser.get<int>("t");
		tSpecified = true;
	}

	//if (parser.exists("nosse")) {
	//	sse = false;
	//}

	if (parser.exists("list")) {
#ifdef WIN64
		GPUEngine::PrintCudaInfo();
#else
		printf("GPU code not compiled, use -DWITHGPU when compiling.\n");
#endif
		return 0;
	}

	//if (parser.exists("rkey")) {
	//	rekey = parser.get<uint64_t>("r");
	//}

	//if (parser.exists("nbit")) {
	//	nbit = parser.get<int>("n");
	//	if (nbit < 1 || nbit > 256) {
	//		printf("Invalid nbit value, must have in range: 1 - 256\n");
	//		exit(-1);
	//	}
	//}

	if (parser.exists("file")) {
		addressFile = parser.get<string>("f");
	}

	if (parser.exists("addr")) {
		address = parser.get<string>("a");
		if (address.length() < 30 || address[0] != '1') {
			printf("Invalid addr argument, must have P2PKH address only\n");
			exit(-1);
		}
		else {
			if (DecodeBase58(address, addressHash)) {
				addressHash.erase(addressHash.begin() + 0);
				addressHash.erase(addressHash.begin() + 20, addressHash.begin() + 24);
				assert(addressHash.size() == 20);
			}
		}
	}

	if (parser.exists("start")) {
		rangeStart = parser.get<string>("s");
		// Validate that rangeStart is a valid hex string
		for (char c : rangeStart) {
			if (!((c >= '0' && c <= '9') || (c >= 'a' && c <= 'f') || (c >= 'A' && c <= 'F'))) {
				printf("Error: Invalid character '%c' in start range. Must be valid hex.\n", c);
				exit(-1);
			}
		}
	}

	if (parser.exists("end")) {
		rangeEnd = parser.get<string>("e");
		// Validate that rangeEnd is a valid hex string
		for (char c : rangeEnd) {
			if (!((c >= '0' && c <= '9') || (c >= 'a' && c <= 'f') || (c >= 'A' && c <= 'F'))) {
				printf("Error: Invalid character '%c' in end range. Must be valid hex.\n", c);
				exit(-1);
			}
		}
	}

	// Validate that if both ranges are provided, end is not smaller than start
	if (rangeStart.length() > 0 && rangeEnd.length() > 0) {
		Int startInt, endInt;
		startInt.SetBase16(rangeStart.c_str());
		endInt.SetBase16(rangeEnd.c_str());

		if (!endInt.IsGreaterOrEqual(&startInt)) {
			printf("Error: End range (%s) is smaller than start range (%s)\n", rangeEnd.c_str(), rangeStart.c_str());
			exit(-1);
		}
	}


	if (gridSize.size() == 0) {
		for (int i = 0; i < gpuId.size(); i++) {
			// Für RTX 4090 optimiert: 16 * MP count (ca. 128) = 2048
			gridSize.push_back(-1);  // Wird automatisch basierend auf MP count berechnet
			gridSize.push_back(128); // Blockgröße bleibt bei 128
		}
	}
	else if (gridSize.size() != gpuId.size() * 2) {
		printf("Invalid gridSize or gpuId argument, must have coherent size\n");
		exit(-1);
	}

	if ((addressHash.size() <= 0) && (addressFile.length() <= 0)) {
		printf("Invalid ripemd160 binary hash file path or invalid address\n");
		exit(-1);
	}

	if ((addressHash.size() > 0) && (addressFile.length() > 0)) {
		printf("Invalid arguments, addr and file, both option can't be used together\n");
		exit(-1);
	}

	if (rangeStart.length() <= 0) {
		printf("Error: Start range must be provided.\n");
		printf("Example: -s 80000000 or --start 80000000\n");
		printf("For maximum range: -s 1\n");
		exit(-1);
	}

	//if (rangeStart.length() > 0 && nbit > 0) {
	//	printf("Invalid arguments, nbit and ranges, both can't be used together\n");
	//	exit(-1);
	//}

	if (nbCPUThread > 0 && gpuEnable) {
		printf("Invalid arguments, CPU and GPU, both can't be used together right now\n");
		exit(-1);
	}

	// Überprüfen, ob alle angegebenen GPU-IDs gültig sind
#ifdef WITHGPU
	if (gpuEnable) {
		int deviceCount = 0;
		cudaGetDeviceCount(&deviceCount);

		for (int i = 0; i < gpuId.size(); i++) {
			if (gpuId[i] >= deviceCount) {
				printf("Invalid GPU ID %d, only %d GPU(s) available\n", gpuId[i], deviceCount);
				exit(-1);
			}
		}

		// Überprüfen, ob mindestens eine GPU gefunden wurde
		if (deviceCount == 0) {
			printf("No CUDA capable devices found\n");
			exit(-1);
		}

		printf("Found %d CUDA capable device(s)\n", deviceCount);
	}
#endif

	// Let one CPU core free per gpu is gpu is enabled
	// It will avoid to hang the system
	if (!tSpecified && nbCPUThread > 1 && gpuEnable)
		nbCPUThread -= (int)gpuId.size();
	if (nbCPUThread < 0)
		nbCPUThread = 0;

	{
		printf("\n");
		printf("Rhocore-extrem v" RELEASE "\n");
		printf("\n");
		printf("MODE         : %s\n", searchMode == SEARCH_COMPRESSED ? "COMPRESSED" : (searchMode == SEARCH_UNCOMPRESSED ? "UNCOMPRESSED" : "COMPRESSED & UNCOMPRESSED"));
		printf("DEVICE       : %s\n", (gpuEnable && nbCPUThread > 0) ? "CPU & GPU" : ((!gpuEnable && nbCPUThread > 0) ? "CPU" : "GPU"));
		printf("CPU THREAD   : %d\n", nbCPUThread);
		printf("GPU IDS      : ");
		for (int i = 0; i < gpuId.size(); i++) {
			printf("%d", gpuId.at(i));
			if (i + 1 < gpuId.size())
				printf(", ");
		}
		printf("\n");
		printf("GPU GRIDSIZE : ");
		for (int i = 0; i < gridSize.size(); i++) {
			printf("%d", gridSize.at(i));
			if (i + 1 < gridSize.size()) {
				if ((i + 1) % 2 != 0) {
					printf("x");
				}
				else {
					printf(", ");
				}
			}
		}
		if (gpuAutoGrid)
			printf(" (grid size will be calculated automatically based on multiprocessor number on GPU device)\n");
		else
			printf("\n");
		printf("SSE          : %s\n", sse ? "YES" : "NO");
		printf("MAX FOUND    : %d\n", maxFound);
		if (addressFile.length() > 0)
			printf("HASH160 FILE : %s\n", addressFile.c_str());
		else
			printf("ADDRESS      : %s (single address mode)\n", address.c_str());
		printf("OUTPUT FILE  : %s\n", outputFile.c_str());
	}
#ifdef WIN64
	if (SetConsoleCtrlHandler(CtrlHandler, TRUE)) {
		KeyHunt* v = new KeyHunt(addressFile, addressHash, searchMode, gpuEnable,
			outputFile, sse, maxFound, rangeStart, rangeEnd, should_exit);

		v->Search(nbCPUThread, gpuId, gridSize, should_exit);

		delete v;
		printf("\n\nBYE\n");
		return 0;
	}
	else {
		printf("error: could not set control-c handler\n");
		return 1;
	}
#else
	signal(SIGINT, CtrlHandler);
	KeyHunt* v = new KeyHunt(addressFile, addressHash, searchMode, gpuEnable,
		outputFile, sse, maxFound, rangeStart, rangeEnd, should_exit);

	v->Search(nbCPUThread, gpuId, gridSize, should_exit);

	delete v;
	return 0;
#endif
}