module:set_global();

local jid_bare, jid_host = require "util.jid".bare, require "util.jid".host;

local st = require "util.stanza";

local xmlns_muc_user = "http://jabber.org/protocol/muc#user";

local trusted_services = module:get_option_inherited_set("muc_ban_ip_trusted_services", {});

local trust_local_restricted_services = module:get_option_boolean("muc_ban_ip_trust_local_restricted_services", true);

local ip_bans = module:shared("bans");

local full_sessions = prosody.full_sessions;

local function is_local_restricted_service(host)

	local muc_service = prosody.hosts[host] and prosody.hosts[host].modules.muc;

	if muc_service and module:context(host):get_option("restrict_room_creation") ~= nil then -- COMPAT: May need updating post-0.12

		return true;

	end

	return false;

end

local function ban_ip(session, from)

	local ip = session.ip;

	if not ip then

		module:log("warn", "Failed to ban IP (IP unknown) for %s", session.full_jid);

		return;

	end

	local from_host = jid_host(from);

	if trusted_services:contains(from_host) or (trust_local_restricted_services and is_local_restricted_service(from_host)) then

		from = from_host; -- Ban from entire host

	end

	local banned_from = ip_bans[ip];

	if not banned_from then

		banned_from = {};

		ip_bans[ip] = banned_from;

	end

	banned_from[from] = true;

	module:log("debug", "Added ban for IP address %s from %s", ip, from);

end

local function check_for_incoming_ban(event)

	local stanza = event.stanza;

	local to_session = full_sessions[stanza.attr.to];

	if to_session then

		local directed = to_session.directed;

		local from = stanza.attr.from;

		if directed and directed[from] and stanza.attr.type == "unavailable" then

			-- This is a stanza from somewhere we sent directed presence to (may be a MUC)

			local x = stanza:get_child("x", xmlns_muc_user);

			if x then

				for status in x:childtags("status") do

					if status.attr.code == '301' then

						ban_ip(to_session, jid_bare(from));

					end

				end

			end

		end

	end

end

local function check_for_ban(event)

	local origin, stanza = event.origin, event.stanza;

	local ip = origin.ip;

	local to, to_host = jid_bare(stanza.attr.to), jid_host(stanza.attr.to);

	if ip_bans[ip] and (ip_bans[ip][to] or ip_bans[ip][to_host]) then

		(origin.log or module._log)("debug", "IP banned: %s is banned from %s", ip, to)

		if stanza.attr.type ~= "error" then

			origin.send(st.error_reply(stanza, "auth", "forbidden")

				:tag("x", { xmlns = xmlns_muc_user })

					:tag("status", { code = '301' }));

		end

		return true;

	end

	(origin.log or module._log)("debug", "IP not banned: %s from %s", ip, to)

end

function module.add_host(module)

	module:hook("presence/full", check_for_incoming_ban, 100);

	module:hook("pre-presence/full", check_for_ban, 100);

end