local s2smanager = require "core.s2smanager";
local config = require "core.configmanager";
local nameprep = require "util.encodings".stringprep.nameprep;

local s2s_blacklist = module:get_option_array("s2s_blacklist");
local s2s_whitelist = module:get_option_array("s2s_whitelist");
local s2s_enable_blackwhitelist = module:get_option_string("s2s_enable_blackwhitelist");
local is_blacklist_enabled = false;
local is_whitelist_enabled = false;

if s2s_enable_blackwhitelist == "blacklist" then
	if type(s2s_blacklist) == "table" then
		is_blacklist_enabled = true;
		module:log("debug", "s2s-blacklist is enabled");
		local count=#s2s_blacklist;
                for i=1,count do
			module:log("debug", "s2s-blacklist adding [%s]", s2s_blacklist[i]);
		end
	end
elseif s2s_enable_blackwhitelist == "whitelist" then
	if type(s2s_whitelist) == "table" then
		is_whitelist_enabled = true;
		module:log("debug", "s2s-whitelist is enabled");
                local count=#s2s_whitelist;
                for i=1,count do
                        module:log("debug", "s2s-whitelist adding [%s]", s2s_whitelist[i]);
                end
	end
end

local function reload_list()
	s2s_blacklist = module:get_option_array("s2s_blacklist");
	s2s_whitelist = module:get_option_array("s2s_whitelist");
	s2s_enable_blackwhitelist = module:get_option_string("s2s_enable_blackwhitelist");

	if s2s_enable_blackwhitelist == "blacklist" then
        	if type(s2s_blacklist) == "table" then
                	is_blacklist_enabled = true;
                	module:log("debug", "s2s-blacklist is enabled");
                	local count=#s2s_blacklist;
                	for i=1,count do
                        	module:log("debug", "s2s-blacklist adding [%s]", s2s_blacklist[i]);
                	end
        	end
	elseif s2s_enable_blackwhitelist == "whitelist" then
        	if type(s2s_whitelist) == "table" then
                	is_whitelist_enabled = true;
                	module:log("debug", "s2s-whitelist is enabled");
                	local count=#s2s_whitelist;
                	for i=1,count do
                        	module:log("debug", "s2s-whitelist adding [%s]", s2s_whitelist[i]);
                	end
        	end
	end
end

local _make_connect = s2smanager.make_connect;
function s2smanager.make_connect(session, connect_host, connect_port)
  local host = session.to_host;
  if not session.s2sValidation then
        if (host and is_blacklist_enabled == true) then
                local count=#s2s_blacklist;
                for i=1,count do
                        if s2s_blacklist[i] == host then
                                module:log ("error", "blacklisted host received %s", s2s_blacklist[i]);
                                s2smanager.destroy_session(session, "This host does not serve "..host);
                                return false;
                        end
                end
        elseif (host and is_whitelist_enabled == true)  then
                local count=#s2s_whitelist;
                local found=false;
                for i=1,count do
                        if s2s_whitelist[i] == host then
                                found=true;
                        end
                end
                if found == false then
                        module:log ("error", "host %s couldn't be found in whitelist", host);
                        s2smanager.destroy_session(session, "This host does not serve "..host);
                        return false;
                end
        end
  end
  return _make_connect(session, connect_host, connect_port);
end

local _stream_opened = s2smanager.streamopened;
function s2smanager.streamopened(session, attr)
        local host = attr.from and nameprep(attr.from);
        if not host then
                session.s2sValidation = false;
        else
                session.s2sValidation = true;
        end

        if (host and is_blacklist_enabled == true) then
                local count=#s2s_blacklist;
                for i=1,count do
                        if s2s_blacklist[i] == host then
                                module:log ("error", "blacklisted host received %s", s2s_blacklist[i]);
                                session:close({condition = "host-unknown", text = "This host does not serve " .. host});
                                return;
                        end
                end
        elseif (host and is_whitelist_enabled == true)  then
                local count=#s2s_whitelist;
                local found=false;
                for i=1,count do
                        if s2s_whitelist[i] == host then
                                found=true;
                        end
                end
                if found == false then
                        module:log ("error", "host %s couldn't be found in whitelist", host);
                        session:close({condition = "host-unknown", text = "This host does not serve " .. host});
                        return;
                end
        end
        _stream_opened(session, attr);
end


local function server_dialback_result_hook (event)
	local origin, stanza = event.origin, event.stanza;

	if origin.type == "s2sin" or origin.type == "s2sin_unauthed" then

		local host = stanza.attr.from;

		if (host and is_blacklist_enabled == true) then
			local count=#s2s_blacklist;
			for i=1,count do
 				if s2s_blacklist[i] == host then
					module:log ("error", "blacklisted host received %s", s2s_blacklist[i]);
      					origin:close({condition = "host-unknown", text = "This host does not serve " .. host});
					return true;
				end
			end
		elseif (host and is_whitelist_enabled == true)  then
			local count=#s2s_whitelist;
			local found=false;
			for i=1,count do
				if s2s_whitelist[i] == host then
					found=true;
				end
			end
			if found == false then
				module:log ("error", "host %s couldn't be found in whitelist", host);
      				origin:close({condition = "host-unknown", text = "This host does not serve " .. host});
				return true;
			end
		end
	
	end

	return nil;
end

local function handle_activated_host (host)
        if (hosts[host] and hosts[host].events) then
                hosts[host].events.add_handler("stanza/jabber:server:dialback:result", server_dialback_result_hook, 100);
                module:log ("debug", "adding hook for %s", host);
        end
end

local function handle_deactivated_host (host)
        if (hosts[host] and hosts[host].events) then
                hosts[host].events.remove_handler("stanza/jabber:server:dialback:result", server_dialback_result_hook);
                module:log ("debug", "removing hook for %s", host);
        end
end

prosody.events.add_handler("host-activated", handle_activated_host);
prosody.events.add_handler("component-activated", handle_activated_host);
prosody.events.add_handler("host-deactivated", handle_deactivated_host);
prosody.events.add_handler("component-deactivated", handle_deactivated_host);
prosody.events.add_handler("config-reloaded", reload_list);

for name, host in pairs(hosts) do
	if host and host.events then
		host.events.add_handler("stanza/jabber:server:dialback:result", server_dialback_result_hook, 100);
                module:log ("debug", "adding hook for %s", name);
	end
end