--[[
Iptables file dumper for Alpine Wall
Copyright (C) 2012-2024 Kaarle Ritvanen
See LICENSE file for license details
]]--


local class = require('awall.class')
local family = require('awall.family')
local raise = require('awall.uerror').raise

local util = require('awall.util')
local printmsg = util.printmsg
local sortedkeys = util.sortedkeys


local lpc = require('lpc')
local posix = require('posix')
local startswith = require('stringy').startswith


local M = {}

local backupdir = '/var/run/awall'

local families = {
	inet={cmd='iptables', file='rules-save'},
	inet6={cmd='ip6tables', file='rules6-save'}
}

local builtin = {
	filter={'FORWARD', 'INPUT', 'OUTPUT'},
	mangle={'FORWARD', 'INPUT', 'OUTPUT', 'POSTROUTING', 'PREROUTING'},
	nat={'INPUT', 'OUTPUT', 'POSTROUTING', 'PREROUTING'},
	raw={'OUTPUT', 'PREROUTING'},
	security={'FORWARD', 'INPUT', 'OUTPUT'}
}

function M.isbuiltin(tbl, chain) return util.contains(builtin[tbl], chain) end


local BaseIPTablesRuleset = class()

function BaseIPTablesRuleset:init(context) self.context = context end

function BaseIPTablesRuleset:print()
	for _, family in ipairs(self.context:conffamilies()) do
		self:dumpfile(family, io.output())
		io.write('\n')
	end
end

function BaseIPTablesRuleset:dump(dir)
	for _, family in ipairs(self.context:conffamilies()) do
		local name = dir..'/'..families[family].file
		local file = io.open(name, 'w')
		if not file then raise('Cannot write to '..name) end
		self:dumpfile(family, file)
		file:close()
	end
end

function BaseIPTablesRuleset:restorecmd(family, test)
	local cmd = {families[family].cmd..'-restore'}
	if test then table.insert(cmd, '-t') end
	return table.unpack(cmd)
end

function BaseIPTablesRuleset:restore(test)
	for _, family in ipairs(self.context:actfamilies()) do
		local pid, stdin, stdout = lpc.run(self:restorecmd(family, test))
		stdout:close()
		self:dumpfile(family, stdin)
		stdin:close()
		assert(lpc.wait(pid) == 0)
	end
end

function BaseIPTablesRuleset:activate()
	self:flush()
	self:restore(false)
end

function BaseIPTablesRuleset:test() self:restore(true) end

function BaseIPTablesRuleset:flush() self.context:flush() end


M.IPTablesRuleset = class(BaseIPTablesRuleset)

function M.IPTablesRuleset:init(context)
	M.IPTablesRuleset.super(self):init(context)

	local function nestedtable(levels)
		return levels > 0 and setmetatable(
			{},
			{
				__index=function(t, k)
					t[k] = nestedtable(getmetatable(t).levels - 1)
					return t[k]
				end,
				levels=levels
			}
		) or {}
	end
	self.rules = nestedtable(3)
end

function M.IPTablesRuleset:dumpfile(family, iptfile)
	iptfile:write('# '..families[family].file..' generated by awall\n')
	local tables = self.rules[family]
	for _, tbl in sortedkeys(tables) do
		iptfile:write('*'..tbl..'\n')
		local chains = tables[tbl]
		for _, chain in sortedkeys(chains) do
			local policy = '-'
			if M.isbuiltin(tbl, chain) then
				policy = tbl == 'filter' and 'DROP' or 'ACCEPT'
			end
			iptfile:write(':'..chain..' '..policy..' [0:0]\n')
		end
		for _, chain in sortedkeys(chains) do
			for _, rule in ipairs(chains[chain]) do
				iptfile:write('-A '..chain..' '..rule..'\n')
			end
		end
		iptfile:write('COMMIT\n')
	end
end


M.PartialIPTablesRuleset = class(M.IPTablesRuleset)

function M.PartialIPTablesRuleset:restorecmd(family, test)
	local cmd = {M.PartialIPTablesRuleset.super(self):restorecmd(family, test)}
	table.insert(cmd, '-n')
	return table.unpack(cmd)
end

function M.PartialIPTablesRuleset:dumpfile(family, iptfile)
	local tables = self.rules[family]
	for tbl, chains in pairs(tables) do
		local builtins = {}
		for chain, _ in pairs(chains) do
			if startswith(chain, 'awall-') then
				local base = chain:sub(7, -1)
				if M.isbuiltin(tbl, base) then table.insert(builtins, base) end
			end
		end
		for _, chain in ipairs(builtins) do
			chains[chain] = {'-j awall-'..chain}
		end
	end
	M.PartialIPTablesRuleset.super(self):dumpfile(family, iptfile)
end

function M.PartialIPTablesRuleset:flush()
	for _, family in ipairs(self.context:actfamilies()) do
		local cmd = families[family].cmd
		for _, tbl in ipairs(self.context:acttables(family)) do
			if builtin[tbl] then
				local pid, stdin, stdout = lpc.run(cmd, '-t', tbl, '-S')
				stdin:close()
				local chains = {}
				local rules = {}
				for line in stdout:lines() do
					if startswith(line, '-N awall-') then
						table.insert(chains, line:sub(4, -1))
					else
						local chain, target = line:match(
							'^%-A (%u+) %-j (awall%-%u+)$'
						)
						if chain then
							table.insert(rules, {chain, '-j', target})
						end
					end
				end
				stdout:close()
				assert(lpc.wait(pid) == 0)

				local function exec(...)
					assert(util.execute(cmd, '-t', tbl, table.unpack{...}) == 0)
				end
				for _, rule in ipairs(rules) do
					exec('-D', table.unpack(rule))
				end
				for _, opt in ipairs{'-F', '-X'} do
					for _, chain in ipairs(chains) do exec(opt, chain) end
				end
			end
		end
	end
end


local CurrentRuleset = class(BaseIPTablesRuleset)

function CurrentRuleset:dumpfile(family, iptfile)
	local pid, stdin, stdout = lpc.run(families[family].cmd..'-save')
	stdin:close()
	for line in stdout:lines() do iptfile:write(line..'\n') end
	stdout:close()
	assert(lpc.wait(pid) == 0)
end


local BackupRuleset = class(BaseIPTablesRuleset)

function BackupRuleset:dumpfile(family, iptfile)
	for line in io.lines(backupdir..'/'..families[family].file) do
		iptfile:write(line..'\n')
	end
end


M.IPTables = class()

function M.IPTables:init(famlist, checkactive)
	if famlist then
		for _, family in ipairs(famlist) do
			if not families[family] then
				raise('Unknown address family configured: '..family)
			end
		end
		self.families = famlist
	end

	self.checkactive = checkactive
	self._acttables = {}
end

function M.IPTables:conffamilies() return self.families or family.ALL end

function M.IPTables:usablefamilies()
	if not self._usablefamilies then
		self._usablefamilies = {}
		for _, fam in util.listpairs(self:conffamilies()) do
			if family.isactive(fam) then table.insert(self._usablefamilies, fam)
			elseif self.families then
				printmsg(
					'Warning: address family not supported by kernel: '..fam
				)
			end
		end
	end
	return self._usablefamilies
end

function M.IPTables:acttables(family)
	if not self._acttables[family] then
		self._acttables[family] = {}
		local pid, stdin, stdout = lpc.run(families[family].cmd..'-save')
		stdin:close()
		for line in stdout:lines() do
			local tbl = string.match(line, "^%*(.*)")
			if tbl then table.insert(self._acttables[family], tbl) end
		end
		stdout:close()
		assert(lpc.wait(pid) == 0)
	end
	return self._acttables[family]
end

function M.IPTables:actfamilies()
	if not self._actfamilies then
		if self.checkactive then
			self._actfamilies = {}
			for _, family in ipairs(self:usablefamilies()) do
				if #self:acttables(family) > 0 then
					table.insert(self._actfamilies, family)
				else printmsg('Warning: firewall not enabled for '..family) end
			end
		else self._actfamilies = self:usablefamilies() end
	end
	return self._actfamilies
end

function M.IPTables:isenabled() return #self:actfamilies() > 0 end

function M.IPTables:backup()
	posix.mkdir(backupdir)
	CurrentRuleset(self):dump(backupdir)
end

function M.IPTables:revert() BackupRuleset(self):activate() end

function M.IPTables:flush()
	local empty = M.IPTablesRuleset(self)
	for _, family in pairs(self:actfamilies()) do
		for _, tbl in ipairs(self:acttables(family)) do
			if builtin[tbl] then
				for _, chain in ipairs(builtin[tbl]) do
					empty.rules[family][tbl][chain] = {}
				end
			else printmsg('Warning: not flushing unknown table: '..tbl) end
		end
	end
	empty:restore(false)
end


return M

-- vim: ts=4
