mod_e2e_policy/mod_e2e_policy.lua

--  This software is free software: you can redistribute it and/or modify
--  it under the terms of the GNU General Public License as published by
--  the Free Software Foundation, version 3 of the License.
--
--  This software is distributed in the hope that it will be useful,
--  but WITHOUT ANY WARRANTY; without even the implied warranty of
--  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
--  GNU General Public License for more details.
--
--  You should have received a copy of the GNU General Public License
--  along with this software.  If not, see <https://www.gnu.org/licenses/>.

local st = require "util.stanza";
local set = require "util.set";
local jid_bare = require "util.jid".bare;
local it = require "util.iterators"
local string_render = require "util.interpolation".new("%b{}", st.xml_escape);

local host = module.host;

-- define allowed values for enum-ish config vars
local enum_scheme = set.new({ "omemo", "otr", "pgp", "pgp_legacy" })
local enum_graceperiod_scope = set.new({ "direct", "group" })
local enum_policy = set.new({ "optional", "required" })
local enum_warn_mechanism = set.new({ "message", "error" })

local store = module:open_store() -- equal to open_store(module.name, "keyvalue")

-- module options

local e2e_policy_accepted_schemes
local e2e_policy_graceperiod
local e2e_policy_graceperiod_scope
local e2e_policy_direct
local e2e_policy_group
local e2e_policy_warn_mechanism
local e2e_policy_whitelist

local e2e_policy_message_plain_optional_direct
local e2e_policy_message_plain_required_direct
local e2e_policy_message_unacceptable_optional_direct
local e2e_policy_message_unacceptable_required_direct

local e2e_policy_message_plain_optional_group
local e2e_policy_message_plain_required_group
local e2e_policy_message_unacceptable_optional_group
local e2e_policy_message_unacceptable_required_group

local e2e_policy_message_graceperiod


-- utilities

function set_len(s)
    -- return length of a set
    return it.count(s)
end

function ripairs(t)
    return it.reverse(ipairs(t))
end

local time_units = { "week", "day", "hour", "minute", "second" }
local time_distances = { second = 1, minute = 60 }

time_distances["hour"] = 60 * time_distances["minute"]
time_distances["day"] = 24 * time_distances["hour"]
time_distances["week"] = 7 * time_distances["day"]
-- No months and years, because they differ too much in length

function extract_time_units(seconds)

    local data = {}

    local leftover = seconds
    for i, unit in ipairs(time_units) do
        local distance = time_distances[unit]
        if leftover < distance then
            data[unit] = 0
        else -- we have at least 1 unit
            local cleanly_divisible = leftover
            leftover = leftover % distance -- number of seconds not divisible by unit
            cleanly_divisible = cleanly_divisible - leftover
            data[unit] = cleanly_divisible / distance
        end
    end

    return data
end

function pretty_time(n)

    if n == 0 then
        return "0 seconds"
    end

    local data = extract_time_units(n)

    local skip_head = set.new({})
    for i, unit in ipairs(time_units) do
        if data[unit] ~= 0 then
            break
        else
            skip_head:add(unit)
        end
    end

    local skip_tail = set.new({})
    for i, unit in ripairs(time_units) do
        if data[unit] ~= 0 then
            break
        else
            skip_tail:add(unit)
        end
    end

    local display_units = set.new(time_units)
    display_units:exclude(skip_head)
    display_units:exclude(skip_tail)
    display_units_num = set_len(display_units) -- set for some reason isn't a sequence, so #display_units doesn't work

    local output = ""
    local i = 0 -- keep track how many units we actually output
    --for unit in display_units do
    for _, unit in ipairs(time_units) do

        if display_units:contains(unit) then
            i = i + 1
            if i > 1 then
                if i == display_units_num then
                    output = output .. " and "
                else
                    output = output .. ", "
                end
            end

            output = output .. tostring(data[unit]) .. " " .. unit
            if data[unit] ~= 1 then -- ~= because *only* 1 implies no pluralization – i.e. 0 is pluralized
                output = output .. "s" -- pluralize unit
            end
        end
    end
    return output
end


-- module plumbing

function module.load()

    -- load config settings
    --
    -- this function is automatically called by Prosody's module loader
    -- after the module has been loaded.

    e2e_policy_accepted_schemes = module:get_option_set("e2e_policy_accepted_schemes", { "omemo", "pgp" }) -- possible values: omemo, otr, pgp and pgp_legacy – which E2EE schemes are accepted as secure
    e2e_policy_graceperiod = module:get_option_number("e2e_policy_graceperiod", 14 * 24 * 3600) -- grace period in seconds, this is used to escalate the per-user policy from "optional" to "required"
    e2e_policy_graceperiod_scope = module:get_option_set("e2e_policy_graceperiod_scope", { "direct" }) -- where to apply graceperiods

    e2e_policy_direct = module:get_option_string("e2e_policy_direct", "optional") -- possible values: none, optional and required
    e2e_policy_group = module:get_option_string("e2e_policy_group", "optional") -- possible values: none, optional and required
    e2e_policy_warn_mechanism = module:get_option_string("e2e_policy_warn_mechanism", "message") -- possible values: message and error - what mechanism to use to tell users about lack of E2ee
    e2e_policy_whitelist = module:get_option_set("e2e_policy_whitelist", {  }) -- make this module ignore messages sent to and from these JIDs or MUCs

    for i, scheme in ipairs(e2e_policy_accepted_schemes) do
        if not enum_scheme:contains(scheme) then
            error("Invalid value in e2e_policy_accepted_schemes: '" .. scheme .. "'. Valid values are: " .. tostring(enum_scheme))
        end
    end

    for i, scope in ipairs(e2e_policy_graceperiod_scope) do
        if not enum_graceperiod_scope:contains(scope) then
            error("Invalid value in e2e_policy_graceperiod_scope: '" .. scope .."'. Valid values are " .. tostring(enum_graceperiod_scope))
        end
    end

    if not enum_policy:contains(e2e_policy_direct) then
        error("Invalid value for e2e_policy_direct: '" .. e2e_policy_direct .. "'. Valid values are: " .. tostring(enum_policy))
    end
    if not enum_policy:contains(e2e_policy_group) then
        error("Invalid value for e2e_policy_group: '" .. e2e_policy_group .. "'. Valid values are: " .. tostring(enum_policy))
    end
    if not enum_warn_mechanism:contains(e2e_policy_warn_mechanism) then
        error("Invalid value for e2e_policy_warn_mechanism: '" .. e2e_policy_warn_mechanism .. "' Valid values are: " .. tostring(enum_warn_mechanism))
    end

    e2e_policy_message_plain_optional_direct = module:get_option_string("e2e_policy_message_plain_optional_direct", "Your message to {recipient} was not end-to-end encrypted. For security reasons, using one of the following E2EE schemes is *STRONGLY* recommended: {accepted_schemes} ")
    e2e_policy_message_plain_required_direct = module:get_option_string("e2e_policy_message_plain_required_direct", "Your message to {recipient} was not end-to-end encrypted. For security reasons, using one of the following E2EE schemes is *REQUIRED* for conversations on this server: {accepted_schemes} ")
    e2e_policy_message_unacceptable_optional_direct = module:get_option_string("e2e_policy_message_unacceptable_optional_direct", "Your message to {recipient} was end-to-end encrypted using the {scheme} scheme, but we recommend using one of the following instead: {accepted_schemes} ")
    e2e_policy_message_unacceptable_required_direct = module:get_option_string("e2e_policy_message_unacceptable_required_direct", "Your message to {recipient} was end-to-end encrypted using the {scheme} scheme, but this server *REQUIRES* one of these: {accepted_schemes} ")

    e2e_policy_message_plain_optional_group = module:get_option_string("e2e_policy_message_plain_optional_group", e2e_policy_message_plain_optional_direct)
    e2e_policy_message_plain_required_group = module:get_option_string("e2e_policy_message_plain_required_group", e2e_policy_message_plain_required_direct)
    e2e_policy_message_unacceptable_optional_group = module:get_option_string("e2e_policy_message_unacceptable_optional_group", e2e_policy_message_unacceptable_optional_direct)
    e2e_policy_message_unacceptable_required_group = module:get_option_string("e2e_policy_message_unacceptable_required_group", e2e_policy_message_unacceptable_required_direct)
    e2e_policy_message_graceperiod = module:get_option_string("e2e_policy_message_graceperiod", "You have {graceperiod_remaining} left before this will be enforced and messages without acceptable E2EE will be discarded.")
end
module:hook_global("config-reloaded", module.load) -- make sure config vars are updated when config is reloaded


-- actual module functionality

function send_warning(session, stanza, scheme, policy, graceperiod_remaining, message)

    local placeholder_values = {
        recipient = stanza.attr.to,
        scheme = scheme,
        accepted_schemes = tostring(e2e_policy_accepted_schemes)
    }

    if graceperiod_remaining then
        placeholder_values["graceperiod_remaining"] = pretty_time(graceperiod_remaining)
    end

    local rendered_message = string_render(message, placeholder_values)

    if e2e_policy_warn_mechanism == "message" then
        session.send(st.message({ from = host, to = stanza.attr.from, type = "headline" }, rendered_message))
    elseif e2e_policy_warn_mechanism == "error" then
        if policy == "required" then
            session.send(st.error_reply(stanza, "modify", "policy-violation", rendered_message))
        elseif policy == "optional" then
            session.send(st.error_reply(stanza, "continue", "policy-violation", rendered_message))
        end
    end
end

function get_e2e_scheme(stanza)

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

    -- check otr
    if body and body:sub(1,4) == "?OTR" then
        return "otr";
    end
    -- 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 "omemo";
    end
    -- check xep27 pgp https://xmpp.org/extensions/xep-0027.html
    if stanza:get_child("x", "jabber:x:encrypted") then
        return "pgp_legacy";
    end
    -- check xep373 pgp (OX) https://xmpp.org/extensions/xep-0373.html
    if stanza:get_child("openpgp", "urn:xmpp:openpgp:0") then
        return "pgp";
    end

    return "none"
end


function enforce_policy(event)

    -- prepend the servers' JID to the whitelist if we're using the "message"
    -- mechanism for warnings. this is done to avoid blocking our own warnings
    -- when using the "required" policy.
    local jid_whitelist
    if e2e_policy_warn_mechanism == "message" then
        jid_whitelist = set.new()
        jid_whitelist:add(module.host)
        for jid in e2e_policy_whitelist do
            jid_whitelist:add(jid)
        end
    else
        -- error stanzas won't trigger this function, thus
        -- prepending the server JID isn't necessary
        jid_whitelist = e2e_policy_whitelist 
    end

    -- check if JID is whitelisted
    if jid_whitelist:contains(jid_bare(event.stanza.attr.from)) or jid_whitelist:contains(jid_bare(event.stanza.attr.to)) then
        return nil;
    end
    local body = event.stanza:get_child_text("body");
    -- do not warn for status messages
    if not body or event.stanza.attr.type == "error" then
        return nil;
    end

    local scheme = get_e2e_scheme(event.stanza)
    if e2e_policy_accepted_schemes:contains(scheme) then -- message is encrypted with an accepted E2EE scheme
        return nil
    end

    -- no valid encryption found

    local now = os.time()
    local user = jid_bare(event.stanza.attr.from)
    local graceperiod_expired = false
    local graceperiod_remaining

    if not e2e_policy_graceperiod_scope:empty() then
        userdata, err = store:get(user)
        if userdata == nil then
            if err then
                module:log("error", "Can't get policy information from store for '%s'. Error was: %s", user, err)
            else -- no error, just no saved data yet, so create it
                userdata = {
                    first_message = os.time()
                }

                success, err = store:set(user, userdata)
                if not success then
                    if err then
                        module:log("error", "Failed to save policy information for '%s'. Error was: %s", user, err)
                    else
                        module:log("error", "Failed to save policy information for '%s'. No specific error was returned.", user)
                    end
                    -- TODO: return nil?
                end
            end
        elseif now >= userdata["first_message"] + e2e_policy_graceperiod then
            graceperiod_expired = true
        else
            graceperiod_remaining = userdata["first_message"] + e2e_policy_graceperiod - now
        end
    end

    local policy
    local message
  
    if scheme == "none" then

        if event.stanza.attr.type == "groupchat" then

            if e2e_policy_graceperiod_scope:contains("group") then
                if graceperiod_expired then
                    policy = "required"
                else
                    policy = "optional"
                end
            else
                policy = e2e_policy_group
            end

            if policy == "required" then
                message = e2e_policy_message_plain_required_group
            else
                message = e2e_policy_message_plain_optional_group
            end
        else

            if e2e_policy_graceperiod_scope:contains("direct") then
                if graceperiod_expired then
                    policy = "required"
                else
                    policy = "optional"
                end
            else
                policy = e2e_policy_direct
            end

            if policy == "required" then
                message = e2e_policy_message_plain_required_direct
            else
                message = e2e_policy_message_plain_optional_direct
            end
        end

        module:log("debug", "Plaintext message from %s!", event.stanza.attr.from)

    else -- unacceptable E2EE

        if event.stanza.attr.type == "groupchat" then
            policy = e2e_policy_group
            if policy == "required" then
                message = e2e_policy_message_unacceptable_required_group
            else
                message = e2e_policy_message_unacceptable_optional_group
            end
        else
            policy = e2e_policy_direct
            if policy == "required" then
                message = e2e_policy_message_unacceptable_required_direct
            else
                message = e2e_policy_message_unacceptable_optional_direct
            end
        end

        module:log("debug", "Message with unacceptable E2EE (%s) from %s!", scheme, event.stanza.attr.from)

    end

    if graceperiod_remaining then
        message = message .. "\n" .. e2e_policy_message_graceperiod
    end

    send_warning(event.origin, event.stanza, scheme, policy, graceperiod_remaining, message)

    if policy == "required" then
        return true -- stops further processing of this message, i.e. blocks it from being sent
    end

end

module:hook("pre-message/bare", enforce_policy, 300);
module:hook("pre-message/full", enforce_policy, 300);
module:hook("pre-message/host", enforce_policy, 300);