/*
 * wgtool.c - minimal SIOCSWG/SIOCGWG driver for DragonFly if_wg
 * (the guest has no wg(8)). Self-contained RFC 7748 X25519.
 *
 * Subcommands:
 *   selftest                     - verify X25519 against RFC 7748 6.1
 *   pubkey <privhex>             - print derived public key hex
 *   set    <if> <privhex> <port> - set private key + listen port
 *   setport <if> <port>          - change listen port only
 *   peer <if> <pubhex> <ip> <port> <cidr> [pka] - add peer
 *   show <if>                    - dump peer count via SIOCGWG
 *
 * Build: cc -O2 -o wgtool wgtool.c
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/time.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <net/if.h>

/* ---- verbatim struct layout from sys/net/wg/if_wg.h (keep in sync) ---- */
#define WG_KEY_SIZE 32
#define WG_PEER_DESCR_SIZE 64
#define SIOCSWG _IOWR('i', 210, struct wg_data_io)
#define SIOCGWG _IOWR('i', 211, struct wg_data_io)

struct wg_aip_io {
	sa_family_t	 a_af;
	int		 a_cidr;
	union {
		struct in_addr		addr_ipv4;
		struct in6_addr		addr_ipv6;
	}		 a_addr;
};
struct wg_peer_io {
	int			p_flags;
	uint8_t			p_public[WG_KEY_SIZE];
	uint8_t			p_psk[WG_KEY_SIZE];
	uint16_t		p_pka;
	union {
		struct sockaddr		sa_sa;
		struct sockaddr_in	sa_sin;
		struct sockaddr_in6	sa_sin6;
	}			p_endpoint;
	uint64_t		p_txbytes;
	uint64_t		p_rxbytes;
	struct timespec		p_last_handshake;
	uint64_t		p_id;
	char			p_description[WG_PEER_DESCR_SIZE];
	size_t			p_aips_count;
	struct wg_aip_io	p_aips[];
};
#define WG_PEER_HAS_PUBLIC	(1 << 0)
#define WG_PEER_HAS_PSK		(1 << 1)
#define WG_PEER_HAS_PKA		(1 << 2)
#define WG_PEER_HAS_ENDPOINT	(1 << 3)
#define WG_PEER_REPLACE_AIPS	(1 << 4)
struct wg_interface_io {
	int			i_flags;
	in_port_t		i_port;
	uint32_t		i_cookie;
	uint8_t			i_public[WG_KEY_SIZE];
	uint8_t			i_private[WG_KEY_SIZE];
	size_t			i_peers_count;
	struct wg_peer_io	i_peers[];
};
#define WG_INTERFACE_HAS_PUBLIC		(1 << 0)
#define WG_INTERFACE_HAS_PRIVATE	(1 << 1)
#define WG_INTERFACE_HAS_PORT		(1 << 2)
#define WG_INTERFACE_HAS_COOKIE		(1 << 3)
#define WG_INTERFACE_REPLACE_PEERS	(1 << 4)
struct wg_data_io {
	char			 wgd_name[IFNAMSIZ];
	size_t			 wgd_size;
	struct wg_interface_io	*wgd_interface;
};

/* ---------------------- X25519 (RFC 7748, fe = 5x51) ------------------- */
typedef uint64_t fe[5];
typedef unsigned __int128 i128;
#define MASK51 ((1ULL << 51) - 1)

static void fe_mul(fe r, const fe a, const fe b)
{
	i128 t0,t1,t2,t3,t4, c;
	uint64_t o0,o1,o2,o3,o4;
	int i;

	t0 = (i128)a[0]*b[0] + 19*((i128)a[1]*b[4] + (i128)a[2]*b[3] +
	     (i128)a[3]*b[2] + (i128)a[4]*b[1]);
	t1 = (i128)a[0]*b[1] + (i128)a[1]*b[0] + 19*((i128)a[2]*b[4] +
	     (i128)a[3]*b[3] + (i128)a[4]*b[2]);
	t2 = (i128)a[0]*b[2] + (i128)a[1]*b[1] + (i128)a[2]*b[0] +
	     19*((i128)a[3]*b[4] + (i128)a[4]*b[3]);
	t3 = (i128)a[0]*b[3] + (i128)a[1]*b[2] + (i128)a[2]*b[1] +
	     (i128)a[3]*b[0] + 19*(i128)a[4]*b[4];
	t4 = (i128)a[0]*b[4] + (i128)a[1]*b[3] + (i128)a[2]*b[2] +
	     (i128)a[3]*b[1] + (i128)a[4]*b[0];

	for (i = 0; i < 2; i++) {
		c = t0 >> 51; o0 = (uint64_t)t0 & MASK51;
		t1 += c; c = t1 >> 51; o1 = (uint64_t)t1 & MASK51;
		t2 += c; c = t2 >> 51; o2 = (uint64_t)t2 & MASK51;
		t3 += c; c = t3 >> 51; o3 = (uint64_t)t3 & MASK51;
		t4 += c; c = t4 >> 51; o4 = (uint64_t)t4 & MASK51;
		t0 = (i128)(o0 + (uint64_t)c * 19);
	}
	r[0] = (uint64_t)t0; r[1] = o1; r[2] = o2; r[3] = o3; r[4] = o4;
}
static void fe_sq(fe r, const fe a) { fe_mul(r, a, a); }
static void fe_add(fe r, const fe a, const fe b)
{
	int i;
	for (i = 0; i < 5; i++) r[i] = a[i] + b[i];
}
static void fe_sub(fe r, const fe a, const fe b)
{
	/* exact a - b (mod p) via borrow chain; inputs are non-negative
	   limbs < 2^52.  If the overall value went negative, add p =
	   2^255 - 19 back in. */
	uint64_t t[5], borrow, nb;
	i128 d;
	int i;

	borrow = 0;
	for (i = 0; i < 5; i++) {
		d = (i128)a[i] - b[i] - borrow;
		if (d < 0) { d += ((i128)1 << 51); nb = 1; } else nb = 0;
		t[i] = (uint64_t)d;
		borrow = nb;
	}
	if (borrow) {
		/* value = sum(t) - 2^255; add p = +2^255 - 19  =>  sum(t) - 19 */
		t[0] -= 19;
		for (i = 0; i < 4; i++) {
			if (t[i] >> 63) {
				t[i] += ((uint64_t)1 << 51);
				t[i + 1] -= 1;
			}
		}
		/* t[4] cannot go negative (sum(t) >= 19 holds here) */
	}
	for (i = 0; i < 5; i++) r[i] = t[i];
}
static void fe_cswap(uint64_t swap, fe a, fe b)
{
	uint64_t m = (uint64_t)-(int64_t)swap, t;
	int i;
	for (i = 0; i < 5; i++) { t = m & (a[i] ^ b[i]); a[i] ^= t; b[i] ^= t; }
}
static void fe_mulsmall(fe r, const fe a, uint64_t s)
{
	fe b = { s, 0, 0, 0, 0 };
	fe_mul(r, a, b);
}
/* exponent p-2 = 2^255 - 21, little-endian bytes */
static const uint8_t fe_pmin2[32] = {
	0xeb,0xff,0xff,0xff,0xff,0xff,0xff,0xff,0xff,0xff,0xff,0xff,0xff,
	0xff,0xff,0xff,0xff,0xff,0xff,0xff,0xff,0xff,0xff,0xff,0xff,0xff,
	0xff,0xff,0xff,0xff,0xff,0x7f
};
static void fe_inv(fe r, const fe a)
{
	fe result;
	int i, first = 1;

	/* MSB-first square-and-multiply; base stays = a (bit 255 is 0) */
	for (i = 254; i >= 0; i--) {
		int byte = i / 8, b = i % 8;
		if (!first)
			fe_sq(result, result);
		if ((fe_pmin2[byte] >> b) & 1) {
			if (first) {
				memcpy(result, a, sizeof(fe));
				first = 0;
			} else {
				fe_mul(result, result, a);
			}
		}
	}
	memcpy(r, result, sizeof(fe));
}
static void fe_frombytes(fe r, const uint8_t *in)
{
	int i;

	for (i = 0; i < 5; i++) r[i] = 0;
	for (i = 0; i < 32; i++) {
		int limb = (i * 8) / 51, off = (i * 8) % 51;
		uint64_t v = (uint64_t)in[i];
		if (off + 8 <= 51) {
			r[limb] |= v << off;	/* byte fits inside the limb */
		} else {
			/* byte straddles limb boundary */
			r[limb] |= (v & ((1ULL << (51 - off)) - 1)) << off;
			r[limb + 1] |= v >> (51 - off);
		}
	}
	r[4] &= MASK51;		/* drop the 2^255 bit */
}
static void fe_tobytes(uint8_t *out, const fe a)
{
	fe x;
	uint64_t c0, c, carry;
	int i, k;

	memcpy(x, a, sizeof(fe));
	/* full carry propagation (3 rounds is provably enough) */
	for (k = 0; k < 3; k++) {
		uint64_t carry = 0;
		for (i = 0; i < 5; i++) {
			uint64_t v = x[i] + carry;
			x[i] = v & MASK51;
			carry = v >> 51;
		}
		x[0] += 19 * carry;	/* < 2^51 + 19 */
	}

	/* conditional subtraction of p: d = (x + 19) mod 2^255 == x - p
	   iff x + 19 wraps past 2^255 (i.e. x >= p).  Two rounds are more
	   than sufficient for x < 2^255 + 2^53. */
	for (k = 0; k < 2; k++) {
		uint64_t d0, d1, d2, d3, d4, c, wrap;

		c0 = x[0] + 19;  d0 = c0 & MASK51; c = c0 >> 51;
		c0 = x[1] + c;   d1 = c0 & MASK51; c = c0 >> 51;
		c0 = x[2] + c;   d2 = c0 & MASK51; c = c0 >> 51;
		c0 = x[3] + c;   d3 = c0 & MASK51; c = c0 >> 51;
		c0 = x[4] + c;   d4 = c0 & MASK51; wrap = c0 >> 51;
		if (wrap) {
			x[0] = d0; x[1] = d1; x[2] = d2; x[3] = d3; x[4] = d4;
		}
	}
	/* pack (limbs are < 2^51) */
	for (i = 0; i < 32; i++) {
		int limb = (i * 8) / 51, off = (i * 8) % 51;
		uint64_t v = x[limb] >> off;
		if (off > 43 && limb < 4)
			v |= x[limb + 1] << (51 - off);
		out[i] = (uint8_t)v;
	}
}
static const uint8_t basepoint[32] = { 9 };

static void x25519(uint8_t out[32], const uint8_t scalar[32],
		   const uint8_t point[32])
{
	uint8_t k[32];
	fe x1, x2, z2, x3, z3, A, AA, B, BB, E, C, D, DA, CB, tmp, t0;
	uint64_t swap = 0;
	int pos;

	memcpy(k, scalar, 32);
	k[0] &= 248; k[31] &= 127; k[31] |= 64;		/* clamp */

	fe_frombytes(x1, point);
	x2[0] = 1; x2[1] = x2[2] = x2[3] = x2[4] = 0;
	z2[0] = z2[1] = z2[2] = z2[3] = z2[4] = 0;
	memcpy(x3, x1, sizeof(fe));
	z3[0] = 1; z3[1] = z3[2] = z3[3] = z3[4] = 0;

	for (pos = 254; pos >= 0; pos--) {
		uint64_t kt = (k[pos / 8] >> (pos % 8)) & 1;
		swap ^= kt;
		fe_cswap(swap, x2, x3);
		fe_cswap(swap, z2, z3);
		swap = kt;

		fe_add(A, x2, z2);
		fe_sq(AA, A);
		fe_sub(B, x2, z2);
		fe_sq(BB, B);
		fe_sub(E, AA, BB);
		fe_add(C, x3, z3);
		fe_sub(D, x3, z3);
		fe_mul(DA, D, A);
		fe_mul(CB, C, B);
		fe_add(tmp, DA, CB);
		fe_sq(x3, tmp);
		fe_sub(tmp, DA, CB);
		fe_sq(t0, tmp);
		fe_mul(z3, x1, t0);
		fe_mul(x2, AA, BB);
		fe_mulsmall(tmp, E, 121665);
		fe_add(tmp, AA, tmp);
		fe_mul(z2, E, tmp);
	}
	fe_cswap(swap, x2, x3);
	fe_cswap(swap, z2, z3);
	fe_inv(t0, z2);
	fe_mul(tmp, x2, t0);
	fe_tobytes(out, tmp);
}

static int hex2bin_(const char *h, uint8_t *out, size_t n)
{
	size_t i;
	if (strlen(h) != n * 2) return -1;
	for (i = 0; i < n; i++) {
		unsigned v;
		if (sscanf(h + 2 * i, "%2x", &v) != 1) return -1;
		out[i] = (uint8_t)v;
	}
	return 0;
}
static void bin2hex_(const uint8_t *b, size_t n, char *out)
{
	size_t i;
	for (i = 0; i < n; i++) sprintf(out + 2 * i, "%02x", b[i]);
	out[2 * n] = 0;
}

static int x25519_selftest(void)
{
	static const char *a_priv =
		"77076d0a7318a57d3c16c17251b26645df4c2f87ebc0992ab177fba51db92c2a";
	static const char *a_pub =
		"8520f0098930a754748b7ddcb43ef75a0dbf3a0d26381af4eba4a98eaa9b4e6a";
	static const char *b_priv =
		"5dab087e624a8a4b79e17f8b83800ee66f3bb1292618b6fd1c2f8b27ff88e0eb";
	static const char *b_pub =
		"de9edb7d7b7dc1b4d35b61c2ece435373f8343c85b78674dadfc7e146f882b4f";
	uint8_t priv[32], pub[32], got[32], sh1[32], sh2[32];
	char hex[65];

	if (hex2bin_(a_priv, priv, 32) < 0) return -1;
	if (hex2bin_(a_pub, pub, 32) < 0) return -1;
	x25519(got, priv, basepoint);
	bin2hex_(got, 32, hex);
	if (strcmp(hex, a_pub) != 0) {
		fprintf(stderr, "X25519 selftest FAIL A: got %s want %s\n",
		    hex, a_pub);
		return -1;
	}
	if (hex2bin_(b_priv, priv, 32) < 0) return -1;
	if (hex2bin_(b_pub, pub, 32) < 0) return -1;
	x25519(got, priv, basepoint);
	bin2hex_(got, 32, hex);
	if (strcmp(hex, b_pub) != 0) {
		fprintf(stderr, "X25519 selftest FAIL B: got %s want %s\n",
		    hex, b_pub);
		return -1;
	}
	/* shared secret both directions */
	if (hex2bin_(a_priv, priv, 32) < 0) return -1;
	if (hex2bin_(b_pub, pub, 32) < 0) return -1;
	x25519(sh1, priv, pub);
	if (hex2bin_(b_priv, priv, 32) < 0) return -1;
	if (hex2bin_(a_pub, pub, 32) < 0) return -1;
	x25519(sh2, priv, pub);
	if (memcmp(sh1, sh2, 32) != 0) {
		fprintf(stderr, "X25519 selftest FAIL: DH asymmetry\n");
		return -1;
	}
	printf("X25519 selftest OK\n");
	return 0;
}

/* ---------------------------- wg ioctl glue ---------------------------- */
static int wg_fd(void)
{
	int s = socket(AF_INET, SOCK_DGRAM, 0);
	if (s < 0) { perror("socket"); exit(1); }
	return s;
}

static int wg_set(const char *ifname, const uint8_t *priv /*32 or NULL*/,
		  int port /* -1 = unchanged */, const uint8_t *peer_pub,
		  const struct sockaddr_in *peer_sin, int aip_cidr, int pka)
{
	struct {
		struct wg_interface_io i;
		struct wg_peer_io p;
		struct wg_aip_io a;
	} buf;
	struct wg_data_io wgd;
	int s;

	memset(&buf, 0, sizeof(buf));
	if (priv) {
		memcpy(buf.i.i_private, priv, 32);
		buf.i.i_flags |= WG_INTERFACE_HAS_PRIVATE;
	}
	if (port >= 0) {
		buf.i.i_port = (in_port_t)port;
		buf.i.i_flags |= WG_INTERFACE_HAS_PORT;
	}
	if (peer_pub) {
		memcpy(buf.i.i_peers[0].p_public, peer_pub, 32);
		buf.i.i_peers[0].p_flags |= WG_PEER_HAS_PUBLIC;
		if (peer_sin) {
			memcpy(&buf.i.i_peers[0].p_endpoint.sa_sin, peer_sin,
			       sizeof(struct sockaddr_in));
			buf.i.i_peers[0].p_flags |= WG_PEER_HAS_ENDPOINT;
		}
		if (aip_cidr > 0) {
			buf.i.i_peers[0].p_aips[0].a_af = AF_INET;
			buf.i.i_peers[0].p_aips[0].a_cidr = aip_cidr;
			buf.i.i_peers[0].p_aips[0].a_addr.addr_ipv4.s_addr =
			    htonl(0x0a070000);		/* 10.7.0.0 */
			buf.i.i_peers[0].p_aips_count = 1;
		}
		if (pka > 0) {
			buf.i.i_peers[0].p_pka = (uint16_t)pka;
			buf.i.i_peers[0].p_flags |= WG_PEER_HAS_PKA;
		}
		buf.i.i_peers_count = 1;
	}

	memset(&wgd, 0, sizeof(wgd));
	strlcpy(wgd.wgd_name, ifname, IFNAMSIZ);
	wgd.wgd_size = sizeof(buf);
	wgd.wgd_interface = (struct wg_interface_io *)&buf;

	s = wg_fd();
	if (ioctl(s, SIOCSWG, &wgd) < 0) {
		perror("SIOCSWG");
		close(s);
		return -1;
	}
	close(s);
	return 0;
}

static int wg_show(const char *ifname)
{
	struct wg_data_io wgd;
	size_t sz = 262144;
	void *buf = malloc(sz);
	int s, ret;

	memset(&wgd, 0, sizeof(wgd));
	strlcpy(wgd.wgd_name, ifname, IFNAMSIZ);
	wgd.wgd_size = sz;
	wgd.wgd_interface = buf;
	s = wg_fd();
	ret = ioctl(s, SIOCGWG, &wgd);
	close(s);
	if (ret < 0) { perror("SIOCGWG"); free(buf); return -1; }
	printf("%s: port=%u peers=%zu (required size %zu)\n", ifname,
	       ((struct wg_interface_io *)buf)->i_port,
	       ((struct wg_interface_io *)buf)->i_peers_count, wgd.wgd_size);
	free(buf);
	return 0;
}

static int wg_set_cookie_io(const char *ifname, uint32_t cookie)
{
	struct wg_interface_io ibuf;
	struct wg_data_io wgd;
	int s;

	memset(&ibuf, 0, sizeof(ibuf));
	ibuf.i_flags = WG_INTERFACE_HAS_COOKIE;
	ibuf.i_cookie = cookie;

	memset(&wgd, 0, sizeof(wgd));
	strlcpy(wgd.wgd_name, ifname, IFNAMSIZ);
	wgd.wgd_size = sizeof(ibuf);
	wgd.wgd_interface = &ibuf;
	s = wg_fd();
	if (ioctl(s, SIOCSWG, &wgd) < 0) {
		perror("SIOCSWG cookie");
		close(s);
		return -1;
	}
	close(s);
	return 0;
}

int main(int argc, char **argv)
{
	uint8_t priv[32], pub[32];
	char hex[65];
	struct sockaddr_in sin;

	if (argc < 2) goto usage;
	if (!strcmp(argv[1], "selftest")) return x25519_selftest() ? 1 : 0;
	if (!strcmp(argv[1], "pubkey")) {
		if (argc < 3) goto usage;
		if (hex2bin_(argv[2], priv, 32) < 0) { fprintf(stderr,"bad hex\n"); return 1; }
		x25519(pub, priv, basepoint);
		bin2hex_(pub, 32, hex);
		printf("%s\n", hex);
		return 0;
	}
	if (!strcmp(argv[1], "set")) {		/* set <if> <privhex> <port> */
		if (argc < 5) goto usage;
		if (hex2bin_(argv[3], priv, 32) < 0) { fprintf(stderr,"bad hex\n"); return 1; }
		return wg_set(argv[2], priv, atoi(argv[4]), NULL, NULL, 0, 0) < 0;
	}
	if (!strcmp(argv[1], "setport")) {	/* setport <if> <port> */
		if (argc < 4) goto usage;
		return wg_set(argv[2], NULL, atoi(argv[3]), NULL, NULL, 0, 0) < 0;
	}
	if (!strcmp(argv[1], "cookie")) {	/* cookie <if> <value> */
		if (argc < 4) goto usage;
		return wg_set_cookie_io(argv[2], (uint32_t)strtoul(argv[3],
		    NULL, 0)) < 0;
	}
	if (!strcmp(argv[1], "peer")) {
		/* peer <if> <pubhex> <ip> <port> <cidr> [pka] */
		if (argc < 7) goto usage;
		if (hex2bin_(argv[3], pub, 32) < 0) { fprintf(stderr,"bad hex\n"); return 1; }
		memset(&sin, 0, sizeof(sin));
		sin.sin_len = sizeof(sin);
		sin.sin_family = AF_INET;
		sin.sin_port = htons((uint16_t)atoi(argv[5]));
		if (inet_pton(AF_INET, argv[4], &sin.sin_addr) != 1) {
			fprintf(stderr, "bad ip\n"); return 1;
		}
		return wg_set(argv[2], NULL, -1, pub, &sin, atoi(argv[6]),
			      argc > 7 ? atoi(argv[7]) : 0) < 0;
	}
	if (!strcmp(argv[1], "show")) return wg_show(argv[2]) < 0;
usage:
	fprintf(stderr,
	    "usage: wgtool selftest | pubkey P | set IF PRIV PORT | setport IF PORT\n"
	    "       wgtool peer IF PUB IP PORT CIDR [PKA] | show IF\n");
	return 1;
}