#define _POSIX_C_SOURCE 200809L

/* general zpm sign function */

/* read key, export key, import key, sign, export signature, import
 * signature, verify signature
 */

/* import and export is base64 encoded */

/* -v verify
 * -s sign
 * -g generate key
 * -e extract public key
 *
 * -o output file
 *
 * -p read key from file
 * -P read public key from argument to -P
 * -S read secret key from argument
 * -m message from argument
 * -h message is hex encoded
 */

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <sys/mman.h>
#include <time.h>

/* linux specific for getrandom */
#include <sys/random.h>

#include "eddsa.h"
#include "sha512.h"

#define VERIFY 1
#define SIGN 2
#define GENKEY 3
#define EXTRACT 4

static void hexdump(void *src, size_t len) {
	unsigned char *b = src;
	while (len--) {
		printf("%02x", *b++);
	}
	printf("\n");
}

static char hexchars[] = "0123456789abcdefABCDEF";

static void hex(char *dst, uint8_t *src, size_t len) {
	while (len--) {
		dst[0] = hexchars[(src[0]>>4)&0xf];
		dst[1] = hexchars[src[0]&0xf];
		dst+=2;
		src++;
	}
}

static void hexbin(uint8_t *dst, unsigned char *src, size_t len) {
	size_t i;
	int x;

	for (i=0; i<len; i+=2) {
		sscanf((const char *)src+i, "%02x", &x);
		dst[i/2] = x;
	}
}

void *map(char *file, size_t *size) {
	int fd, rv;
	void *m;
	struct stat st;

	fd = open(file, O_RDONLY);
	if (fd == -1) {
		return NULL;
	}

	rv = fstat(fd, &st);
	if (rv == -1) {
		close(fd);
		return NULL;
	}

	m = mmap(0, st.st_size, PROT_READ, MAP_PRIVATE, fd, 0);
	if (m == MAP_FAILED) {
		return 0;
	}

	close(fd);
	*size = st.st_size;
	return m;
}

ssize_t readbytes(char *dest, char *file, size_t n) {
	char *m;
	size_t fs;

	m = map(file, &fs);
	if (!m) {
		return -1;
	}
	if (n > fs) {
		n = fs;
	}

	size_t pk = strspn(m, hexchars);
	pk /= 2;
	if (n > pk) {
		n = pk;
	}

	hexbin(dest, m, n*2);
	munmap(m, fs);
	return n;
}


/* private key format is:
 * "ZPMS" 4 byte magic "ZPMS"
 * byte 0x01 = version 1
 * byte 0x01 = private key
 * byte 0x01 = chacha encrypted, 0x00 = plain/raw key bytes
 * byte 0x00 reserved, probably "key usage"
 * 32 bytes private key
 * 24 bytes reserved
 * sub keys, public keys?  No point to the public key, since it's
 * always derivable.  could have 16 bytes of IV for chacha to encrypt.
 *
 * entire thing is then hexencoded, but 
 */

/* public key format is:
 * "ZPMS" 4 byte magic "ZPMS"
 * 0x01   1 byte version
 * 0x02   1 byte "public key packet"
 * 0x00   1 byte reserved
 * 0x00   1 byte reserved
 * 0x..  32 bytes public key
 * 0x..  24 bytes reserved
 * 0x..     revocation signatures
 * 0x..     user id packets
 * 0x..     certification signatures
 * 0x..     sub key packets
 */

/* timestamps are 8 byte little endian integers of seconds since
 * the epoch.  or unix time, or some such.  Times are a mess.
 * Unix time is leap second unaware, and ambiguous during leap seconds.
 * UTC requires a database of leap seconds for past time, and the
 * specified second in the future possibly changes.
 * TAI is good, but currently 37 seconds behind UTC, which is just weird
 */

/* signature items from openpgp */
/* revokable flag
 * 0-255 trust level
 *
 */

/* signature format:
 * "ZPMS" 4 byte magic "ZPMS"
 * 0x01   1 byte version
 * 0x03   1 byte "signature packet"
 * 0x00   1 byte reserved
 * 0x00   1 byte reserved
 * 0x..   8 bytes timestamp of signature time
 * 0x..   8 bytes timestamp of signature expiration, 0 if none
 * 0x..  64 bytes of actual signature
 * 0x..  32 bytes of public key making the signature
 * 0x..   8 bytes reserved
 * Total is 128 bytes.
 */

/* expires time could probably be 4 bytes as an offset from created.
 * That would free up 4 bytes.  Sign time could probably be four
 * bytes with an epoch of say 2019-01-01, since negative times don't
 * really make any sense, that give more than 100 years, and frees
 * up another 4 bytes
 *
 * version is zero until this proof of concept is better validated
 */
struct signature {
	char magic[4]; /* magic ZPMS */
	char info[4]; /* version = 0, type=3, 1 byte trust, 1 byte reserved */
	uint64_t signtime; /* unix time */
	uint64_t expires; /* unix time, 0 = none */
	char sig[64]; 
	char pub[32];
	char reserved[8];
};

struct key {
	char magic[4]; /* "ZPMS" */
	char info[4]; /* version, type = 0x02, reserved, reserved */
	char key[32];
	uint64_t created;
	uint64_t expires;
	char reserved[8];
	char sig[64]; /* zeroes for a private key, or fill in, doesn't matter*/
	/* could fill in the first 32 bytes of sig with the public key
	 * if this is a private key, then you don't need another structure */
};

int create_key(struct key *sk) {
	char info[] = { 0, 1, 0, 0 };
	memcpy(sk->magic, "ZPMS", 4);
	memcpy(sk->info, info, 4);
	getrandom(sk->key, sizeof sk->key, 0);
	ed25519_genpub(sk->sig, sk->key);
	sk->created = time(NULL);
	memset(sk->reserved, 0, 8);
	memset(sk->sig+32, 0, 32);
	return 1;
}

int derive_pubkey(struct key *pk, struct key *sk) {
	char info[] = { 0, 2, 0, 0 };

	memcpy(sk->magic, "ZPMS", 4);
	memcpy(sk->info, info, 4);
	ed25519_genpub(pk->key, sk->key);
	pk->created = sk->created;
	pk->expires = sk->expires;
	memset(sk->reserved, 0, 8);

	ed25519_sign(pk->sig, sk->key, pk->sig, sk->info, 52);
	return 1;

}

/* if reserved used 16 bytes for signtime and expires, and we added
 * 64 bytes as a signature, a public key could include it's own self
 * signature, and a key would be 128 bytes, as would a signature.  This
 * would make all structures 128 bytes.
 */

int read_secret_key(uint8_t *sec, uint8_t *pub, char *file) {
	uint8_t sk[64];
	
	readbytes(sk, file, sizeof sk);
	if (sk[4] != 1 || sk[5] != 1) {
		printf("magic: ");
		hexdump(sk, 16);
		return 0;
	}
	memcpy(sec, sk+8, 32);
	//printf("sec: ");
	//hexdump(sec, 32);

	if (pub) {
		ed25519_genpub(pub, sec);
	//	printf("pub: ");
	//	hexdump(pub, 32);
	}

	return 1;
}

/* private key format is:
 * "ZPMS" 4 byte magic "ZPMS"
 * byte 0x01 = version 1
 * byte 0x01 = private key
 * byte 0x01 = chacha encrypted, 0x00 = plain/raw key bytes
 * byte 0x00 reserved, probably "key usage"
 * 32 bytes private key
 * 24 bytes reserved
 * sub keys, public keys?  No point to the public key, since it's
 * always derivable.  could have 16 bytes of IV for chacha to encrypt.
 *
 * entire thing is then hexencoded, but 
 */

int write_secret_key(uint8_t *sec, int flags, int fd) {
	uint8_t skh[128] = "5A504D5301010000"; /* ZPMS 1 1 0 0 */

	if (flags & 0x01) {
		skh[13] = '1';
	}
	//hexdump(sec, 32);
	hex(skh + 16, sec, 32);
	memset(skh + 80, '0', 48);
	write(fd, skh, sizeof skh);
	memset(skh, 0, sizeof skh);

	return 1;
}

int main(int ac, char *av[]) {
	int option;
	int mode = 0;
	char *outfile = 0;
	uint8_t pub[ED25519_KEY_LEN];
	/* TODO should use mlock()ed memory for the secret key */
	uint8_t sec[ED25519_KEY_LEN];
	uint8_t sig[ED25519_SIG_LEN];
	void *message = 0;
	char *messagefile = 0;
	char *messagestring = 0;
	char messagehash[SHA512_HASH_LENGTH];
	char *keystring = 0, *keyfile = 0;
	unsigned char *sigstring = 0, *sigfile = 0;
	char *passphrase = 0;
	size_t mlen;
	int rawsig = 0, hexencoded = 0;
	int debug = 0;

	while ((option = getopt(ac, av, "o:vS:f:sgek:K:rm:hd")) != -1) {
		switch (option) {
			case 'o': outfile = optarg; break;
			case 'v': mode = VERIFY; break;
			case 'S': sigstring = optarg; break;
			case 'f': sigfile = optarg; break;
			case 's': mode = SIGN; break;
			case 'g': mode = GENKEY; break;
			case 'e': mode = EXTRACT; break;
			case 'k': keystring = optarg; break;
			case 'K': keyfile = optarg; break;
			case 'p': passphrase = optarg; break;
			case 'm': messagestring = optarg; break;
			case 'h': hexencoded = 1; break;
			case 'r': rawsig = 1; break;
			case 'd': debug++; break;
			default:
				  exit(EXIT_FAILURE);
				  break;
		}
	}

	if (sigfile) {
		ssize_t r;
		memset(sig, 0, sizeof sig);
		r = readbytes(sig, sigfile, sizeof sig);
		if (r == -1) {
			return 8;
		}
		if ((size_t)r < sizeof sig) {
			return 7;
		}
	} else if (sigstring) {
		memset(sig, 0, sizeof sig);
		hexbin(sig, sigstring, sizeof sig);
	}

	if (keystring) {
		memset(sec, 0, sizeof sec);
		hexbin(sec, keystring, strlen(keystring));
		ed25519_genpub(pub, sec);
		if (debug) {
			printf("sec:\n"); hexdump(sec, sizeof sec);
		}
	} else if (keyfile) {
		read_secret_key(sec, pub, keyfile);
		//hexdump(sec, sizeof sec);
		//hexdump(pub, sizeof pub);
	}

	if (mode == EXTRACT) {
		/* want to print the pubkey only */
		hexdump(pub, sizeof pub);
		exit(0);
	}

	/* need to be able to encrypt the private keys */
	/* use chacha */
	if (mode == GENKEY) {
		/* a private key is just 32 random bytes */
		getrandom(sec, sizeof sec, 0);
		ed25519_genpub(pub, sec);
		ed25519_sign(sig, sec, pub, sec, sizeof sec);
		//int rv = ed25519_verify(sig, pub, sec, sizeof sec);
		if (outfile) {
			int fd;
			fd = open(outfile, O_WRONLY|O_CREAT, 0600);
			if (fd == -1) {
				perror("can't open outfile");
				exit(EXIT_FAILURE);
			}
			write_secret_key(sec, 0, fd);
			close(fd);
		} else {
			write_secret_key(sec, 0, 1);
		}
	}

	/* set up the message data */
	if (mode == SIGN || mode == VERIFY) {
		messagefile = av[optind];
		if (messagefile) {
			message = map(messagefile, &mlen);
			if (!message) {
				return 9;
			}
		} else if (messagestring) {
			message = messagestring;
			if (hexencoded) {
				mlen = strlen(messagestring)/2;
				message = malloc(mlen);
				if (!message) {
					return 10;
				}
				hexbin(message, messagestring, mlen*2);
			} else {
				mlen = strlen(messagestring);
			}
		} else {
			return 2;
		}
		if (!message) {
			return 3;
		}
	}

	if (mode == SIGN) {
		struct signature s;
		s.magic[0] = 'Z';
		s.magic[1] = 'P';
		s.magic[2] = 'M';
		s.magic[3] = 'S';
		s.info[0] = 1;
		s.info[1] = 1;
		s.info[2] = 0;
		s.info[3] = 0;
		s.signtime = time(NULL);
		s.expires = 0;
		memset(s.reserved, 0, 8);
		if (!rawsig) {
			struct sha512 hash;
			sha512_init(&hash);
			sha512_add(&hash, message, mlen);
			sha512_add(&hash, (uint8_t *)&s.signtime, sizeof s.signtime);
			sha512_add(&hash, (uint8_t *)&s.expires, sizeof s.expires);
			sha512_add(&hash, s.info, sizeof s.info);
			sha512_final(&hash, messagehash);
			message = messagehash;
		}

		if (message == 0) {
			return 5;
		}

		ed25519_sign(sig, sec, pub, message, mlen);
		memcpy(s.sig, sig, sizeof sig);
		memcpy(s.pub, pub, sizeof pub);
		if (rawsig) {
			hexdump(sig, sizeof sig);
		} else {
			hexdump(&s, sizeof s);
		}
	}

	if (mode == VERIFY) {
		fprintf(stderr, "sig: ");
		hexdump(sig, sizeof sig);
		int rv = ed25519_verify(sig, pub, message, mlen);
		if (rv) {
			fprintf(stderr, "verified\n");
			return 0;
		} else {
			fprintf(stderr, "not verified\n");
			return 6;
		}
	}

	if (message && messagefile) {
		munmap(message, mlen);
	}

	return 0;
}