add PoC WireGuard script

1 files changed, 339 insertions, 0 deletions

diff --git a/dhcp-subscribe.lua b/dhcp-subscribe.lua
new file mode 100755
index 0000000..bceac43
--- /dev/null
+++ b/dhcp-subscribe.lua
@@ -0,0 +1,339 @@
+#!/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()