-- mod_auth_anon v0.1
-- Copyright (C) 2011 Matthieu Lalonde
--
--    All rights reserved.
--
--    Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
--
--        * Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
--        * Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
--        * Neither the name of Tobias Markmann nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
--
--    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

-- TODO: Minimum length username+pass
-- TODO: Add pretty nicknames to users (must be random and unique!)
-- TODO: Drop client that are idle for X time in order to help protect their privacy on Tor
-- DONE: On new signup, publish to the roster so no login cycle is needed
-- 	TODO: When a new user is pushed to the roster of online users, they don't seem to be authorized (or at least presence is not pushed)
-- TODO?: Change all Version IQ to INFOXMPP-APPROVED
-- TODO?: Add <sub xmlns='urn:xmpp:features:pre-approval'/> to stream negociation
-- TODO?: Change to a whitelisting system for IQ filtering in case there is more than bytestreams and jingle/content/transport that leak ips
-- TODO?: Move the datastore callbacks to the new mechanism?! (Need to figure out what that mechanism is but callbacks are deprecated apparently)
-- TODO?: Block locale from leaking (done in core/stanza_router.lua but needs to move to filters in here see: http://prosody-modules.googlecode.com/hg/mod_smacks/mod_smacks.lua )

local mod_auth_anon_changelog = {
	{0.3, [[
]]
	},
	{0.2, [[
- Added: Changelog and version
- Added: Generated user nicknames
- Fixed: Bytestreams filtering cleaned up and fixed for Psi
- Fixed: resource binding in Psi]]},
	{0.1, [[
- First versions before the changelog, a brief of current functionalities
  - Anonymous login with user ids based on hmac:sha256(salt, username+password),
    once logged in a user will receive their unique hashed username, resources are also anonymized.
    Users exist so long as they login every X days (default 1 day);
    after a user expires any of its stored data is deleted.
  - Automatic roster of all current users
  - Offline messages permitted only for encrypted messages (detects OTR, RSA and x:encrypted GPG)
  - Filtering of bytestream and jingle (at least libpurple) negotiation leaking a users IP
  - Filtering of a user's locale leaked by some clients (namely Tkabber)]]}
}

local hmac_sha256 = require "util.hmac".sha256
local hashes = require "util.hashes"
local uuid_generate = require "util.uuid".generate
local datamanager = require "util.datamanager"
local usermanager_create_user = require "core.usermanager".create_user
local json = require "util.json"
local st = require "util.stanza"

require "luasql.sqlite3"
local sqlite3 = luasql.sqlite3()
local connection

local resolve_relative_path = require "core.configmanager".resolve_relative_path
local sm_bind_resource = require "core.sessionmanager".bind_resource

local full_sessions = full_sessions

local data_path = (prosody and prosody.paths and prosody.paths.data) or "."

local xmlns_bind ='urn:ietf:params:xml:ns:xmpp-bind'

local module_host = module:get_host()

local params = module:get_option("anons")
if not params then -- Set defaults
	params	 = {}
end

params.expiry_time = (params.expiry_time or 1) -- Days
params.salt	= (params.salt or "iamsalt")

local expiry_time = params.expiry_time * 24 * 3600

local name_list_path = (prosody and prosody.paths and prosody.paths.config) or "."

local function open_list(file)
	local list = {}
	local fd = assert(io.open(file, "r"))
	for line in fd:lines() do
		table.insert(list, line)
	end
	
	return list
end

local animals = open_list(resolve_relative_path(name_list_path, "names/animals.list"))
local colors = open_list(resolve_relative_path(name_list_path, "names/colors.list"))
local locations = open_list(resolve_relative_path(name_list_path, "names/locations.list"))
local adjectives = open_list(resolve_relative_path(name_list_path, "names/adjectives.list"))

function round(num, idp)
  return tonumber(string.format("%." .. (idp or 0) .. "f", num))
end

local function tchelper(first, rest)
  return string.upper(first) .. string.lower(rest)
end

local function capitalize(str)
	return string.gsub(str, "(%a)([%w_']*)", tchelper)
end

local function trim(str)
	return string.gsub(str, "^%s*(.-)%s*$", "%1")
end

local function num_superscript(nbr)
	local str_nbr = tostring(nbr):sub(-1)
	local suffix = ""
	
	if string.len(tostring(nbr)) > 1 and tostring(nbr):sub(-2) == "13" then
		suffix = "th"
	elseif str_nbr == "1" then
		suffix = "st"
	elseif str_nbr == "2" then
		suffix = "nd"
	elseif str_nbr == "3" then
		suffix = "rd"
	else
		suffix = "th"
	end
	
	return suffix
end


local function get_hash_name(hash)
	local number = tonumber(string.format("0x%s", hash))
	local animal = math.fmod(number, table.getn(animals)) +1
	number = round(number / table.getn(animals))
	
	if number == 0 then
		return capitalize(trim(animals[animal]))
	end
	
	local color = math.fmod(number, table.getn(colors)) + 1
	number = round(number / table.getn(colors))
	
	if number == 0 then
		return capitalize(trim(colors[color])) .. " " .. capitalize(trim(animals[animal]))
	end
	
	local location = math.fmod(number, table.getn(locations)) + 1
	number = round(number / table.getn(locations))
	
	if number == 0 then
		return capitalize(trim(colors[color])) .. " " .. capitalize(trim(animals[animal])) .. " of " .. capitalize(trim(locations[location]))
	end
	
	local adjective = math.fmod(number, table.getn(adjectives)) + 1
	number = round(number / table.getn(adjectives))
	
	if number == 0 then
		return capitalize(trim(colors[color])) .. " " ..  capitalize(trim(animals[animal])) .. " `The " .. capitalize(trim(adjectives[adjective])) .. "` of " .. capitalize(trim(locations[location]))
	end
	
	local str_nbr = "The " .. number .. num_superscript(number)
	
	return capitalize(trim(colors[color])) .. 
			" " .. capitalize(trim(animals[animal])) .. 
			" " .. str_nbr .. 
			" `The " .. capitalize(trim(adjectives[adjective])) ..
			"` of " .. capitalize(trim(locations[location]))
end

-- URL Encode
local function encode(s)
	return s and (s:gsub("%W", function (c) return string.format("%%%02x", c:byte()); end))
end

-- Initiate SQL connections
local function connect()
	local database	= resolve_relative_path(data_path, (module_host and encode(module_host)) .. "/anons.sqlite")
	
	assert(datamanager.getpath("", module_host, "", "", true))
	
	local dbh = assert(sqlite3:connect(database))
	
	if not dbh then
		module:log("error", "Database connection failed: %s", tostring(err))
		return nil
	end
	
	module:log("debug", "Successfully connected to database")
	
	dbh:setautocommit(true);
	connection = dbh;
	
	return connection;
end

-- Drop SQL Connection
local function disconnect()
	connection:close();
	sqlite3:close();
end

-- Create anon table
local function create_table()
	local res = assert(connection:execute[[
		CREATE TABLE IF NOT EXISTS `anon_users` (TEXT, `user` TEXT, `name` TEXT, `last` TIMESTAMP, `stores` TEXT);
	]]);
	
	if not res then
		module:log("error", "Failed to create SQLite Anon Database");
	else
		module:log("info", "Initialized new %s database with anon_users", "SQLite");
		local res = connection:execute[[
			CREATE INDEX IF NOT EXISTS `anon_users_index` ON `anon_users` (`user`, `last`, `name`);
		]];
		
		if not res then
			module:log("warn", "Failed to create anon indexes, lookups may not be optimised");
		end
		
		return true;
	end
end

-- Lookup a user
local function is_anon_user(username, fields)
	if fields == nil then
		fields = "`user`"
	end
	
	local sql = string.format("SELECT %s FROM `anon_users` WHERE user = %q LIMIT 1;", fields, username)
	local res = connection:execute(sql)
	
	if not res then
		return nil
	else
		return res:fetch({}, "a");
	end
end

-- Log a user datastore so we can purge them later on
local function log_user_datastore(username, datastore)
	local user_result	= is_anon_user(username, "`user`, `stores`")
	
	if user_result ~= nil then
		local user_stores = json.decode(user_result.stores)
		
		if not user_stores[datastore] then
			user_stores[datastore]	= true
			local sql_str = "UPDATE `anon_users` SET stores  = '%s', last = %q WHERE user = %q;";
			local sql = string.format(sql_str, json.encode(user_stores), tostring(os.time()), username);
			
			if not connection:execute(sql) then
				module:log("error", "Failed to write stores %q status for %q: %s", datastore, username, sql);
				return nil
			else
				return true
			end
		else
			return true
		end
	else
		module:log("debug", "User %q doesn't exist, not storing datastore!", username)
		
		return nil
	end
end

-- Purges expired users by deleting any datastores they might have and their database entry 
local function purge_users()
	-- Fetch all the stale users
	sql = string.format("SELECT * FROM `anon_users` WHERE last < %s;", (os.time() - expiry_time));
	local res = connection:execute(sql);
	
	if not res then
		module:log("error", "Error fetching stale users: %s", sql);
		return nil;
	end
	
	local row = res:fetch ({}, "a");
	local errors = false;
	
	-- There are no expired users
	if not row then
		return true;
	end
	
	while row do	
		-- First make sure we update any expired logins that are still connected
		if prosody.bare_sessions[row.user .. "@" .. module_host] ~= nil then
			sql = string.format("UPDATE `anon_users` SET last = %s WHERE user = %q;", 
								tostring(os.time()), row.user);
			
			if not connection:execute(sql) then
				errors = true;
			end
		-- Otherwise we delete the user's datastores
		else
			local user_stores	= json.decode(row.stores);
			local str_stores		= ""
			
			if string.len(row.stores) > 0 then
				for t, v in pairs(user_stores) do
					-- Remove stores and list stores
					datamanager.store(row.user, module_host, tostring(t), nil)
					datamanager.list_store(row.user, module_host, tostring(t), nil)
					str_stores	= str_stores .. tostring(t) .. " "
				end
				
				module:log("debug", "Removed datastores %q for stale user %q", str_stores, row.user)
			end
		end
		
		row = res:fetch (row, "a")
	end
	
	-- Deleting any expired users left in the DB
	if not errors then
		sql = string.format("DELETE FROM `anon_users` WHERE last < %s;", (os.time() - expiry_time));
		res	= connection:execute(sql)
		
		if not res then
			module:log("error", "Failed to purge anon users, check the database!")
		else
			module:log("info", "Purged %s user(s) from anons!", res)
		end
		
		return true
	else
		module:log("error", "There was an error updating current expired users, purging failed!")
	end
end

-- Log an anon to the anon DB and clean the old ones.
local function log_user(username)
	local sql = string.format("SELECT `user` FROM `anon_users` WHERE user = %q;", username);
	local res = connection:execute(sql);
	
	-- The entry doesn't exist, let's create it
	if not res:fetch() then
		local user_longname = get_hash_name(username)
		sql = string.format("INSERT INTO `anon_users` (user, last, name, stores) VALUES (%q, %s, %q, \"{}\");", 
							username, os.time(), user_longname);
		module:log("debug", "Inserting new anon %s", username);
	-- The user already exists, update his last login time
	else
		sql = string.format("UPDATE `anon_users` SET last = %s WHERE user = %q;", os.time(), username);
		module:log("debug", "Updating current anon %s", username);
	end
	
	if not connection:execute(sql) then
		module:log("error", "Failed to write anon login to database: %s", sql)
	else -- Let's purge the old users here
		purge_users()
	end
end

-- Inject ALL THE anons into the roster
local function inject_roster_contacts(username, host, roster)
	local res = connection:execute("SELECT `user`, `name` FROM `anon_users`;");
	local row = res:fetch ({}, "a")
	
	while row do
		if username ~= row.user then
			local jid	= row.user .. "@" .. module_host;
			module:log("debug", "Adding user %s to roster %s.", jid, username);
			if not roster[jid] then roster[jid] = {}; end
			roster[jid].subscription = "both";
			roster[jid].published = true;
			roster[jid].name = row.name;
			roster[jid].persist = false;
			roster[jid].groups = { [module_host] = true };
		end
		
		row = res:fetch (row, "a")
	end
	
	roster[module_host]	= {}
	roster[module_host].subscription = "both"
	roster[module_host].published = true
	roster[module_host].name = module_host
	roster[module_host].persist = false
	roster[module_host].groups = { [module_host] = true }
end

local function push_roster_contact(username, full_jid)
	local user_longname = get_hash_name(username)
	local roster_attr = {jid = username.."@"..module_host, name = user_longname, 
							subscription = "both", approved = "true"}
	module:log("debug", "Sending new roster item for user: %s", username)
	
	for _, session in pairs(prosody.bare_sessions) do
		if _ ~= username.."@"..module_host then
			for __, current_session in pairs(session.sessions) do
				local roster_item = 
					st.iq({to = current_session.full_jid, type = "set"})
						:tag("query", {xmlns = "jabber:iq:roster", ver = os.time()})
							:tag("item", roster_attr)
								:tag("group"):text(module_host)
							:up()
						:up()
				
				core_route_stanza(module_host, roster_item)
			end
		end
	end
end

-- Catching the remove requests for data to block roster access
local function handle_datamanager(username, host, datastore, data)
	if host ~= module_host then
		return false;
	end
	
	if datastore == "roster" then
		-- Don't allow users to actually save rosters
		return false;
	else
		-- Register all datastores for a user so we can purge them later
		if data ~= nil then
			if log_user_datastore(username, datastore) ~= nil then
				return username, host, datastore, data;
			else
				return false;
			end
		else
			return username, host, datastore, data;
		end
	end
end

-- Anon Auth Provider, let ALL THE anons in!
module:add_item("auth-provider", {
	name = module.name:gsub("^auth_",""),
	test_password = function() return true end,
	user_exists = function(node, host) return (is_anon_user(node) ~= nil) end,
	get_sasl_handler = function () return {
		mechanisms = function() return { PLAIN = true, } end,
		plain = function(self, message)
			if string.len(message) < 18 then
				return "failure", "malformed-request", "Invalid username or password must be 18 characters or more.";
			else
				self.username = hmac_sha256(params.salt, message:match("%z(.*)$"), true);
				self.username = self.username:sub(-10);
			
				return "success"
			end
		end,
		select = function(self, mech) self.process = self[mech:lower()] return true end,
		clean_clone = function(self) self.username = nil; self.process = nil; return self end,
	} end,
})

-- Override the IQ Bind hook from saslauth in order to send the welcome message
module:hook("resource-bind", function(event)
	local session = event.session
	local bare_jid = session.username .. "@" .. session.host
	
	if is_anon_user(session.username) == nil then
		prosody.bare_sessions[bare_jid].is_new_user = true
	end
	
	log_user(session.username)
	
	local user_longname = is_anon_user(session.username, "`name`")
	user_longname = user_longname.name
	
	local welcome_stanza =
		st.message({ to = bare_jid, from = session.host })
			:tag("body")
				:text(string.format("Welcome in %s, your username is: %s@%s", 
									user_longname, session.username, session.host))
			:up()
	
	module:log("debug", "Welcomed user %s %q", session.full_jid, tostring(welcome_stanza))
	core_route_stanza(session.host, welcome_stanza)
	
	return nil -- Go on with the original hook if there is any
end, 9000) -- Priority over 9000, nothing should be over 9000, really!

module:hook("presence/host", function(event)
	local origin, stanza = event.origin, event.stanza;
	
	-- Catch the fist time a user comes online
	if stanza and stanza.attr and stanza.attr.type == "probe" and stanza.attr.to == module_host 
		and prosody.bare_sessions[origin.username.."@"..origin.host].is_new_user ~= nil then
		
		module:log("debug", "***** PRESENCE PUSHING NEW USER TO USERS ******")
		prosody.bare_sessions[origin.username.."@"..origin.host].is_new_user = nil
		push_roster_contact(origin.username, origin.full_jid)
	end
	
	return nil
end, 9000)

-- Override the IQ Bind hook from saslauth in order to anonymize resource
module:hook("iq/self/urn:ietf:params:xml:ns:xmpp-bind:bind", function(event)
	local origin, stanza = event.origin, event.stanza;
	
	if stanza.attr.type == "set" then
		local resource_tag = stanza.tags[1]:child_with_name("resource")
		local hashed_resource = tostring(hmac_sha256((resource_tag and #resource_tag.tags == 0 and resource_tag[1] or "wearelegions"), (resource_tag and #resource_tag.tags == 0 and resource_tag[1] or uuid_generate()), true)):sub(-16);
		
		if resource_tag == nil then
			stanza.tags[1]:tag("resource"):text(hashed_resource);
		else
			stanza.tags[1]:child_with_name("resource")[1] = hashed_resource
		end
	end
	
	return nil
end, 9000) -- Priority over 9000, nothing should be over 9000, really!

module:hook("iq/full", function(data)
	-- IQ to full JID recieved
	local origin, stanza = data.origin, data.stanza;
	local xmlns_bytestream = "http://jabber.org/protocol/bytestreams"
	
	-- Sanitize bytestreams
	if stanza:get_child("query", xmlns_bytestream) 
		and stanza:get_child("query", xmlns_bytestream):get_child("streamhost") then
		local sanitize_streamhost = function(tag)
			if tag.name == "streamhost" then
				local is_proxy = false
			
				if tag.attr and tag.attr.jid and string.find(tag.attr.jid, "proxy.") ~= nil then
					is_proxy	 = true
				elseif tag.attr and tag.attr.host and string.find(tag.attr.host, "proxy.") ~= nil then
					is_proxy	 = true
				elseif tag:get_child("proxy") ~= nil then
					is_proxy = true
				end
			
				if is_proxy == false then
					return nil
				end
			end
			
			return tag
		end
		
		stanza:get_child("query", xmlns_bytestream):maptags(sanitize_streamhost)
		
		local num_bytestreams = 0
		for child in stanza:get_child("query", xmlns_bytestream):childtags("bytestream") do
			num_bytestreams = num_bytestreams + 1
		end
		
		local num_tags = #stanza:get_child("query", xmlns_bytestream).tags
		if num_bytestreams < num_tags then
			local error_stanza = st.error_reply(stanza, "cancel", "service-unavailable")
			error_stanza.attr.to = stanza.attr.to
			error_stanza.attr.from = stanza.attr.from
			core_route_stanza(module_host, error_stanza);
			core_route_stanza(module_host, st.error_reply(stanza, "cancel", "service-unavailable")) --"error", nil, "Couldn't connect to any streamhosts"));
			return true;
		end
	-- Sanitize Jingle
	-- FIXME: This should use maptags as above
	elseif stanza and stanza.tags and stanza.tags[1] then
		local child	= stanza.tags[1]
		
		if child and child.name == "jingle" and child.attr.xmlns == "urn:xmpp:jingle:1" then
			for a, b in ipairs(stanza.tags[1]) do
				if b.name == "content" then
					for i, v in ipairs(stanza.tags[1].tags[a]) do
						if v.name == "transport" then
							local remove_ids = {}
							for ii, vv in ipairs(v) do
								if vv.name == "candidate" and vv.attr.type and vv.attr.type ~= "relay" then
									table.insert(remove_ids, ii)
								end
							end
						
							local ib = #remove_ids
							while ib > 0 do
								table.remove(stanza.tags[1].tags[a].tags[i], ib)
								ib = ib-1
							end
						end
					end
				end
			end
		end
	end
	
	return nil; -- Let the original hook go through
end, 9000);

-- Catch offline messages and allow only those that are encrypted
module:hook("message/offline/handle", function(event)
	local origin, stanza = event.origin, event.stanza;
	
	if not stanza:child_with_name("body") then
		return true; -- Only log actual messages
	end
	
	-- Let host and bare messages go through anyway
	if not (stanza.attr.from == "" or stanza.attr.from == module_host or stanza.attr.to == module_host) then
		if stanza and stanza.name and stanza.name == "message" then
			if stanza.tags and stanza.tags[2] then
				local is_encrypted	= false
				
				if stanza.tags[2].name == "body" then
					if string.find(stanza.tags[2]:get_text(), "?OTR:") == 1 then
						is_encrypted	 = true;
					end
					if string.find(stanza.tags[2]:get_text(), "*** Encrypted with the Gaim-Encryption plugin", 1, true) == 1 then
						is_encrypted	 = true;
					end
				end
				
				if stanza.tags[2].name == "x" and stanza.tags[2].attr.xmlns == "jabber:x:encrypted" then
					is_encrypted = true;
				end
				
				if is_encrypted == false then
					local error_stanza =
						st.message({ to = origin.full_jid, from = stanza.attr.to })
							:tag("body"):text("User offline; offline messages are only allowed when you use encryption");
					
					origin.send(	error_stanza);
					
					return true; -- Block the message from getting stored
				end
			end
		end
	end
	
	return nil; -- let the hook go along
end, 9000); -- Priority over 9000, nothing should be over 9000, really!

module:hook("message/host", function(event)
	local origin, stanza = event.origin, event.stanza;
	local message_body = stanza:child_with_name("body"):get_text()
	
	if string.find(message_body, "#!help") == 1 then
		local help_str = [[#!help		This message
#!version	Returns mod_auth_anon's current version
#!changelog	Returns the full changelog (one message per version)]]
		
		local reply_stanza = st.message({to = origin.full_jid, from = module_host, type = "chat"})
								:tag("body"):text(help_str)
		
		origin.send(reply_stanza)
		
		return true
	elseif string.find(message_body, "#!version") == 1 then
		local current_version = mod_auth_anon_changelog[1][1]
		local reply_stanza = st.message({to = origin.full_jid, from = module_host, type = "chat"})
								:tag("body"):text("Current version of Mod Auth Anon: " .. current_version	)
		
		origin.send(reply_stanza)
		
		return true
	elseif string.find(message_body, "#!changelog") == 1 then
		local changelog	= ""
		
		for i, v in ipairs(mod_auth_anon_changelog) do
			local version_body = "Version "..v[1].."\n"..v[2]
			local reply_stanza = st.message({to = origin.full_jid, from = module_host, type = "chat"})
									:tag("body"):text(version_body)
			
			origin.send(reply_stanza)
		end
		
		return true
--	elseif string.find(message_body, "#!hash") == 1 then
--		local file_path = string.sub(debug.getinfo(1).source, 2)
--		module:log("debug", file_path)
--		local file = assert(io.open(file_path, "r"))
--		local x = ""
--		for line in file:lines() do
--			x = x..line
--		end
--		file:close()
--		
--		local file_hash = hashes.md5(x, true)
--		
--		module:log("debug", file_path .. " " .. file_hash)
--		
--		return true
	else
		return nil
	end
end, 9000)

function module.load()
	assert (connect() ~= nil)
	
	assert (create_table() ~= nil)
	
	assert (purge_users() ~= nil)
	
	module:hook("roster-load", inject_roster_contacts);
	datamanager.add_callback(handle_datamanager);
end

function module.unload()
	datamanager.remove_callback(handle_datamanager);
	
	disconnect();
end