#!/usr/bin/lua

require("luci.util")
require("luci.model.uci")
require("luci.sys.iptparser")

-- Init state session
local uci = luci.model.uci.cursor_state()
local ipt = luci.sys.iptparser.IptParser()

local splash_interfaces = { }
local limit_up = 0
local limit_down = 0
local limit_down_burst = 0

function main(argv)
	local cmd = argv[1]
	local arg = argv[2]

	limit_up = tonumber(uci:get("luci_splash", "general", "limit_up") or 0)
	limit_down = tonumber(uci:get("luci_splash", "general", "limit_down") or 0)
	limit_down_burst = tonumber(uci:get("luci_splash", "general", "limit_down_burst") or limit_down * 2)
	
	uci:foreach("luci_splash", "iface", function(s)
		if s.network then
			splash_interfaces[#splash_interfaces+1] = uci:get("network", s.network, "ifname")
		end
	end)

	if cmd == "status" and arg then
		if islisted("whitelist", arg) then
			print("whitelisted")
		elseif islisted("blacklist", arg) then
			print("blacklisted")
		else		
			local lease = haslease(arg)
			if lease and lease.kicked then
				print("kicked")
			elseif lease then
				print("lease")
			else
				print("unknown")
			end
		end
		os.exit(0)
	elseif cmd == "add" and arg then
		if not haslease(arg) then
			add_lease(arg)
		else
			print("already leased!")
			os.exit(2)
		end
		os.exit(0)
	elseif cmd == "remove" and arg then
		remove_lease(arg)
		os.exit(0)		
	elseif cmd == "sync" then
		sync()
		os.exit(0)
	else
		print("Usage: " .. argv[0] .. " <status|add|remove|sync> [MAC]")
		os.exit(1)	
	end
end

-- Add a lease to state and invoke add_rule
function add_lease(mac)
	uci:section("luci_splash", "lease", nil, {
		mac = mac,
		start = os.time()
	})
	add_rule(mac)
	
	uci:save("luci_splash")
end


-- Remove a lease from state and invoke remove_rule
function remove_lease(mac)
	mac = mac:lower()
	remove_rule(mac)

	uci:delete_all("luci_splash", "lease",
		function(s) return ( s.mac:lower() == mac ) end)
		
	uci:save("luci_splash")
end


-- Add an iptables rule
function add_rule(mac)
	local a, b, c, d, e, f = mac:match("(%w+):(%w+):(%w+):(%w+):(%w+):(%w+)")
	local mac_pre  = "%s%s" %{ a, b }
	local mac_post = "%s%s%s%s" %{ c, d, e, f }
	local handle   = f

	if limit_up > 0 and limit_down > 0 then
		os.execute("iptables -t mangle -I luci_splash_mark -m mac --mac-source %q -j MARK --set-mark 79" % mac)

		for _, i in ipairs(splash_interfaces) do
			os.execute("tc filter add dev %q parent 77:0 protocol ip prio 2 " % i .. 
				"handle ::%q u32 " % handle ..
				"match u16 0x0800 0xFFFF at -2 match u32 0x%q 0xFFFFFFFF at -12 " % mac_post ..
				"match u16 0x%q 0xFFFF at -14 flowid 77:10" % mac_pre)
		end
	end

	os.execute("iptables -t filter -I luci_splash_counter -m mac --mac-source %q -j RETURN" % mac)
	return os.execute("iptables -t nat -I luci_splash_leases -m mac --mac-source %q -j RETURN" % mac)
end


-- Remove an iptables rule
function remove_rule(mac)
	local handle = mac:match("%w+:%w+:%w+:%w+:%w+:(%w+)")

	local function ipt_delete_foreach(args)
		for _, r in ipairs(ipt:find(args)) do
			if r.options and #r.options >= 2 and r.options[1] == "MAC" and
			   r.options[2]:lower() == mac:lower()
			then
				os.execute("iptables -t %q -D %q -m mac --mac-source %q %s 2>/dev/null"
					%{ r.table, r.chain, mac,
						r.target == "MARK" and "-j MARK --set-mark 79" or
							r.target and "-j %q" % r.target or "" })
			end
		end
	end

	ipt_delete_foreach({table="filter", chain="luci_splash_counter"})
	ipt_delete_foreach({table="mangle", chain="luci_splash_mark"})
	ipt_delete_foreach({table="nat", chain="luci_splash_leases"})

	for _, i in ipairs(splash_interfaces) do
		os.execute("tc filter del dev %q parent 77:0 protocol ip prio 2 " % i .. 
			"handle 800::%q u32 2>/dev/null" % handle)
	end

	ipt:resync()
end


-- Check whether a MAC-Address is listed in the lease state list
function haslease(mac)
	mac = mac:lower()
	local lease = nil

	uci:foreach("luci_splash", "lease",
		function (section)
			if section.mac:lower() == mac then
				lease = section
			end
		end)

	return lease
end


-- Check whether a MAC-Address is in given list
function islisted(what, mac)
	mac = mac:lower()

	uci:foreach("luci_splash", what,
		function (section)
			if section.mac:lower() == mac then
				stat = true
				return
			end
		end)

	return false
end


-- Returns a list of MAC-Addresses for which a rule is existing
function listrules()
	local macs = { }
	for i, r in ipairs(ipt:find({table="nat", chain="luci_splash_leases"})) do
		if r.options and #r.options >= 2 and r.options[1] == "MAC" then
			macs[r.options[2]:lower()] = true
		end
	end
	return luci.util.keys(macs)
end


-- Synchronise leases, remove abandoned rules
function sync()
	local written = {}
	local time = os.time()
	
	-- Current leases in state files
	local leases = uci:get_all("luci_splash")
	
	-- Convert leasetime to seconds
	local leasetime = tonumber(uci:get("luci_splash", "general", "leasetime")) * 3600
	
	-- Clean state file
	uci:load("luci_splash")
	uci:revert("luci_splash")
	
	
	-- For all leases
	for k, v in pairs(leases) do
		if v[".type"] == "lease" then
			if os.difftime(time, tonumber(v.start)) > leasetime then
				-- Remove expired
				remove_rule(v.mac, v.kicked)
			else
				-- Rewrite state
				uci:section("luci_splash", "lease", nil, {		
					mac    = v.mac,
					start  = v.start,
					kicked = v.kicked
				})
				written[v.mac:lower()] = 1
			end
		elseif v[".type"] == "whitelist" or v[".type"] == "blacklist" then
			written[v.mac:lower()] = 1
		end
	end
	
	
	-- Delete rules without state
	for i, r in ipairs(listrules()) do
		if #r > 0 and not written[r:lower()] then
			remove_rule(r)
		end
	end
	
	uci:save("luci_splash")
end

main(arg)