#!/usr/bin/env lua

local json = require "luci.jsonc"
local util = require "luci.util"
local sys = require "luci.sys"
local io = require "io"
local uci = require "uci"
local fs = require "nixio.fs"

local methods = {
	generateKeyPair = {
		call = function()
			local prv = sys.exec("wg genkey 2>/dev/null"):sub(1, -2)
			local pub = sys.exec("echo %s | wg pubkey 2>/dev/null" % util.shellquote(prv)):sub(1, -2)

			return {keys = {priv = prv, pub = pub}}
		end
	},
	getPublicAndPrivateKeyFromPrivate = {
		args = {privkey = "privkey"},
		call = function(args)
			local pubkey = sys.exec("echo %s | wg pubkey 2>/dev/null" % util.shellquote(args.privkey)):sub(1, -2)

			return {keys = {priv = args.privkey, pub = pubkey}}
		end
	},
	generateQrCode = {
		args = {privkey = "privkey", psk = "psk", allowed_ips = {"allowed_ips"}},
		call = function(args)
			local qr_code

			if fs.access("/usr/bin/qrencode") then
				local psk = args.psk
				local listen_port = args.listen_port
				local allowed_ips = args.allowed_ips

				local pubkey = sys.exec("echo %s | wg pubkey 2>/dev/null" % util.shellquote(args.privkey)):sub(1, -2)
				local client_privkey = sys.exec("wg genkey 2>/dev/null"):sub(1, -2)

				local iface_qr = {
					"[Interface]",
					"PrivateKey = " .. client_privkey,
				}

				local peer_qr = {
					"[Peer]",
					"PublicKey = " .. pubkey,
				}

				if not allowed_ips or next(allowed_ips) == nil then
					allowed_ips = {"0.0.0.0/0", "::/0"}
				end
				table.insert(peer_qr, "AllowedIPs = " .. table.concat(allowed_ips, ", "))

				if psk then
					table.insert(peer_qr, "PresharedKey = " .. psk)
				end

				qr_enc = table.concat(iface_qr, "\n") .. "\n\n" .. table.concat(peer_qr, "\n")
				qr_code = sys.exec("/usr/bin/qrencode --inline --8bit --type=SVG --output=- %s 2>/dev/null" % util.shellquote(qr_enc))
			end

			return {qr_code = qr_code}
		end
	},
	getWgInstances = {
		call = function()
			local data = {}
			local last_device = ""
			local qr_pubkey = {}

			local wg_dump = io.popen("wg show all dump 2>/dev/null")
			if wg_dump then
				local line
				for line in wg_dump:lines() do
					local line = string.split(line, "\t")
					if not (last_device == line[1]) then
						last_device = line[1]
						data[line[1]] = {
							name = line[1],
							public_key = line[3],
							listen_port = line[4],
							fwmark = line[5],
							peers = {}
						}
						if not line[3] or line[3] == "" or line[3] == "(none)" then
							qr_pubkey[line[1]] = ""
						else
							qr_pubkey[line[1]] = "PublicKey = " .. line[3]
						end
					else
						local peer_name
						local cur = uci.cursor()

						cur:foreach(
							"network",
							"wireguard_" .. line[1],
							function(s)
								if s.public_key == line[2] then
									peer_name = s.description
								end
							end
						)

						local peer = {
							name = peer_name,
							public_key = line[2],
							endpoint = line[4],
							allowed_ips = {},
							latest_handshake = line[6],
							transfer_rx = line[7],
							transfer_tx = line[8],
							persistent_keepalive = line[9]
						}

						if not (line[4] == "(none)") then
							local ipkey, ipvalue
							for ipkey, ipvalue in pairs(string.split(line[5], ",")) do
								if #ipvalue > 0 then
									table.insert(peer["allowed_ips"], ipvalue)
								end
							end
						end

						table.insert(data[line[1]].peers, peer)
					end
				end
			end

			return data
		end
	}
}

local function parseInput()
	local parse = json.new()
	local done, err

	while true do
		local chunk = io.read(4096)
		if not chunk then
			break
		elseif not done and not err then
			done, err = parse:parse(chunk)
		end
	end

	if not done then
		print(json.stringify({error = err or "Incomplete input"}))
		os.exit(1)
	end

	return parse:get()
end

local function validateArgs(func, uargs)
	local method = methods[func]
	if not method then
		print(json.stringify({error = "Method not found"}))
		os.exit(1)
	end

	if type(uargs) ~= "table" then
		print(json.stringify({error = "Invalid arguments"}))
		os.exit(1)
	end

	uargs.ubus_rpc_session = nil

	local k, v
	local margs = method.args or {}
	for k, v in pairs(uargs) do
		if margs[k] == nil or (v ~= nil and type(v) ~= type(margs[k])) then
			print(json.stringify({error = "Invalid arguments"}))
			os.exit(1)
		end
	end

	return method
end

if arg[1] == "list" then
	local _, method, rv = nil, nil, {}
	for _, method in pairs(methods) do
		rv[_] = method.args or {}
	end
	print((json.stringify(rv):gsub(":%[%]", ":{}")))
elseif arg[1] == "call" then
	local args = parseInput()
	local method = validateArgs(arg[2], args)
	local result, code = method.call(args)
	print((json.stringify(result):gsub("^%[%]$", "{}")))
	os.exit(code or 0)
end