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

#include <stdio.h>
#include <string.h>
#include <stdint.h>
#include <stdlib.h>

#ifdef VERIFY_SIG
#include <openssl/obj_mac.h>
#include <openssl/rsa.h>
#include <openssl/sha.h>
#endif

void usage(void)
{
	fprintf(stderr,
		"Usage:\n"
		"\tmec1322 dump [-s <flash size>] [-o <outfile>] <infile>\n"
		"\tmec1322 insert [-s <flash size>] [-f <image>] <infile> <header_loc>\n");

	exit(1);
}

const int SPI_CLOCK_LIST[] = {48, 24, 12, 8};
const uint8_t SPI_READ_CMD_LIST[] = {0x3, 0xb, 0x3b};

// FIXME: uint32_t cannot be used directly on big endian systems
struct mec1322_header {
	uint8_t sig[6]; // "CSMS\0\0"
	uint8_t spi_clk;
	uint8_t spi_read_cmd;
	uint32_t load_addr;
	uint32_t entry_point;
	uint16_t payload_blks; // payload_len >> 6
	uint8_t padding0[2];
	uint32_t payload_offset;
	uint8_t padding1[8];
	uint8_t payload_key_exp[16];	  // @0x20
	uint8_t payload_key_modulus[256]; // @0x30
};

uint8_t crc8(const uint8_t *data, size_t len)
{
	const uint8_t CRC_TABLE[] = {0x00, 0x07, 0x0e, 0x09, 0x1c, 0x1b, 0x12, 0x15,
				     0x38, 0x3f, 0x36, 0x31, 0x24, 0x23, 0x2a, 0x2d};
	uint8_t crc = 0;
	for (size_t i = 0; i < len; i++) {
		uint8_t v = data[i];
		crc = (crc << 4) ^ (CRC_TABLE[(crc >> 4) ^ (v >> 4)]);
		crc = (crc << 4) ^ (CRC_TABLE[(crc >> 4) ^ (v & 0xf)]);
	}
	return crc ^ 0x55;
}

// return the total length of the firmware, including header,
// header_signature, and payload
// return 0 if there's error
size_t parse_header(const struct mec1322_header *hdr)
{
	if (memcmp(hdr->sig, "CSMS", 4) != 0) {
		fprintf(stderr, "MEC1322 header signature error!\n");
		return 0;
	}

	printf("SPI clock = %hhd MHz\n", SPI_CLOCK_LIST[hdr->spi_clk]);
	printf("SPI read cmd = 0x%hhx\n", SPI_READ_CMD_LIST[hdr->spi_read_cmd]);
	printf("load address = 0x%x\n", hdr->load_addr);
	printf("entry point = 0x%x\n", hdr->entry_point);
	size_t payload_len = (size_t)hdr->payload_blks * 64;
	printf("payload length = 0x%lx\n", payload_len);
	printf("payload offset = 0x%x\n", hdr->payload_offset);
	return payload_len + hdr->payload_offset;
}

int verify_payload_signature(const struct mec1322_header *hdr, void *payload,
			     const uint8_t *sig)
{
#ifdef VERIFY_SIG
	uint8_t hash[32];
	uint8_t sigrev[256];

	for (size_t i = 0; i < 256; i++)
		sigrev[i] = sig[255 - i];

	RSA *rsa = RSA_new();
	BIGNUM *modulus = BN_lebin2bn(hdr->payload_key_modulus, 256, NULL);
	BIGNUM *exp = BN_lebin2bn(hdr->payload_key_exp, 8, NULL);
	RSA_set0_key(rsa, modulus, exp, NULL);

	SHA256(payload, (size_t)(hdr->payload_blks) * 64, hash);

	int v = RSA_verify(NID_sha256, hash, 32, sigrev, 256, rsa);
	if (!v)
		fprintf(stderr, "MEC1322 payload RSA signature error!\n");
	else
		printf("MEC1322 payload RSA signature matches the payload.\n");

	return v;
#else
	fprintf(stderr,
		"** MEC1322 payload signature verification is not enabled.\n"
		"** You can run `make VERIFY_SIG=1` to enable it.\n");
	return 1;
#endif
}

int mec1322_dump(int argc, char *argv[])
{
	FILE *infile, *outfile = NULL;
	long spisize = 0;

	while (argc > 2) {
		if (strcmp(argv[0], "-s") == 0) {
			spisize = strtol(argv[1], NULL, 0);
			argc -= 2;
			argv += 2;
			continue;
		}
		if (strcmp(argv[0], "-o") == 0) {
			outfile = fopen(argv[1], "w");
			if (outfile == NULL) {
				fprintf(stderr, "Fail to open file %s!\n", argv[1]);
				exit(1);
			}
			argc -= 2;
			argv += 2;
			continue;
		}
		// otherwise, print usage and exit
		usage();
	}

	if (argc != 1)
		usage();

	// without the above options, it's the input file
	infile = fopen(argv[0], "r");
	if (infile == NULL) {
		fprintf(stderr, "Fail to open file %s!\n", argv[0]);
		exit(1);
	}

	if (spisize == 0) {
		fseek(infile, -1, SEEK_END);
		spisize = ftell(infile) + 1;

		printf("Input file size is 0x%lx.\n", spisize);
	}

	// read MEC1322 firmware tag
	fseek(infile, -256, SEEK_END);
	uint8_t fwTag[4];
	uint32_t header_loc;
	fread(fwTag, 1, 4, infile);

	if (crc8(fwTag, 3) == fwTag[3]) {
		header_loc = (((uint32_t)fwTag[0]) << 8) | (((uint32_t)fwTag[1]) << 16)
			     | (((uint32_t)fwTag[2]) << 24);
		printf("MEC1322 firmware header location is 0x%x.\n", header_loc);
	} else {
		fprintf(stderr, "MEC1322 tag CRC8 error!\n");
		return 1;
	}

	// we use the distance of header_loc to end when seeking
	struct mec1322_header hdr;
	fseek(infile, header_loc - spisize, SEEK_END);
	fread(&hdr, 1, sizeof(hdr), infile);
	size_t fwlen = parse_header(&hdr);
	if (fwlen == 0) {
		fprintf(stderr, "Error parsing MEC1322 firmware header!\n");
		return 1;
	}

	if (outfile) {
		fseek(infile, header_loc - spisize, SEEK_END);
		void *image = malloc(fwlen);
		fread(image, 1, fwlen, infile);
		fwrite(image, 1, fwlen, outfile);

		uint8_t rsasig[256];
		fread(rsasig, 1, 256, infile);
		fwrite(rsasig, 1, 256, outfile);

		if (!verify_payload_signature(&hdr, image + hdr.payload_offset, rsasig))
			return 1;

		free(image);
		fclose(outfile);
	}

	fclose(infile);
	return 0;
}

int mec1322_insert(int argc, char *argv[])
{
	FILE *infile, *ecfile = NULL;
	void *image = NULL;
	size_t fwlen;
	long spisize = 0;
	long header_loc;

	if (argc < 2)
		usage();

	while (argc > 2) {
		if (strcmp(argv[0], "-s") == 0) {
			spisize = strtol(argv[1], NULL, 0);
			argc -= 2;
			argv += 2;
			continue;
		}
		if (strcmp(argv[0], "-f") == 0) {
			ecfile = fopen(argv[1], "r");
			if (ecfile == NULL) {
				fprintf(stderr, "Fail to open EC file %s!\n", argv[1]);
				return 1;
			}
			argc -= 2;
			argv += 2;
			continue;
		}
		usage();
	}

	if (argc != 2)
		usage();

	infile = fopen(argv[0], "r+");
	header_loc = strtol(argv[1], NULL, 0);

	if (infile == NULL) {
		fprintf(stderr, "Fail to open firmware image %s!\n", argv[1]);
		return 1;
	}

	if (spisize == 0) {
		fseek(infile, -1, SEEK_END);
		spisize = ftell(infile) + 1;

		printf("Input file size is 0x%lx.\n", spisize);
	}

	if (ecfile != NULL) {
		struct mec1322_header hdr;
		fread(&hdr, 1, sizeof(hdr), ecfile);
		fwlen = parse_header(&hdr);

		if (fwlen == 0) {
			fprintf(stderr, "Error parsing MEC1322 firmware header!\n");
			return 1;
		}
		fseek(ecfile, 0, SEEK_SET);
		size_t totallen = fwlen + 256; // firmware length with signature
		image = malloc(totallen);
		fread(image, 1, totallen, ecfile);
		fclose(ecfile);

		if (!verify_payload_signature(&hdr, image + hdr.payload_offset, image + fwlen))
			return 1;

		fwlen = totallen;
	}

	if (header_loc & 0xff) {
		fprintf(stderr, "Header location should be aligned to 256 bytes!\n");
		return 1;
	}
	uint8_t fwTag[4] = {header_loc >> 8, header_loc >> 16, header_loc >> 24};
	fwTag[3] = crc8(fwTag, 3);

	// write tag
	fseek(infile, -256, SEEK_END);
	fwrite(fwTag, 1, 4, infile);

	// if EC image is given, write EC image
	if (image) {
		fseek(infile, header_loc - spisize, SEEK_END);
		fwrite(image, 1, fwlen, infile);
		free(image);
	}

	fclose(infile);

	return 0;
}

int main(int argc, char *argv[])
{
	if (argc < 3)
		usage();

	if (strcmp(argv[1], "dump") == 0)
		return mec1322_dump(argc - 2, argv + 2);

	if (strcmp(argv[1], "insert") == 0)
		return mec1322_insert(argc - 2, argv + 2);

	// no action provided
	usage();
}