/* SPDX-License-Identifier: GPL-2.0-or-later */

#include <assert.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/stat.h>
#include <unistd.h>
#include <LIEF/ELF.hpp>
#include "nls-root_ca-certificates.hpp"

using namespace LIEF::ELF;

FILE *gridd_filep = NULL;
uint8_t *gridd_data = NULL;
FILE *cert_fp = NULL;
void *user_root_ca = NULL;

void usage(char *program_name)
{
	printf("Usage: %s [OPTIONS]\n", program_name);
	printf("\n"
	       "  -h\n"
	       "  -g <nvidia-gridd binary>\n"
	       "  -c <root CA certificate>\n"
	       "\n");
}

int parse_args(int argc, char *argv[])
{
	int opt;

	while ((opt = getopt(argc, argv, "hg:c:")) != -1) {
		// Required parameter is in global "optarg"
		switch (opt) {
		case '?':
		/* fall through */
		case 'h':
			usage(argv[0]);
			return -1;
		case 'g':
			gridd_filep = fopen(optarg, "rb+");
			break;
		case 'c':
			cert_fp = fopen(optarg, "rb");
			break;
		}
	}

	if (gridd_filep == NULL || cert_fp == NULL) {
		usage(argv[0]);
		return -1;
	}

	return 0;
}
void cleanup(void)
{
	if (user_root_ca)
		free(user_root_ca);
	if (gridd_data)
		free(gridd_data);
	if (cert_fp)
		fclose(cert_fp);
	if (gridd_filep)
		fclose(gridd_filep);
}

int main(int argc, char *argv[])
{
	struct stat gridd_fp_stats, cert_fp_stats;
	size_t status;

	if (parse_args(argc, argv) == -1) {
		cleanup();
		return -1;
	}

	/* Find the hardcoded certificates */
	// Read the nvidia-gridd binary
	fstat(fileno(gridd_filep), &gridd_fp_stats);
	gridd_data = (uint8_t *)malloc(gridd_fp_stats.st_size);
	assert(gridd_data != NULL);

	status = fread((void *)gridd_data, gridd_fp_stats.st_size, 1, gridd_filep);
	assert(status > 0);

	// Parse the ELF
	std::vector<uint8_t> gridd_vec(gridd_data, gridd_data + gridd_fp_stats.st_size);
	auto gridd_bin = Parser::parse(gridd_vec);
	auto s_rodata = gridd_bin->get_section(".rodata");
	auto s_data = gridd_bin->get_section(".data");

	// There are two hardcoded certificates
	size_t cert_one_offset = s_rodata->search(gridd_hardcoded_cert_one);
	size_t cert_two_offset = s_rodata->search(gridd_hardcoded_cert_two);
	printf("Found the two hardcoded NLS certificates at 0x%x and 0x%x.\n", cert_one_offset, cert_two_offset);

	/* Validate the size of the provided root CA */
	fstat(fileno(cert_fp), &cert_fp_stats);
	// FIXME: Align up to cover the padding?
	size_t total_cert_size = cert_two_offset + (strlen(gridd_hardcoded_cert_two) + 1) - cert_one_offset;
	if (cert_fp_stats.st_size > total_cert_size) {
		printf("The provided certificate (size %d) is larger than the available space (size %d)!\n",
			cert_fp_stats.st_size, total_cert_size);
		cleanup();
		return -1;
	}

	/* Patch in the provided certificate */
	// Read the provided root CA
	user_root_ca = malloc(cert_fp_stats.st_size);
	assert(user_root_ca != NULL);

	status = fread(user_root_ca, cert_fp_stats.st_size, 1, cert_fp);
	assert(status > 0);

	// Overwrite the first certificate, NULL the second
	uint8_t *cert_start = gridd_data + s_rodata->offset() + cert_one_offset;
	memcpy((void *)cert_start, user_root_ca, cert_fp_stats.st_size);
	memset((void *)(cert_start + cert_fp_stats.st_size), 0, total_cert_size - cert_fp_stats.st_size);

	printf("Replaced the hardcoded certificates with the provided one.\n");

	// Erase the XREF to the second certificate
	size_t cert_xrefs_array = s_data->offset() + s_data->search(s_rodata->virtual_address() + cert_one_offset);
	memset((void *)(gridd_data + cert_xrefs_array + sizeof(uint64_t)), 0, sizeof(uint64_t));

	printf("Fixed up the list of certificates. Done!\n");

	// Write the nvidia-gridd binary
	FILE *test = fopen("gridd_bin", "wb");
	fwrite((void *)gridd_data, gridd_fp_stats.st_size, 1, test);
	fclose(test);

	cleanup();
	return 0;
}