// SPDX-License-Identifier: WTFPL
// cc -Wall -O2 -I/usr/local/include -L/usr/local/lib -lsecp256k1 nostr-key.c -o nostr-key

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

#define unpacked_size(x) (((x) * 8 + 4) / 5)	// round-up
#define packed_size(x) (((x) * 5) / 8)		// round-down
#define expanded_hrp_size(x) (((x) * 2) + 1)

#define BECH32_CHECKSUM_LEN 6
#define LOWER_CASE_USED (1 << 0)
#define UPPER_CASE_USED (1 << 1)

static void encode_hex(char *out, int outsize, unsigned char *in, int insize)
{
	static char c[] = "0123456789abcdef";
	int i, j;

	for (i = j = 0; i < outsize - 1 && j < insize; j++) {
		out[i++] = c[(in[j] >> 4) & 0x0f];
		if (i < outsize - 1)
			out[i++] = c[in[j] & 0x0f];
	}
	out[i] = '\0';
}

static int decode_hex(unsigned char *buf, int buflen, const char *str)
{
	int i, p, n;
	char *tmp;

	n = buflen * 2;
	tmp = alloca(n);
	memset(tmp, 0, n);
	for (i = 0, p = strlen(str) - 1; i < n; i++, p--) {
		if (str[p] >= '0' && str[p] <= '9')
			tmp[i] = str[p] - '0';
		else if (str[p] >= 'a' && str[p] <= 'f')
			tmp[i] = str[p] - 'a' + 10;
		else if (str[p] >= 'A' && str[p] <= 'F')
			tmp[i] = str[p] - 'A' + 10;
		else return -1;
	}
				
	for (i = 0, p = n - 1; i < buflen; i++, p -= 2)
		buf[i] = (tmp[p] << 4) | tmp[p - 1];

	return 0;
}

static int check_string_case(char *str, int len)
{
	int i, m;

	m = 0;
	for (i = 0; i < len; i++) {
		if (str[i] >= 'a' && str[i] <= 'z')
			m |= LOWER_CASE_USED;
		else if (str[i] >= 'A' && str[i] <= 'Z')
			m |= UPPER_CASE_USED;
	}

	if (m == 0)
		m = LOWER_CASE_USED; // as default

	return m;
}

static uint32_t bech32_polymod(uint8_t *values, int len)
{
	uint32_t b, chk = 1;
	int i;

	for (i = 0; i < len; i++) {
		b = (chk >> 25);
		chk = (chk & 0x01ffffff) << 5 ^ (values[i] & 0x01f);
		chk ^= (b & 0x01) ? 0x3b6a57b2 : 0;
		chk ^= (b & 0x02) ? 0x26508e6d : 0;
		chk ^= (b & 0x04) ? 0x1ea119fa : 0;
		chk ^= (b & 0x08) ? 0x3d4233dd : 0;
		chk ^= (b & 0x10) ? 0x2a1462b3 : 0;
	}

	return chk;
}

static uint8_t bech32_char_decode(int c)
{
	const static uint8_t decode_table[] = {
		0x0f, 0xff, 0x0a, 0x11, 0x15, 0x14, 0x1a, 0x1e, // 01234567
		0x07, 0x05, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, // 89:;<=>?
		0xff, 0x1d, 0xff, 0x18, 0x0d, 0x19, 0x09, 0x08, // @ABCDEFG
		0x17, 0xff, 0x12, 0x16, 0x1f, 0x1b, 0x13, 0xff, // HIJKLMNO
		0x01, 0x00, 0x03, 0x10, 0x0b, 0x1c, 0x0c, 0x0e, // PQRSTUVW
		0x06, 0x04, 0x02, 0xff, 0xff, 0xff, 0xff, 0xff, // XYZ[\]^_
		0xff, 0x1d, 0xff, 0x18, 0x0d, 0x19, 0x09, 0x08, // `abcdefg
		0x17, 0xff, 0x12, 0x16, 0x1f, 0x1b, 0x13, 0xff, // hijklmno
		0x01, 0x00, 0x03, 0x10, 0x0b, 0x1c, 0x0c, 0x0e, // pqrstuvw
		0x06, 0x04, 0x02, 0xff, 0xff, 0xff, 0xff, 0xff, // xyz{|}~
	};

	return (c < '0' || c > 'z') ? 0xff : decode_table[c - '0'];
}

static int pack5to8(uint8_t *out, int outlen, uint8_t *in, int inlen)
{
	int inpos, outpos, olen;
	uint64_t v;

	olen = packed_size(inlen);
	outlen = (outlen < olen) ? outlen : olen;

	for (inpos = outpos = 0; inpos < inlen && outpos < outlen; ) {
		v = (uint64_t)(in[inpos++] & 0x1f) << 35;
		if (inpos < inlen) v |= (uint64_t)(in[inpos++] & 0x1f) << 30;
		if (inpos < inlen) v |= (uint64_t)(in[inpos++] & 0x1f) << 25;
		if (inpos < inlen) v |= (uint64_t)(in[inpos++] & 0x1f) << 20;
		if (inpos < inlen) v |= (uint64_t)(in[inpos++] & 0x1f) << 15;
		if (inpos < inlen) v |= (uint64_t)(in[inpos++] & 0x1f) << 10;
		if (inpos < inlen) v |= (uint64_t)(in[inpos++] & 0x1f) << 5;
		if (inpos < inlen) v |= (uint64_t)(in[inpos++] & 0x1f);

		out[outpos++] = v >> 32;
		if (outpos < outlen) out[outpos++] = v >> 24;
		if (outpos < outlen) out[outpos++] = v >> 16;
		if (outpos < outlen) out[outpos++] = v >> 8;
		if (outpos < outlen) out[outpos++] = v;
	}

	return outpos;
}

static int unpack8to5(uint8_t *out, int outlen, uint8_t *in, int inlen)
{
	int inpos, outpos, olen;
	uint64_t v;

	olen = unpacked_size(inlen);
	outlen = (outlen < olen) ? outlen : olen;

	for (inpos = outpos = 0; inpos < inlen && outpos < outlen; ) {
		v = (uint64_t)in[inpos++] << 32;
		if (inpos < inlen) v |= (uint64_t)in[inpos++] << 24;
		if (inpos < inlen) v |= (uint64_t)in[inpos++] << 16;
		if (inpos < inlen) v |= (uint64_t)in[inpos++] << 8;
		if (inpos < inlen) v |= (uint64_t)in[inpos++];

		out[outpos++] = (v >> 35) & 0x1f;
		if (outpos < outlen) out[outpos++] = (v >> 30) & 0x1f;
		if (outpos < outlen) out[outpos++] = (v >> 25) & 0x1f;
		if (outpos < outlen) out[outpos++] = (v >> 20) & 0x1f;
		if (outpos < outlen) out[outpos++] = (v >> 15) & 0x1f;
		if (outpos < outlen) out[outpos++] = (v >> 10) & 0x1f;
		if (outpos < outlen) out[outpos++] = (v >> 5) & 0x1f;
		if (outpos < outlen) out[outpos++] = v & 0x1f;
	}

	return outpos;
}

static int expand_hrp(uint8_t *out, int outlen, char *in, int inlen)
{
// do not use pre/post incremented argument in macro
#define lower_case(c) (((c) >= 'A' && (c) <= 'Z') ? ((c) - 'A' + 'a') : (c))

	int inpos, outpos, olen;

	olen = expanded_hrp_size(inlen);
	outlen = (outlen < olen) ? outlen : olen;

	for (inpos = outpos = 0; inpos < inlen && outpos < outlen;
	     inpos++, outpos++)
		out[outpos] = (lower_case(in[inpos]) >> 5) & 0x07;

	if (outpos < outlen)
		out[outpos++] = 0;

	for (inpos = 0; inpos < inlen && outpos < outlen; inpos++, outpos++)
		out[outpos] = lower_case(in[inpos]) & 0x1f;

	return outpos;
}

int nostr_key_decode(uint8_t *out, int outlen, char *in)
{
	int i, j, hrplen, tmpsize, v = -1;
	char *sep;
	uint8_t *tmp;

	/* get HRP length, check upper/lower case */
	if ((sep = strrchr(in, '1')) == NULL)
		goto fin0;
	hrplen = sep - in;
	if (check_string_case(in, strlen(in)) == (LOWER_CASE_USED |
						  UPPER_CASE_USED))
		goto fin0;

	/* allocate working buffer */
	tmpsize = expanded_hrp_size(hrplen) +
		unpacked_size(outlen) + BECH32_CHECKSUM_LEN;
	if ((tmp = calloc(tmpsize, sizeof(*tmp))) == NULL)
		goto fin0;

	/* expand HRP string to calculate checksum */
	expand_hrp(tmp, tmpsize, in, hrplen);

	/* decode bech32 string with checksum */
	for (i = hrplen + 1, j = expanded_hrp_size(hrplen); in[i]; i++, j++)
		if ((tmp[j] = bech32_char_decode(in[i])) == 0xff)
			goto fin1;

	/* validate checksum */
	if (bech32_polymod(tmp, tmpsize) != 1)
		goto fin1;

	/* output result */
	pack5to8(out, outlen,
		 tmp + expanded_hrp_size(hrplen), unpacked_size(outlen));
	v = 0;
fin1:
	memset(tmp, 0, tmpsize);
	free(tmp);
fin0:
	return v;
}

static void bech32_set_checksum(uint8_t *p, uint32_t cs)
{
	int i;

	for (i = 0; i < BECH32_CHECKSUM_LEN; i++)
		p[i] = (cs >> (5 * (5 - i))) & 0x1f;
}

int nostr_key_encode(char *out, int outlen, uint8_t *in, int inlen, char *hrp)
{
	int i, hrplen, tmpsize, v = -1;
	uint8_t *tmp, *p;
	const char *t;
	const static char *encode_table[] = {
		[0] = NULL,
		[LOWER_CASE_USED] = "qpzry9x8gf2tvdw0s3jn54khce6mua7l",
		[UPPER_CASE_USED] = "QPZRY9X8GF2TVDW0S3JN54KHCE6MUA7L",
		[LOWER_CASE_USED | UPPER_CASE_USED] = NULL,
	};

	/* get HRP length, check upper/lower case */
	hrplen = strlen(hrp);
	if (hrplen < 1 || hrplen > 83)
		goto fin0;
	if ((t = encode_table[check_string_case(hrp, hrplen)]) == NULL)
		goto fin0;

	/* allocate working buffer (with terminate character for string) */
	tmpsize = expanded_hrp_size(hrplen) +
		unpacked_size(inlen) + BECH32_CHECKSUM_LEN;
	if ((tmp = calloc(tmpsize + 1, sizeof(*tmp))) == NULL)
		goto fin0;

	/* expand HRP string to calculate checksum */
	expand_hrp(tmp, tmpsize, hrp, hrplen);

	/* convert hex payload */
	unpack8to5(tmp + expanded_hrp_size(hrplen), unpacked_size(inlen),
		   in, inlen);

	/* calculate checksum (checksum region is 0 filled) */
	bech32_set_checksum(tmp +
			    expanded_hrp_size(hrplen) + unpacked_size(inlen),
			    bech32_polymod(tmp, tmpsize) ^ 1);

	/* output result */
	p = tmp + expanded_hrp_size(hrplen) - hrplen - 1;
	memcpy(p, hrp, hrplen);
	p[hrplen] = '1';
	for (i = expanded_hrp_size(hrplen); i < tmpsize; i++)
		tmp[i] = t[tmp[i]];
	snprintf(out, outlen, "%s", p); // strlcpy() should be better
	v = 0;
//fin1:
	memset(tmp, 0, tmpsize);
	free(tmp);
fin0:
	return v;
}

int main(int argc, char *argv[])
{
#define NOSTR_KEY_LEN 32

	uint8_t seckey[NOSTR_KEY_LEN], pubkey[NOSTR_KEY_LEN];
	char out[256];
	secp256k1_context *ctx;
	secp256k1_keypair keypair;
	secp256k1_xonly_pubkey xpubkey;

	if (argc < 2) {
		printf("usage: %s [private key]\n", argv[0]);
		goto fin0;
	}

	if (nostr_key_decode(seckey, sizeof(seckey), argv[1]) &&
	    decode_hex(seckey, sizeof(seckey), argv[1])) {
		printf("invalid key\n");
		goto fin0;
	}

	encode_hex(out, sizeof(out), seckey, sizeof(seckey));
	printf("private key (hex): %s\n", out);

	nostr_key_encode(out, sizeof(out), seckey, sizeof(seckey), "nsec");
	printf("private key (nsec): %s\n", out);

	if ((ctx = secp256k1_context_create(SECP256K1_CONTEXT_SIGN)) == NULL) {
		printf("secp256k1_context_create\n");
		goto fin0;
	}
	if (!secp256k1_keypair_create(ctx, &keypair, seckey) ||
	    !secp256k1_keypair_xonly_pub(ctx, &xpubkey, NULL, &keypair) ||
	    !secp256k1_xonly_pubkey_serialize(ctx, pubkey, &xpubkey)) {
		printf("pubkey create error\n");
		goto fin1;
	}

	encode_hex(out, sizeof(out), pubkey, sizeof(pubkey));
	printf("public key (hex): %s\n", out);

	nostr_key_encode(out, sizeof(out), pubkey, sizeof(pubkey), "npub");
	printf("public key (npub) %s\n", out);

fin1:
	secp256k1_context_destroy(ctx);
fin0:
	return 0;
}