local lfs = require "lfs";

local resolve_relative_path = require "core.configmanager".resolve_relative_path;

local logger = require "util.logger".init;

local it = require "util.iterators";

local set = require "util.set";

-- [definition_type] = definition_factory(param)

local definitions = module:shared("definitions");

-- When a definition instance has been instantiated, it lives here

-- [definition_type][definition_name] = definition_object

local active_definitions = {

	ZONE = {

		-- Default zone that includes all local hosts

		["$local"] = setmetatable({}, { __index = prosody.hosts });

	};

};

local default_chains = {

	preroute = {

		type = "event";

		priority = 0.1;

		"pre-message/bare", "pre-message/full", "pre-message/host";

		"pre-presence/bare", "pre-presence/full", "pre-presence/host";

		"pre-iq/bare", "pre-iq/full", "pre-iq/host";

	};

	deliver = {

		type = "event";

		priority = 0.1;

		"message/bare", "message/full", "message/host";

		"presence/bare", "presence/full", "presence/host";

		"iq/bare", "iq/full", "iq/host";

	};

	deliver_remote = {

		type = "event"; "route/remote";

		priority = 0.1;

	};

};

local extra_chains = module:get_option("firewall_extra_chains", {});

local chains = {};

for k,v in pairs(default_chains) do

	chains[k] = v;

end

for k,v in pairs(extra_chains) do

	chains[k] = v;

end

-- Returns the input if it is safe to be used as a variable name, otherwise nil

function idsafe(name)

	return name:match("^%a[%w_]*$");

end

local meta_funcs = {

	bare = function (code)

		return "jid_bare("..code..")", {"jid_bare"};

	end;

	node = function (code)

		return "(jid_split("..code.."))", {"jid_split"};

	end;

	host = function (code)

		return "(select(2, jid_split("..code..")))", {"jid_split"};

	end;

	resource = function (code)

		return "(select(3, jid_split("..code..")))", {"jid_split"};

	end;

};

-- Run quoted (%q) strings through this to allow them to contain code. e.g.: LOG=Received: $(stanza:top_tag())

function meta(s, deps, extra)

	return (s:gsub("$(%b())", function (expr)

			expr = expr:gsub("\\(.)", "%1");

			return [["..tostring(]]..expr..[[).."]];

		end)

		:gsub("$(%b<>)", function (expr)

			expr = expr:sub(2,-2);

			local default = "<undefined>";

			expr = expr:gsub("||(%b\"\")$", function (default_string)

				default = stripslashes(default_string:sub(2,-2));

				return "";

			end);

			local func_chain = expr:match("|[%w|]+$");

			if func_chain then

				expr = expr:sub(1, -1-#func_chain);

			end

			local code;

			if expr:match("^@") then

				-- Skip stanza:find() for simple attribute lookup

				local attr_name = expr:sub(2);

				if deps and (attr_name == "to" or attr_name == "from" or attr_name == "type") then

					-- These attributes may be cached in locals

					code = attr_name;

					table.insert(deps, attr_name);

				else

					code = "stanza.attr["..("%q"):format(attr_name).."]";

				end

			elseif expr:match("^%w+#$") then

				code = ("stanza:get_child_text(%q)"):format(expr:sub(1, -2));

			else

				code = ("stanza:find(%q)"):format(expr);

			end

			if func_chain then

				for func_name in func_chain:gmatch("|(%w+)") do

					-- to/from are already available in local variables, use those if possible

					if (code == "to" or code == "from") and func_name == "bare" then

						code = "bare_"..code;

						table.insert(deps, code);

					elseif (code == "to" or code == "from") and (func_name == "node" or func_name == "host" or func_name == "resource") then

						table.insert(deps, "split_"..code);

						code = code.."_"..func_name;

					else

						assert(meta_funcs[func_name], "unknown function: "..func_name);

						local new_code, new_deps = meta_funcs[func_name](code);

						code = new_code;

						if new_deps and #new_deps > 0 then

							assert(deps, "function not supported here: "..func_name);

							for _, dep in ipairs(new_deps) do

								table.insert(deps, dep);

							end

						end

					end

				end

			end

			return "\"..tostring("..code.." or "..("%q"):format(default)..")..\"";

		end)

		:gsub("$$(%a+)", extra or {})

		:gsub([[^""%.%.]], "")

		:gsub([[%.%.""$]], ""));

end

function metaq(s, ...)

	return meta(("%q"):format(s), ...);

end

local escape_chars = {

	a = "\a", b = "\b", f = "\f", n = "\n", r = "\r", t = "\t",

	v = "\v", ["\\"] = "\\", ["\""] = "\"", ["\'"] = "\'"

};

function stripslashes(s)

	return (s:gsub("\\(.)", escape_chars));

end

-- Dependency locations:

-- <type lib>

-- <type global>

-- function handler()

--   <local deps>

--   if <conditions> then

--     <actions>

--   end

-- end

local available_deps = {

	st = { global_code = [[local st = require "util.stanza";]]};

	it = { global_code = [[local it = require "util.iterators";]]};

	it_count = { global_code = [[local it_count = it.count;]], depends = { "it" } };

	current_host = { global_code = [[local current_host = module.host;]] };

	jid_split = {

		global_code = [[local jid_split = require "util.jid".split;]];

	};

	jid_bare = {

		global_code = [[local jid_bare = require "util.jid".bare;]];

	};

	to = { local_code = [[local to = stanza.attr.to or jid_bare(session.full_jid);]]; depends = { "jid_bare" } };

	from = { local_code = [[local from = stanza.attr.from;]] };

	type = { local_code = [[local type = stanza.attr.type;]] };

	name = { local_code = [[local name = stanza.name;]] };

	split_to = { -- The stanza's split to address

		depends = { "jid_split", "to" };

		local_code = [[local to_node, to_host, to_resource = jid_split(to);]];

	};

	split_from = { -- The stanza's split from address

		depends = { "jid_split", "from" };

		local_code = [[local from_node, from_host, from_resource = jid_split(from);]];

	};

	bare_to = { depends = { "jid_bare", "to" }, local_code = "local bare_to = jid_bare(to)"};

	bare_from = { depends = { "jid_bare", "from" }, local_code = "local bare_from = jid_bare(from)"};

	group_contains = {

		global_code = [[local group_contains = module:depends("groups").group_contains]];

	};

	is_admin = { global_code = [[local is_admin = require "core.usermanager".is_admin;]]};

	core_post_stanza = { global_code = [[local core_post_stanza = prosody.core_post_stanza;]] };

	zone = { global_code = function (zone)

		assert(idsafe(zone), "Invalid zone name: "..zone);

		return ("local zone_%s = zones[%q] or {};"):format(zone, zone);

	end };

	date_time = { global_code = [[local os_date = os.date]]; local_code = [[local current_date_time = os_date("*t");]] };

	time = { local_code = function (what)

		local defs = {};

		for field in what:gmatch("%a+") do

			table.insert(defs, ("local current_%s = current_date_time.%s;"):format(field, field));

		end

		return table.concat(defs, " ");

	end, depends = { "date_time" }; };

	timestamp = { global_code = [[local get_time = require "socket".gettime;]]; local_code = [[local current_timestamp = get_time();]]; };

	globalthrottle = {

		global_code = function (throttle)

			assert(idsafe(throttle), "Invalid rate limit name: "..throttle);

			assert(active_definitions.RATE[throttle], "Unknown rate limit: "..throttle);

			return ("local global_throttle_%s = rates.%s:single();"):format(throttle, throttle);

		end;

	};

	multithrottle = {

		global_code = function (throttle)

			assert(pcall(require, "util.cache"), "Using LIMIT with 'on' requires Prosody 0.10 or higher");

			assert(idsafe(throttle), "Invalid rate limit name: "..throttle);

			assert(active_definitions.RATE[throttle], "Unknown rate limit: "..throttle);

			return ("local multi_throttle_%s = rates.%s:multi();"):format(throttle, throttle);

		end;

	};

	rostermanager = {

		global_code = [[local rostermanager = require "core.rostermanager";]];

	};

	roster_entry = {

		local_code = [[local roster_entry = (to_node and rostermanager.load_roster(to_node, to_host) or {})[bare_from];]];

		depends = { "rostermanager", "split_to", "bare_from" };

	};

	list = { global_code = function (list)

			assert(idsafe(list), "Invalid list name: "..list);

			assert(active_definitions.LIST[list], "Unknown list: "..list);

			return ("local list_%s = lists[%q];"):format(list, list);

		end

	};

	search = {

		local_code = function (search_name)

			local search_path = assert(active_definitions.SEARCH[search_name], "Undefined search path: "..search_name);

			return ("local search_%s = tostring(stanza:find(%q) or \"\")"):format(search_name, search_path);

		end;

	};

	pattern = {

		local_code = function (pattern_name)

			local pattern = assert(active_definitions.PATTERN[pattern_name], "Undefined pattern: "..pattern_name);

			return ("local pattern_%s = %q"):format(pattern_name, pattern);

		end;

	};

	tokens = {

		local_code = function (search_and_pattern)

			local search_name, pattern_name = search_and_pattern:match("^([^%-]+)-(.+)$");

			local code = ([[local tokens_%s_%s = {};

			if search_%s then

				for s in search_%s:gmatch(pattern_%s) do

					tokens_%s_%s[s] = true;

				end

			end

			]]):format(search_name, pattern_name, search_name, search_name, pattern_name, search_name, pattern_name);

			return code, { "search:"..search_name, "pattern:"..pattern_name };

		end;

	};

	scan_list = {

		global_code = [[local function scan_list(list, items) for item in pairs(items) do if list:contains(item) then return true; end end end]];

	}

};

local function include_dep(dependency, code)

	local dep, dep_param = dependency:match("^([^:]+):?(.*)$");

	local dep_info = available_deps[dep];

	if not dep_info then

		module:log("error", "Dependency not found: %s", dep);

		return;

	end

	if code.included_deps[dependency] ~= nil then

		if code.included_deps[dependency] ~= true then

			module:log("error", "Circular dependency on %s", dep);

		end

		return;

	end

	code.included_deps[dependency] = false; -- Pending flag (used to detect circular references)

	for _, dep_dep in ipairs(dep_info.depends or {}) do

		include_dep(dep_dep, code);

	end

	if dep_info.global_code then

		if dep_param ~= "" then

			local global_code, deps = dep_info.global_code(dep_param);

			if deps then

				for _, dep_dep in ipairs(deps) do

					include_dep(dep_dep, code);

				end

			end

			table.insert(code.global_header, global_code);

		else

			table.insert(code.global_header, dep_info.global_code);

		end

	end

	if dep_info.local_code then

		if dep_param ~= "" then

			local local_code, deps = dep_info.local_code(dep_param);

			if deps then

				for _, dep_dep in ipairs(deps) do

					include_dep(dep_dep, code);

				end

			end

			table.insert(code, "\n\t\t-- "..dep.."\n\t\t"..local_code.."\n");

		else

			table.insert(code, "\n\t\t-- "..dep.."\n\t\t"..dep_info.local_code.."\n");

		end

	end

	code.included_deps[dependency] = true;

end

local definition_handlers = module:require("definitions");

local condition_handlers = module:require("conditions");

local action_handlers = module:require("actions");

local function new_rule(ruleset, chain)

	assert(chain, "no chain specified");

	local rule = { conditions = {}, actions = {}, deps = {} };

	table.insert(ruleset[chain], rule);

	return rule;

end

local function parse_firewall_rules(filename)

	local line_no = 0;

	local function errmsg(err)

		return "Error compiling "..filename.." on line "..line_no..": "..err;

	end

	local ruleset = {

		deliver = {};

	};

	local chain = "deliver"; -- Default chain

	local rule;

	local file, err = io.open(filename);

	if not file then return nil, err; end

	local state; -- nil -> "rules" -> "actions" -> nil -> ...

	local line_hold;

	for line in file:lines() do

		line = line:match("^%s*(.-)%s*$");

		if line_hold and line:sub(-1,-1) ~= "\\" then

			line = line_hold..line;

			line_hold = nil;

		elseif line:sub(-1,-1) == "\\" then

			line_hold = (line_hold or "")..line:sub(1,-2);

		end

		line_no = line_no + 1;

		if line_hold or line:find("^[#;]") then -- luacheck: ignore 542

			-- No action; comment or partial line

		elseif line == "" then

			if state == "rules" then

				return nil, ("Expected an action on line %d for preceding criteria")

					:format(line_no);

			end

			state = nil;

		elseif not(state) and line:sub(1, 2) == "::" then

			chain = line:gsub("^::%s*", "");

			local chain_info = chains[chain];

			if not chain_info then

				if chain:match("^user/") then

					chains[chain] = { type = "event", priority = 1, pass_return = false };

				else

					return nil, errmsg("Unknown chain: "..chain);

				end

			elseif chain_info.type ~= "event" then

				return nil, errmsg("Only event chains supported at the moment");

			end