#!/usr/bin/lua
-- Copyright (C) 2020-2023 mikma

local json = require("luci.jsonc")
require "ubus"
require "uloop"

local IFACE = 'wg'
local TIMEOUT = 1000

local wg_peers = {}
local wg_ll = {}

local function add_to_set(set, key)
	set[key] = true
end

local function remove_from_set(set, key)
	set[key] = nil
end

local function add_to_allowed_ips(allowed_ips, ip)
	add_to_set(allowed_ips, ip)
end

local function remove_from_allowed_ips(allowed_ips, ip)
	remove_from_set(allowed_ips, ip)
end

local function delete_ip(allowed_ips, ip)
	for k, v in pairs(allowed_ips) do
		if v == ip then
			allowed_ips[k] = nil
			print('deleted ' .. ip)
			return
		end
	end
	print('not deleted ' .. ip)
end

local function lookup_peer(iface, addr)
	local cmd = 'wg show ' .. iface .. ' allowed-ips'
	print('cmd: ' .. cmd .. ' addr: ' .. addr)
	local f = assert(io.popen(cmd, 'r'))

	peers = {}
	while true do
		local line = f:read('*line')
		if line == nil then break end

--		print('line: ' .. line)

		for peer, v in string.gmatch(line, "([^\t]+)\t([^t]+)") do
--			print(peer	.. '#' .. v)
			local allowed_ips = {}
			local found = false
			for allowed_ip in string.gmatch(v, "[^ ]+") do
--				print('addr: "' .. allowed_ip .. '"')
				add_to_allowed_ips(allowed_ips, allowed_ip)
				if addr == allowed_ip then
					print('found ' .. addr)
					found = true
				end
			end
			if found then
				f:close()
				return peer, allowed_ips
			end
			peers[peer] = allowed_ips
		end
	end

	f:close()
	return nil
end

local function set_allowed_ips(iface, peer_key, allowed_ips)
	str = nil

	for ip, _ in pairs(allowed_ips) do
		if str then
			str = str .. ',' .. ip
		else
			str = ip
		end
	end
	cmd = 'wg set ' .. iface .. ' peer ' .. peer_key .. ' allowed-ips ' .. str

	print('cmd: ' .. cmd)
	os.execute(cmd)
end

local function dhcpv4(method, msg)
	print("dhcpv4 " .. method)

	local iface = msg.interface
	local peer_addr = msg['peer-4o6']
	local ip = msg.ip .. '/32'

	print('ack: ' .. iface .. ' ' .. peer_addr .. ' ' .. ip)
	local peer, allowed_ips = lookup_peer(iface, peer_addr .. '/128')
	if not peer then
		return
	end
	print('peer: ' .. peer)
	if method == 'dhcp.ack' then
		add_to_allowed_ips(allowed_ips, ip)
	else if method == 'dhcp.release' then
		remove_from_allowed_ips(allowed_ips, ip)
	end end
	set_allowed_ips(iface, peer, allowed_ips)
end

local function dhcpv6(method, msg)
	print("dhcpv6 " .. method)

	local iface = msg.interface
	local peer_addr = msg.peer
	print('ack: ' .. iface .. ' ' .. peer_addr)
	local peer, allowed_ips = lookup_peer(iface, peer_addr .. '/128')

	if not peer then
		return
	end
	print('peer: ' .. peer)

	local add_ips = nil
	for _, addr in ipairs(msg.ips) do
		ip = addr.ip
		print('ip: ' .. ip)
		if method == 'dhcpv6.ack' then
			add_to_allowed_ips(allowed_ips, ip)
		else if method == 'dhcpv6.release' then
			remove_from_allowed_ips(allowed_ips, ip)
		end end
	end

	set_allowed_ips(iface, peer, allowed_ips)
end

local function dhcpv6_release(msg)
	print("dhcpv6_release")
end

local methods = {}
methods['dhcp.release'] = dhcpv4
methods['dhcp.ack'] = dhcpv4
methods['dhcpv6.release'] = dhcpv6
methods['dhcpv6.ack'] = dhcpv6

local function subscribe_dhcp(conn, iface)
	print("Subscribing to dhcp")
	local sub = {
		notify = function(msg, method)
			local ifname = ifname,
			print(method)
			print(json.stringify(msg))
			if ifname ~= iface then
				print('Unknown interface: ' .. ifname)
				return
			end
			action = methods[method]
			if action then
				action(method, msg)
			else
				print("Unknown method")
			end
		end,
	}
	conn:subscribe("dhcp", sub)
end

function do_tables_match( a, b )
    return table.concat(a) == table.concat(b)
end

-- Scan ipv6 leases
local function scan_leases(conn, iface, name)
	print("Scan: " .. name .. ' ' .. iface)
	local status = conn:call("dhcp", name, {})

	local peers = {}
	local peers_dirty = {}

	for _, v in pairs(status.device[iface].leases) do
		print(v)
		local peer_ll = v.peer

		if not peer_ll then
			peer_ll = v['peer-4o6']
		end
		print("peer=" .. peer_ll)

		local pos = string.find(peer_ll, '%%')
		if pos then
			peer_ll = string.sub(peer_ll, 1, pos - 1)
		end

		local peer_key = wg_ll[peer_ll]
		local wg_peer = wg_peers[peer_key]
		local peer = peers[peer_key]
		if not peer then
			peer = {}
			peers[peer_key] = peer
		end

		print("peer=" .. tostring(peer_ll) .. ' ' .. tostring(peer_key))
		local ipv4_addr = v['address']
		local ipv6_prefix = v['ipv6-prefix']
		local ipv6_addr = v['ipv6-addr']

		if ipv4_addr then
			local ip = ipv4_addr .. '/32'
			local status = 'found'
			if not wg_peer[ip] then
				status = 'not found'
				add_to_set(peers_dirty, peer_key)
				add_to_set(peer, ip)
			end

			print("address=" .. ip .. ' ' .. status)
		end

		if ipv6_addr then
			for _, addr in pairs(ipv6_addr) do
				local ip = addr.address .. '/128'
				local status = 'found'
				if not wg_peer[ip] then
					status = 'not found'
					add_to_set(peers_dirty, peer_key)
					add_to_set(peer, ip)
				end

				print("address=" .. ip .. ' ' .. status)
			end
		end

		if ipv6_prefix then
			for _, prefix in pairs(ipv6_prefix) do
				local ip = prefix.address .. '/' .. prefix['prefix-length']
				local status = 'found'
				if not wg_peer[ip] then
					status = 'not found'
					add_to_set(peers_dirty, peer_key)
					add_to_set(peer, ip)
				end

				print("prefix=" .. ip .. ' ' .. status)
			end
		end
	end

	for peer_key, _ in pairs(peers_dirty) do
		print('Dirty ' .. peer_key)

		local add_allowed_ips = peers[peer_key]
		local allowed_ips = wg_peers[peer_key]

		for allowed_ip, _ in pairs(add_allowed_ips) do
			print('Add ip ' .. allowed_ip)
			add_to_set(allowed_ips, allowed_ip)
		end

		set_allowed_ips(iface, peer_key, allowed_ips)
	end
	print("Done " .. name)
end

-- Scan wg allowed_ips and update wg_peers, and wg_ll
local function scan_wg(iface)
	print("Scan wg")
	local cmd = 'wg show ' .. iface .. ' allowed-ips'
	print('cmd: ' .. cmd)
	local f = assert(io.popen(cmd, 'r'))

	local peers = {}
	local ll = {}
	while true do
		local line = f:read('*line')
		if line == nil then break end

--		print('line: ' .. line)

		for peer, v in string.gmatch(line, "([^\t]+)\t([^t]+)") do
--			print(peer	.. '#' .. v)
			local allowed_ips = {}
			for allowed_ip in string.gmatch(v, "[^ ]+") do
				print('addr: "' .. allowed_ip .. '"')
				add_to_allowed_ips(allowed_ips, allowed_ip)

				print(string.sub(allowed_ip, 1, 6))
				if string.sub(allowed_ip, 1, 6) == 'fe80::' then
					local peer_ll = allowed_ip
					local pos = string.find(peer_ll, '/')
					if pos then
						peer_ll = string.sub(peer_ll, 1, pos - 1)
					end
					print("Found LL " .. peer_ll .. ' = ' .. peer)
					ll[peer_ll] = peer
				end
			end
			peers[peer] = allowed_ips
		end
	end

	wg_peers = peers
	wg_ll = ll
	f:close()
	print("Done wg")
end

local function listen_interface(conn, iface, timeout)
	local event = {}

	event['network.interface'] = function(msg)
		if msg.action == 'ifup' and msg.interface == iface then
			print("Up " .. iface)
			uloop.timer(function() scan_wg(iface); scan_lease(conn, iface, 'ipv4leases'); scan_leases(conn, iface, 'ipv6leases') end, timeout)
		end
	end
	conn:listen(event)
end

uloop.init()

local conn = ubus.connect()
if not conn then
	error("Failed to connect to ubusd")
end

scan_wg(IFACE)

scan_leases(conn, IFACE, 'ipv4leases')
scan_leases(conn, IFACE, 'ipv6leases')

listen_interface(conn, IFACE, TIMEOUT)
subscribe_dhcp(conn)

uloop.run()