#!/usr/bin/env ucode

'use strict';

import { access, basename, dirname, mkstemp, open, writefile } from 'fs';

function assert(cond, message) {
	if (!cond) {
		warn(message, "\n");
		exit(1);
	}

	return true;
}

let unet_tool = "unet-tool";
let script_dir = sourcepath(0, true);

if (basename(script_dir) == "scripts") {
	unet_tool = `${dirname(script_dir)}/unet-tool`;
	assert(access(unet_tool, "x"), "unet-tool missing");
}

let args = {};

const defaults = {
	port: 51830,
	pex_port: 51831,
	keepalive: 10,
};

const usage_message = `
Usage: ${basename(sourcepath())} [<flags>] <file> <command> [<args>] [<option>=<value> ...]

     Commands:
      - create:					Create a new network file
      - set-config:				Change network config parameters
      - add-host <name>:			Add a host
      - add-ssh-host <name> <host>:		Add a remote OpenWrt host via SSH
						(<host> can contain SSH options as well)
      - set-host <name>:			Change host settings
      - set-ssh-host <name> <host>:		Update local and remote host settings
      - add-service <name>:			Add a service
      - set-service <name>:			Change service settings
      - sign					Sign network data

     Flags:
      -p:					Print modified JSON instead of updating file

     Options:
      - config options (create, set-config):
	port=<val>				set tunnel port (default: ${defaults.port})
	pex_port=<val>				set peer-exchange port (default: ${defaults.pex_port}, 0: disabled)
	keepalive=<val>				set keepalive interval (seconds, 0: off, default: ${defaults.keepalive})
	stun=[+|-]<host:port>[,<host:port>...]	set/add/remove STUN servers
      - host options (add-host, add-ssh-host, set-host):
	key=<val>				set host public key (required for add-host)
	port=<val>				set host tunnel port number
	pex_port=<val>				set host peer-exchange port (default: network pex_port, 0: disabled)
	groups=[+|-]<val>[,<val>...]		set/add/remove groups that the host is a member of
	ipaddr=[+|-]<val>[,<val>...]		set/add/remove host ip addresses
	subnet=[+|-]<val>[,<val>...]		set/add/remove host announced subnets
	endpoint=<val>				set host endpoint address
	gateway=<name>				set host gateway (using name of other host)
     - ssh host options (add-ssh-host, set-ssh-host)
	auth_key=<key>				use <key> as public auth key on the remote host
	priv_key=<key>				use <key> as private host key on the remote host (default: generate a new key)
	interface=<name>			use <name> as interface in /etc/config/network on the remote host
	domain=<name>				use <name> as hosts file domain on the remote host (default: unet)
	connect=<val>[,<val>...]		set IP addresses that the host will contact for network updates
	tunnels=<ifname>:<service>[,...]	set active tunnel devices
	dht=0|1					set DHT enabled
     - service options (add-service, set-service):
	type=<val>				set service type (required for add-service)
	members=[+|-]<val>[,<val>...]		set/add/remove service member hosts/groups
     - vxlan service options (add-service, set-service):
	id=<val>				set VXLAN ID
	port=<val>				set VXLAN port
	mtu=<val>				set VXLAN device MTU
	forward_ports=[+|-]<val>[,<val>...]	set members allowed to receive broadcast/multicast/unknown-unicast
     - sign options:
	upload=<ip>[,<ip>...]			upload signed file to hosts

`;

function usage() {
	warn(usage_message);
	return 1;
}

if (length(ARGV) < 2)
	exit(usage());

let file = shift(ARGV);
let command = shift(ARGV);

const field_types = {
	int: function(object, name, val) {
		object[name] = int(val);
	},
	string: function(object, name, val) {
		object[name] = val;
	},
	array: function(object, name, val) {
		let op = substr(val, 0, 1);

		if (op == "+" || op == "-") {
			val = substr(val, 1);
			object[name] ??= [];
		} else {
			op = "=";
			object[name] = [];
		}

		let vals = split(val, ",");
		for (val in vals) {
			object[name] = filter(object[name], function(v) {
				return v != val
			});
			if (op != "-")
				push(object[name], val);
		}

		if (!length(object[name]))
			delete object[name];
	},
};

const service_field_types = {
	vxlan: {
		id: "int",
		port: "int",
		mtu: "int",
		forward_ports: "array",
	},
};

const ssh_script = `

set_list() {
	local field="$1"
	local val="$2"

	first=1
	for cur in $val; do
		if [ -n "$first" ]; then
			cmd=set
		else
			cmd=add_list
		fi
		uci $cmd "network.$INTERFACE.$field=$cur"
		first=
	done
}
set_interface_attrs() {
	[ -n "$AUTH_KEY" ] && uci set "network.$INTERFACE.auth_key=$AUTH_KEY"
	[ -n "$DHT" ] && uci set "network.$INTERFACE.dht=$DHT"
	set_list connect "$CONNECT"
	set_list tunnels "$TUNNELS"
	uci set "network.$INTERFACE.domain=$DOMAIN"
}

check_interface() {
	[ "$(uci -q get "network.$INTERFACE")" = "interface" -a "$(uci -q get "network.$INTERFACE.proto")" = "unet" ] && return 0
	uci batch <<EOF
set network.$INTERFACE=interface
set network.$INTERFACE.proto=unet
set network.$INTERFACE.device=$INTERFACE
EOF
}

check_interface_key() {
	key="$(uci -q get "network.$INTERFACE.key" | unet-tool -q -H -K -)"
	[ -n "$key" ] || {
		uci set "network.$INTERFACE.key=$(unet-tool -G)"
		key="$(uci get "network.$INTERFACE.key" | unet-tool -H -K -)"
	}
	echo "key=$key"
}

check_interface
check_interface_key
set_interface_attrs
uci commit
reload_config
ifup $INTERFACE
`;

let print_only = false;

function fetch_args() {
	for (let arg in ARGV) {
		let vals = match(arg, /^(.[[:alnum:]_-]*)=(.*)$/);
		assert(vals, `Invalid argument: ${arg}`);
		args[vals[1]] = vals[2]
	}
}

function set_field(typename, object, name, val) {
	if (!field_types[typename]) {
		warn(`Invalid type ${type}\n`);
		return;
	}

	if (type(val) != "string")
		return;

	if (val == "") {
		delete object[name];
		return;
	}

	field_types[typename](object, name, val);
}

function set_fields(object, list) {
	for (let f in list)
		set_field(list[f], object, f, args[f]);
}

function set_host(host) {
	set_fields(host, {
		key: "string",
		endpoint: "string",
		gateway: "string",
		port: "int",
		ipaddr: "array",
		subnet: "array",
		groups: "array",
	});
	set_field("int", host, "peer-exchange-port", args.pex_port);
}

function set_service(service) {
	set_fields(service, {
		type: "string",
		members: "array",
	});

	if (service_field_types[service.type])
		set_fields(service.config, service_field_types[service.type]);
}

function sync_ssh_host(host) {
	let interface = args.interface ?? "unet";
	let connect = replace(args.connect ?? "", ",", " ");
	let auth_key = args.auth_key;
	let tunnels = replace(replace(args.tunnels ?? "", ",", " "), ":", "=");
	let domain = args.domain ?? "unet";
	let dht;

	if (args.dht == "1" || args.dht == "0")
		dht = args.dht;
	else
		dht = "";

	if (!auth_key) {
		let fh = mkstemp();
		system(`${unet_tool} -q -P -K ${file}.key >&${fh.fileno()}`);
		fh.seek();
		auth_key = fh.read("line");
		fh.close();
		auth_key = replace(auth_key, "\n", "");
		if (auth_key == "") {
			warn("Could not read auth key\n");
			exit(1);
		}
	}

	let fh = mkstemp();
	fh.write(`INTERFACE='${interface}'\n`);
	fh.write(`CONNECT='${connect}'\n`);
	fh.write(`AUTH_KEY='${auth_key}'\n`);
	fh.write(`TUNNELS='${tunnels}'\n`);
	fh.write(`DOMAIN='${domain}'\n`);
	fh.write(`DHT='${dht}'\n`);
	fh.write(ssh_script);
	fh.flush();
	fh.seek();

	let fh2 = mkstemp();
	system(`ssh ${host} sh <&${fh.fileno()} >&${fh2.fileno()}`);
	fh.close();

	let data = {}, line;

	fh2.seek();
	while (line = fh2.read("line")) {
		let vals = match(line, /^(.[[:alnum:]_-]*)=(.*)\n$/);
		assert(vals, `Invalid argument: ${line}`);
		data[vals[1]] = vals[2]
	}
	fh2.close();

	assert(data.key, "Could not read host key from SSH host");

	args.key = data.key;
}

while (substr(ARGV[0], 0, 1) == "-") {
	let opt = shift(ARGV);
	if (opt == "--")
		break;
	else if (opt == "-p")
		print_only = true;
	else
		exit(usage());
}

let hostname, ssh_host, servicename;

if (command in [ "add-host", "set-host", "add-ssh-host", "set-ssh-host" ]) {
	hostname = shift(ARGV);
	assert(hostname, "Missing host name argument");
}

if (command in [ "add-ssh-host", "set-ssh-host" ]) {
	ssh_host = shift(ARGV);
	assert(ssh_host, "Missing SSH host/user argument");
}

if (command in [ "add-service", "set-service" ]) {
	servicename = shift(ARGV);
	assert(servicename, "Missing service name argument");
}

fetch_args();

if (command in [ "add-ssh-host", "set-ssh-host" ]) {
	sync_ssh_host(ssh_host);
	command = replace(command, "ssh-", "");
}

let net_data;

if (command == "create") {
	net_data = {
		config: {},
		hosts: {},
		services: {}
	};
} else {
	let fh = open(file);
	assert(fh, `Could not open input file ${file}`);

	try {
		net_data = json(fh);
	} catch(e) {
		assert(false, `Could not parse input file ${file}`);
	}
}

if (command == "create") {
	for (let key, val in defaults)
		args[key] ??= `${val}`;
	if (!access(`${file}.key`))
		system(`${unet_tool} -G > ${file}.key`);
}

if (command == "sign") {
	let ret = system(`${unet_tool} -S -K ${file}.key -o ${file}.bin ${file}`);
	if (ret != 0)
		exit(ret);

	if (args.upload) {
		for (let host in split(args.upload, ",")) {
			warn(`Uploading ${file}.bin to ${host}\n`);
			ret = system(`${unet_tool} -U ${host} -K ${file}.key ${file}.bin`);
			if (ret)
				warn("Upload failed\n");
		}
	}
	exit(0);
}

switch (command) {
case 'create':
case 'set-config':
	set_fields(net_data.config, {
		port: "int",
		keepalive: "int",
	});
	set_field("int", net_data.config, "peer-exchange-port", args.pex_port);
	set_field("array", net_data.config, "stun-servers", args.stun);
	break;

case 'add-host':
	net_data.hosts[hostname] = {};
	assert(args.key, "Missing host key");
	set_host(net_data.hosts[hostname]);
	break;

case 'set-host':
	assert(net_data.hosts[hostname], `Host '${hostname}' does not exist`);
	set_host(net_data.hosts[hostname]);
	break;

case 'add-service':
	net_data.services[servicename] = {
		config: {},
		members: [],
	};
	assert(args.type, "Missing service type");
	set_service(net_data.services[servicename]);
	break;

case 'set-service':
	assert(net_data.services[servicename], `Service '${servicename}' does not exist`);
	set_service(net_data.services[servicename]);
	break;

default:
	assert(false, "Unknown command");
}

const net_data_json = sprintf("%.J\n", net_data);

if (print_only)
	print(net_data_json);
else
	writefile(file, net_data_json);