-- XEP-0357: Push (aka: My mobile OS vendor won't let me have persistent TCP connections)

-- Copyright (C) 2015-2016 Kim Alvefur

-- Copyright (C) 2017-2019 Thilo Molitor

--

-- This file is MIT/X11 licensed.

local os_time = os.time;

local st = require"util.stanza";

local jid = require"util.jid";

local dataform = require"util.dataforms".new;

local hashes = require"util.hashes";

local random = require"util.random";

local cache = require"util.cache";

local xmlns_push = "urn:xmpp:push:0";

-- configuration

local include_body = module:get_option_boolean("push_notification_with_body", false);

local include_sender = module:get_option_boolean("push_notification_with_sender", false);

local max_push_errors = module:get_option_number("push_max_errors", 16);

local max_push_devices = module:get_option_number("push_max_devices", 5);

local dummy_body = module:get_option_string("push_notification_important_body", "New Message!");

local host_sessions = prosody.hosts[module.host].sessions;

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

local id2node = {};

local id2identifier = {};

-- For keeping state across reloads while caching reads

-- This uses util.cache for caching the most recent devices and removing all old devices when max_push_devices is reached

local push_store = (function()

	local store = module:open_store();

	local push_services = {};

	local api = {};

	--luacheck: ignore 212/self

	function api:get(user)

		if not push_services[user] then

			local loaded, err = store:get(user);

			if not loaded and err then

				module:log("warn", "Error reading push notification storage for user '%s': %s", user, tostring(err));

				push_services[user] = cache.new(max_push_devices):table();

				return push_services[user], false;

			end

			if loaded then

				push_services[user] = cache.new(max_push_devices):table();

				-- copy over plain table loaded from disk into our cache

				for k, v in pairs(loaded) do push_services[user][k] = v; end

			else

				push_services[user] = cache.new(max_push_devices):table();

			end

		end

		return push_services[user], true;

	end

	function api:flush_to_disk(user)

		local plain_table = {};

		for k, v in pairs(push_services[user]) do plain_table[k] = v; end

		local ok, err = store:set(user, plain_table);

		if not ok then

			module:log("error", "Error writing push notification storage for user '%s': %s", user, tostring(err));

			return false;

		end

		return true;

	end

	function api:set_identifier(user, push_identifier, data)

		local services = self:get(user);

		services[push_identifier] = data;

	end

	return api;

end)();

-- Forward declarations, as both functions need to reference each other

local handle_push_success, handle_push_error;

function handle_push_error(event)

	local stanza = event.stanza;

	local error_type, condition = stanza:get_error();

	local node = id2node[stanza.attr.id];

	local identifier = id2identifier[stanza.attr.id];

	if node == nil then return false; end		-- unknown stanza? Ignore for now!

	local from = stanza.attr.from;

	local user_push_services = push_store:get(node);

	local changed = false;

	for push_identifier, _ in pairs(user_push_services) do

		if push_identifier == identifier then

			if user_push_services[push_identifier] and user_push_services[push_identifier].jid == from and error_type ~= "wait" then

				push_errors[push_identifier] = push_errors[push_identifier] + 1;

				module:log("info", "Got error of type '%s' (%s) for identifier '%s': "

					.."error count for this identifier is now at %s", error_type, condition, push_identifier,

					tostring(push_errors[push_identifier]));

				if push_errors[push_identifier] >= max_push_errors then

					module:log("warn", "Disabling push notifications for identifier '%s'", push_identifier);

					-- remove push settings from sessions

					if host_sessions[node] then

						for _, session in pairs(host_sessions[node].sessions) do

							if session.push_identifier == push_identifier then

								session.push_identifier = nil;

								session.push_settings = nil;

								session.first_hibernated_push = nil;

							end

						end

					end

					-- save changed global config

					changed = true;

					user_push_services[push_identifier] = nil

					push_errors[push_identifier] = nil;

					-- unhook iq handlers for this identifier (if possible)

					module:unhook("iq-error/host/"..stanza.attr.id, handle_push_error);

					module:unhook("iq-result/host/"..stanza.attr.id, handle_push_success);

					id2node[stanza.attr.id] = nil;

					id2identifier[stanza.attr.id] = nil;

				end

			elseif user_push_services[push_identifier] and user_push_services[push_identifier].jid == from and error_type == "wait" then

				module:log("debug", "Got error of type '%s' (%s) for identifier '%s': "

					.."NOT increasing error count for this identifier", error_type, condition, push_identifier);

			end

		end

	end

	if changed then

		push_store:flush_to_disk(node);

	end

	return true;

end

function handle_push_success(event)

	local stanza = event.stanza;

	local node = id2node[stanza.attr.id];

	local identifier = id2identifier[stanza.attr.id];

	if node == nil then return false; end		-- unknown stanza? Ignore for now!

	local from = stanza.attr.from;

	local user_push_services = push_store:get(node);

	for push_identifier, _ in pairs(user_push_services) do

		if push_identifier == identifier then

			if user_push_services[push_identifier] and user_push_services[push_identifier].jid == from and push_errors[push_identifier] > 0 then

				push_errors[push_identifier] = 0;

				-- unhook iq handlers for this identifier (if possible)

				module:unhook("iq-error/host/"..stanza.attr.id, handle_push_error);

				module:unhook("iq-result/host/"..stanza.attr.id, handle_push_success);

				id2node[stanza.attr.id] = nil;

				id2identifier[stanza.attr.id] = nil;

				module:log("debug", "Push succeeded, error count for identifier '%s' is now at %s again",

					push_identifier, tostring(push_errors[push_identifier])

				);

			end

		end

	end

	return true;

end

-- http://xmpp.org/extensions/xep-0357.html#disco

local function account_dico_info(event)

	(event.reply or event.stanza):tag("feature", {var=xmlns_push}):up();

end

module:hook("account-disco-info", account_dico_info);

-- http://xmpp.org/extensions/xep-0357.html#enabling

local function push_enable(event)

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

	local enable = stanza.tags[1];

	origin.log("debug", "Attempting to enable push notifications");

	-- MUST contain a 'jid' attribute of the XMPP Push Service being enabled

	local push_jid = enable.attr.jid;

	-- SHOULD contain a 'node' attribute

	local push_node = enable.attr.node;

	-- CAN contain a 'include_payload' attribute

	local include_payload = enable.attr.include_payload;

	if not push_jid then

		origin.log("debug", "Push notification enable request missing the 'jid' field");

		origin.send(st.error_reply(stanza, "modify", "bad-request", "Missing jid"));

		return true;

	end

	if push_jid == stanza.attr.from then

		origin.log("debug", "Push notification enable request 'jid' field identical to our own");

		origin.send(st.error_reply(stanza, "modify", "bad-request", "JID must be different from ours"));

		return true;

	end

	local publish_options = enable:get_child("x", "jabber:x:data");

	if not publish_options then

		-- Could be intentional

		origin.log("debug", "No publish options in request");

	end

	local push_identifier = push_jid .. "<" .. (push_node or "");

	local push_service = {

		jid = push_jid;

		node = push_node;

		include_payload = include_payload;

		options = publish_options and st.preserialize(publish_options);

		timestamp = os_time();

	};

	local allow_registration = module:fire_event("cloud_notify/registration", {

		origin = origin, stanza = stanza, push_info = push_service;

	});

	if allow_registration == false then

		return true; -- Assume error reply already sent

	end

	push_store:set_identifier(origin.username, push_identifier, push_service);

	local ok = push_store:flush_to_disk(origin.username);

	if not ok then

		origin.send(st.error_reply(stanza, "wait", "internal-server-error"));

	else

		origin.push_identifier = push_identifier;

		origin.push_settings = push_service;

		origin.first_hibernated_push = nil;

		origin.log("info", "Push notifications enabled for %s (%s)", tostring(stanza.attr.from), tostring(origin.push_identifier));

		origin.send(st.reply(stanza));

	end

	return true;

end

module:hook("iq-set/self/"..xmlns_push..":enable", push_enable);

-- http://xmpp.org/extensions/xep-0357.html#disabling

local function push_disable(event)

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

	local push_jid = stanza.tags[1].attr.jid; -- MUST include a 'jid' attribute

	local push_node = stanza.tags[1].attr.node; -- A 'node' attribute MAY be included

	if not push_jid then

		origin.send(st.error_reply(stanza, "modify", "bad-request", "Missing jid"));

		return true;

	end

	local user_push_services = push_store:get(origin.username);

	for key, push_info in pairs(user_push_services) do

		if push_info.jid == push_jid and (not push_node or push_info.node == push_node) then

			origin.log("info", "Push notifications disabled (%s)", tostring(key));

			if origin.push_identifier == key then

				origin.push_identifier = nil;

				origin.push_settings = nil;

				origin.first_hibernated_push = nil;

			end

			user_push_services[key] = nil;

			push_errors[key] = nil;

			for stanza_id, identifier in pairs(id2identifier) do

				if identifier == key then

					module:unhook("iq-error/host/"..stanza_id, handle_push_error);

					module:unhook("iq-result/host/"..stanza_id, handle_push_success);

					id2node[stanza_id] = nil;

					id2identifier[stanza_id] = nil;

				end

			end

		end

	end

	local ok = push_store:flush_to_disk(origin.username);

	if not ok then

		origin.send(st.error_reply(stanza, "wait", "internal-server-error"));

	else

		origin.send(st.reply(stanza));

	end

	return true;

end

module:hook("iq-set/self/"..xmlns_push..":disable", push_disable);

-- is this push a high priority one (this is needed for ios apps not using voip pushes)

local function is_important(stanza)

	local st_name = stanza and stanza.name or nil;

	if not st_name then return false; end -- nonzas are never important here

	if st_name == "presence" then

		return false; -- same for presences

	elseif st_name == "message" then

		-- unpack carbon copied message stanzas

		local carbon = stanza:find("{urn:xmpp:carbons:2}/{urn:xmpp:forward:0}/{jabber:client}message");

		local stanza_direction = carbon and stanza:child_with_name("sent") and "out" or "in";

		if carbon then stanza = carbon; end

		local st_type = stanza.attr.type;

		-- headline message are always not important

		if st_type == "headline" then return false; end

		-- carbon copied outgoing messages are not important

		if carbon and stanza_direction == "out" then return false; end

		-- We can't check for body contents in encrypted messages, so let's treat them as important

		-- Some clients don't even set a body or an empty body for encrypted messages

		-- check omemo https://xmpp.org/extensions/inbox/omemo.html

		if stanza:get_child("encrypted", "eu.siacs.conversations.axolotl") or stanza:get_child("encrypted", "urn:xmpp:omemo:0") then return true; end

		-- check xep27 pgp https://xmpp.org/extensions/xep-0027.html

		if stanza:get_child("x", "jabber:x:encrypted") then return true; end

		-- check xep373 pgp (OX) https://xmpp.org/extensions/xep-0373.html

		if stanza:get_child("openpgp", "urn:xmpp:openpgp:0") then return true; end

		-- XEP-0353: Jingle Message Initiation (incoming call request)

		if stanza:get_child("propose", "urn:xmpp:jingle-message:0") then return true; end

		local body = stanza:get_child_text("body");

		-- groupchat subjects are not important here

		if st_type == "groupchat" and stanza:get_child_text("subject") then

			return false;

		end

		-- empty bodies are not important

		return body ~= nil and body ~= "";

	end

	return false;		-- this stanza wasn't one of the above cases --> it is not important, too

end

local push_form = dataform {

	{ name = "FORM_TYPE"; type = "hidden"; value = "urn:xmpp:push:summary"; };

	{ name = "message-count"; type = "text-single"; };

	{ name = "pending-subscription-count"; type = "text-single"; };

	{ name = "last-message-sender"; type = "jid-single"; };

	{ name = "last-message-body"; type = "text-single"; };

};

-- http://xmpp.org/extensions/xep-0357.html#publishing

local function handle_notify_request(stanza, node, user_push_services, log_push_decline)

	local pushes = 0;

	if not #user_push_services then return pushes end

	for push_identifier, push_info in pairs(user_push_services) do

		local send_push = true;		-- only send push to this node when not already done for this stanza or if no stanza is given at all

		if stanza then

			if not stanza._push_notify then stanza._push_notify = {}; end

			if stanza._push_notify[push_identifier] then

				if log_push_decline then

					module:log("debug", "Already sent push notification for %s@%s to %s (%s)", node, module.host, push_info.jid, tostring(push_info.node));

				end

				send_push = false;

			end

			stanza._push_notify[push_identifier] = true;

		end

		if send_push then

			-- construct push stanza

			local stanza_id = hashes.sha256(random.bytes(8), true);

			local push_notification_payload = st.stanza("notification", { xmlns = xmlns_push });

			local form_data = {

				-- hardcode to 1 because other numbers are just meaningless (the XEP does not specify *what exactly* to count)

				["message-count"] = "1";

			};

			if stanza and include_sender then

				form_data["last-message-sender"] = stanza.attr.from;

			end

			if stanza and include_body then

				form_data["last-message-body"] = stanza:get_child_text("body");

			elseif stanza and dummy_body and is_important(stanza) then

				form_data["last-message-body"] = tostring(dummy_body);

			end

			push_notification_payload:add_child(push_form:form(form_data));

			local push_publish = st.iq({ to = push_info.jid, from = module.host, type = "set", id = stanza_id })

				:tag("pubsub", { xmlns = "http://jabber.org/protocol/pubsub" })

					:tag("publish", { node = push_info.node })

						:tag("item")

							:add_child(push_notification_payload)

						:up()

					:up();

			if push_info.options then