/*
 * Copyright (C) 2017-present ScyllaDB
 */

/*
 * SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.1
 */

#include "auth/standard_role_manager.hh"

#include <optional>
#include <stdexcept>
#include <unordered_set>
#include <vector>

#include <seastar/core/future-util.hh>
#include <seastar/core/on_internal_error.hh>
#include <seastar/core/format.hh>
#include <seastar/core/sleep.hh>
#include <seastar/core/sstring.hh>
#include <seastar/core/thread.hh>

#include "auth/common.hh"
#include "auth/config.hh"
#include "auth/role_manager.hh"
#include "auth/roles-metadata.hh"
#include "cql3/query_processor.hh"
#include "cql3/description.hh"
#include "cql3/untyped_result_set.hh"
#include "cql3/util.hh"
#include "db/consistency_level_type.hh"
#include "db/system_keyspace.hh"
#include "exceptions/exceptions.hh"
#include "utils/error_injection.hh"
#include "utils/log.hh"
#include <seastar/core/loop.hh>
#include <seastar/coroutine/maybe_yield.hh>
#include "service/raft/raft_group0_client.hh"
#include "service/migration_manager.hh"
#include "utils/managed_string.hh"

namespace auth {


static logging::logger log("standard_role_manager");

future<std::optional<standard_role_manager::record>> standard_role_manager::find_record(std::string_view role_name) {
    auto role = _cache.get(role_name);
    if (!role) {
        return make_ready_future<std::optional<record>>(std::nullopt);
    }
    return make_ready_future<std::optional<record>>(std::make_optional(record{
        .name = sstring(role_name),
        .is_superuser = role->is_superuser,
        .can_login = role->can_login,
        .member_of = role->member_of
    }));
}

future<standard_role_manager::record> standard_role_manager::require_record(std::string_view role_name) {
    return find_record(role_name).then([role_name](std::optional<record> mr) {
        if (!mr) {
            throw nonexistant_role(role_name);
        }

        return make_ready_future<record>(*mr);
   });
}

static bool has_can_login(const cql3::untyped_result_set_row& row) {
    return row.has("can_login") && !(boolean_type->deserialize(row.get_blob_unfragmented("can_login")).is_null());
}

standard_role_manager::standard_role_manager(cql3::query_processor& qp, ::service::raft_group0_client& g0, ::service::migration_manager& mm, cache& cache, const config& cfg)
    : _qp(qp)
    , _group0_client(g0)
    , _migration_manager(mm)
    , _cache(cache)
    , _superuser_name(cfg.auth_superuser_name)
    , _stopped(make_ready_future<>())
{}

std::string_view standard_role_manager::qualified_java_name() const noexcept {
    return "org.apache.cassandra.auth.CassandraRoleManager";
}

const resource_set& standard_role_manager::protected_resources() const {
    static const resource_set resources({
            make_data_resource(meta::legacy::AUTH_KS, meta::roles_table::name),
            make_data_resource(meta::legacy::AUTH_KS, ROLE_MEMBERS_CF)});

    return resources;
}

future<> standard_role_manager::maybe_create_default_role() {
    if (_superuser_name.empty()) {
        co_return;
    }
    auto has_superuser = [this] () -> future<bool> {
        const sstring query = seastar::format("SELECT * FROM {}.{} WHERE is_superuser = true ALLOW FILTERING", db::system_keyspace::NAME, meta::roles_table::name);
        auto results = co_await _qp.execute_internal(query, db::consistency_level::LOCAL_ONE,
                internal_distributed_query_state(), cql3::query_processor::cache_internal::yes);
        for (const auto& result : *results) {
            if (has_can_login(result)) {
                co_return true;
            }
        }
        co_return false;
    };
    if (co_await has_superuser()) {
        co_return;
    }
    // We don't want to start operation earlier to avoid quorum requirement in
    // a common case.
    ::service::group0_batch batch(
            co_await _group0_client.start_operation(_as, get_raft_timeout()));
    // Check again as the state may have changed before we took the guard (batch).
    if (co_await has_superuser()) {
        co_return;
    }
    // There is no superuser which has can_login field - create default role.
    // Note that we don't check if can_login is set to true.
    const sstring insert_query = seastar::format("INSERT INTO {}.{} ({}, is_superuser, can_login) VALUES (?, true, true)",
            db::system_keyspace::NAME,
            meta::roles_table::name,
            meta::roles_table::role_col_name);
    co_await collect_mutations(_qp, batch, insert_query, {sstring(_superuser_name)});
    co_await std::move(batch).commit(_group0_client, _as, get_raft_timeout());
    log.info("Created default superuser role '{}'.", _superuser_name);
}

future<> standard_role_manager::maybe_create_default_role_with_retries() {
    size_t retries = _migration_manager.get_concurrent_ddl_retries();
    while (true)  {
        try {
            co_return co_await maybe_create_default_role();
        } catch (const ::service::group0_concurrent_modification& ex) {
            log.warn("Failed to execute maybe_create_default_role due to guard conflict.{}.", retries ? " Retrying" : " Number of retries exceeded, giving up");
            if (retries--) {
                continue;
            }
            // Log error but don't crash the whole node startup sequence.
            log.error("Failed to create default superuser role due to guard conflict.");
            co_return;
        } catch (const ::service::raft_operation_timeout_error& ex) {
            log.error("Failed to create default superuser role due to exception: {}", ex.what());
            co_return;
        }
    }
}

future<> standard_role_manager::start() {
    return once_among_shards([this] () -> future<> {
        auto handler = [this] () -> future<> {
            co_await maybe_create_default_role_with_retries();
            if (!_superuser_created_promise.available()) {
                _superuser_created_promise.set_value();
            }
        };

        _stopped = auth::do_after_system_ready(_as, handler);
        co_return;
    });
}

future<> standard_role_manager::stop() {
    _as.request_abort();
    return _stopped.handle_exception_type([] (const sleep_aborted&) { }).handle_exception_type([](const abort_requested_exception&) {});;
}

future<> standard_role_manager::ensure_superuser_is_created() {
    SCYLLA_ASSERT(this_shard_id() == 0);
    return _superuser_created_promise.get_shared_future();
}

future<> standard_role_manager::create_or_replace(std::string_view role_name, const role_config& c, ::service::group0_batch& mc) {
    const sstring query = seastar::format("INSERT INTO {}.{} ({}, is_superuser, can_login) VALUES (?, ?, ?)",
            db::system_keyspace::NAME,
            meta::roles_table::name,
            meta::roles_table::role_col_name);
    co_await collect_mutations(_qp, mc,  query, {sstring(role_name), c.is_superuser, c.can_login});
}

future<>
standard_role_manager::create(std::string_view role_name, const role_config& c, ::service::group0_batch& mc) {
    return exists(role_name).then([this, role_name, &c, &mc](bool role_exists) {
        if (role_exists) {
            throw role_already_exists(role_name);
        }

        return create_or_replace(role_name, c, mc);
    });
}

future<>
standard_role_manager::alter(std::string_view role_name, const role_config_update& u, ::service::group0_batch& mc) {
    static const auto build_column_assignments = [](const role_config_update& u) -> sstring {
        std::vector<sstring> assignments;

        if (u.is_superuser) {
            assignments.push_back(sstring("is_superuser = ") + (*u.is_superuser ? "true" : "false"));
        }

        if (u.can_login) {
            assignments.push_back(sstring("can_login = ") + (*u.can_login ? "true" : "false"));
        }

        return fmt::to_string(fmt::join(assignments, ", "));
    };

    return require_record(role_name).then([this, role_name, &u, &mc](record) {
        if (!u.is_superuser && !u.can_login) {
            return make_ready_future<>();
        }
        const sstring query = seastar::format("UPDATE {}.{} SET {} WHERE {} = ?",
            db::system_keyspace::NAME,
            meta::roles_table::name,
            build_column_assignments(u),
            meta::roles_table::role_col_name);
        return collect_mutations(_qp, mc, std::move(query), {sstring(role_name)});
    });
}

future<> standard_role_manager::drop(std::string_view role_name, ::service::group0_batch& mc) {
    if (!co_await exists(role_name)) {
        throw nonexistant_role(role_name);
    }
    // First, revoke this role from all roles that are members of it.
    const auto revoke_from_members = [this, role_name, &mc] () -> future<> {
        const sstring query = seastar::format("SELECT member FROM {}.{} WHERE role = ?",
                db::system_keyspace::NAME,
                ROLE_MEMBERS_CF);
        const auto members = co_await _qp.execute_internal(
                query,
                db::consistency_level::LOCAL_ONE,
                internal_distributed_query_state(),
                {sstring(role_name)},
                cql3::query_processor::cache_internal::no);
        co_await parallel_for_each(
                members->begin(),
                members->end(),
                [this, role_name, &mc] (const cql3::untyped_result_set_row& member_row) -> future<> {
                    const sstring member = member_row.template get_as<sstring>("member");
                    co_await modify_membership(member, role_name, membership_change::remove, mc);
                }
        );
    };
    // In parallel, revoke all roles that this role is members of.
    const auto revoke_members_of = [this, grantee = role_name, &mc] () -> future<> {
        const role_set granted_roles = co_await query_granted(
                grantee,
                recursive_role_query::no);
        co_await parallel_for_each(
                granted_roles.begin(),
                granted_roles.end(),
                [this, grantee, &mc](const sstring& role_name) {
            return modify_membership(grantee, role_name, membership_change::remove, mc);
        });
    };
    // Delete all attributes for that role
    const auto remove_attributes_of = [this, role_name, &mc] () -> future<> {
        const sstring query = seastar::format("DELETE FROM {}.{} WHERE role = ?",
                db::system_keyspace::NAME,
                ROLE_ATTRIBUTES_CF);
        co_await collect_mutations(_qp, mc, query, {sstring(role_name)});
    };
    // Finally, delete the role itself.
    const auto delete_role = [this, role_name, &mc] () -> future<> {
        const sstring query = seastar::format("DELETE FROM {}.{} WHERE {} = ?",
                db::system_keyspace::NAME,
                meta::roles_table::name,
                meta::roles_table::role_col_name);

        co_await collect_mutations(_qp, mc, query, {sstring(role_name)});
    };

    co_await when_all_succeed(revoke_from_members, revoke_members_of, remove_attributes_of);
    co_await delete_role();
}

future<>
standard_role_manager::modify_membership(
        std::string_view grantee_name,
        std::string_view role_name,
        membership_change ch,
        ::service::group0_batch& mc) {
    const auto modify_roles = seastar::format(
            "UPDATE {}.{} SET member_of = member_of {} ? WHERE {} = ?",
            db::system_keyspace::NAME,
            meta::roles_table::name,
            (ch == membership_change::add ? '+' : '-'),
            meta::roles_table::role_col_name);
    co_await collect_mutations(_qp, mc, modify_roles,
            {role_set{sstring(role_name)}, sstring(grantee_name)});

    sstring modify_role_members;
    switch (ch) {
    case membership_change::add:
        modify_role_members = seastar::format("INSERT INTO {}.{} (role, member) VALUES (?, ?)",
                db::system_keyspace::NAME,
                ROLE_MEMBERS_CF);
        break;
    case membership_change::remove:
        modify_role_members = seastar::format("DELETE FROM {}.{} WHERE role = ? AND member = ?",
                db::system_keyspace::NAME,
                ROLE_MEMBERS_CF);
        break;
    default:
        on_internal_error(log, format("unknown membership_change value: {}", int(ch)));
    }
    co_await collect_mutations(_qp, mc, modify_role_members,
            {sstring(role_name), sstring(grantee_name)});
}

future<>
standard_role_manager::grant(std::string_view grantee_name, std::string_view role_name, ::service::group0_batch& mc) {
    const auto check_redundant = [this, role_name, grantee_name] {
        return query_granted(
                grantee_name,
                recursive_role_query::yes).then([role_name, grantee_name](role_set roles) {
            if (roles.contains(sstring(role_name))) {
                throw role_already_included(grantee_name, role_name);
            }

            return make_ready_future<>();
        });
    };

    const auto check_cycle = [this, role_name, grantee_name] {
        return query_granted(
                role_name,
                recursive_role_query::yes).then([role_name, grantee_name](role_set roles) {
            if (roles.contains(sstring(grantee_name))) {
                throw role_already_included(role_name, grantee_name);
            }

            return make_ready_future<>();
        });
    };

    return when_all_succeed(check_redundant(), check_cycle()).then_unpack([this, role_name, grantee_name, &mc] {
        return modify_membership(grantee_name, role_name, membership_change::add, mc);
    });
}

future<>
standard_role_manager::revoke(std::string_view revokee_name, std::string_view role_name, ::service::group0_batch& mc) {
    return exists(role_name).then([role_name](bool role_exists) {
        if (!role_exists) {
            throw nonexistant_role(sstring(role_name));
        }
    }).then([this, revokee_name, role_name, &mc] {
        return query_granted(
                revokee_name,
                recursive_role_query::no).then([revokee_name, role_name](role_set roles) {
            if (!roles.contains(sstring(role_name))) {
                throw revoke_ungranted_role(revokee_name, role_name);
            }

            return make_ready_future<>();
        }).then([this, revokee_name, role_name, &mc] {
            return modify_membership(revokee_name, role_name, membership_change::remove, mc);
        });
    });
}

future<> standard_role_manager::collect_roles(
        std::string_view grantee_name,
        bool recurse,
        role_set& roles) {
    return require_record(grantee_name).then([this, &roles, recurse](standard_role_manager::record r) {
        return do_with(std::move(r.member_of), [this, &roles, recurse](const role_set& memberships) {
            return do_for_each(memberships.begin(), memberships.end(), [this, &roles, recurse](const sstring& role_name) {
                roles.insert(role_name);

                if (recurse) {
                    return collect_roles(role_name, true, roles);
                }

                return make_ready_future<>();
            });
        });
    });
}

future<role_set> standard_role_manager::query_granted(std::string_view grantee_name, recursive_role_query m) {
    const bool recurse = (m == recursive_role_query::yes);

    return do_with(
            role_set{sstring(grantee_name)},
            [this, grantee_name, recurse](role_set& roles) {
        return collect_roles(grantee_name, recurse, roles).then([&roles] { return roles; });
    });
}

future<role_to_directly_granted_map> standard_role_manager::query_all_directly_granted(::service::query_state& qs) {
    role_to_directly_granted_map roles_map;
    _cache.for_each_role([&roles_map] (const cache::role_name_t& name, const cache::role_record& record) {
        for (const auto& granted_role : record.member_of) {
            roles_map.emplace(name, granted_role);
        }
    });
    co_return roles_map;
}

future<role_set> standard_role_manager::query_all(::service::query_state& qs) {
    role_set roles;
    roles.reserve(_cache.roles_count());
    _cache.for_each_role([&roles] (const cache::role_name_t& name, const cache::role_record&) {
        roles.insert(name);
    });
    co_return roles;
}

future<bool> standard_role_manager::exists(std::string_view role_name) {
    return find_record(role_name).then([](std::optional<record> mr) {
        return static_cast<bool>(mr);
    });
}

future<bool> standard_role_manager::is_superuser(std::string_view role_name) {
    return require_record(role_name).then([](record r) {
        return r.is_superuser;
    });
}

future<bool> standard_role_manager::can_login(std::string_view role_name) {
    return require_record(role_name).then([](record r) {
        return r.can_login;
    });
}

future<std::optional<sstring>> standard_role_manager::get_attribute(std::string_view role_name, std::string_view attribute_name, ::service::query_state& qs) {
    auto role = _cache.get(role_name);
    if (!role) {
        co_return std::nullopt;
    }
    auto it = role->attributes.find(attribute_name);
    if (it != role->attributes.end()) {
        co_return it->second;
    }
    co_return std::nullopt;
}

future<role_manager::attribute_vals> standard_role_manager::query_attribute_for_all(std::string_view attribute_name, ::service::query_state& qs) {
    attribute_vals result;
    _cache.for_each_role([&result, attribute_name] (const cache::role_name_t& name, const cache::role_record& record) {
        auto it = record.attributes.find(attribute_name);
        if (it != record.attributes.end()) {
            result.emplace(name, it->second);
        }
    });
    co_return result;
}

future<> standard_role_manager::set_attribute(std::string_view role_name, std::string_view attribute_name, std::string_view attribute_value, ::service::group0_batch& mc) {
    if (!co_await exists(role_name)) {
        throw auth::nonexistant_role(role_name);
    }
    const sstring query = seastar::format("INSERT INTO {}.{} (role, name, value)  VALUES (?, ?, ?)",
            db::system_keyspace::NAME,
            ROLE_ATTRIBUTES_CF);
    co_await collect_mutations(_qp, mc, query,
            {sstring(role_name), sstring(attribute_name), sstring(attribute_value)});
}

future<> standard_role_manager::remove_attribute(std::string_view role_name, std::string_view attribute_name, ::service::group0_batch& mc) {
    if (!co_await exists(role_name)) {
        throw auth::nonexistant_role(role_name);
    }
    const sstring query = seastar::format("DELETE FROM {}.{} WHERE role = ? AND name = ?",
            db::system_keyspace::NAME,
            ROLE_ATTRIBUTES_CF);
    co_await collect_mutations(_qp, mc, query,
            {sstring(role_name), sstring(attribute_name)});
}

future<std::vector<cql3::description>> standard_role_manager::describe_role_grants() {
    std::vector<cql3::description> result{};

    const auto grants = co_await query_all_directly_granted(internal_distributed_query_state());
    result.reserve(grants.size());

    for (const auto& [grantee_role, granted_role] : grants) {
        const auto formatted_grantee = cql3::util::maybe_quote(grantee_role);
        const auto formatted_granted = cql3::util::maybe_quote(granted_role);

        sstring create_statement = seastar::format("GRANT {} TO {};", formatted_granted, formatted_grantee);

        result.push_back(cql3::description {
            // Role grants do not belong to any keyspace.
            .keyspace = std::nullopt,
            .type = "grant_role",
            .name = granted_role,
            .create_statement = managed_string(create_statement)
        });

        co_await coroutine::maybe_yield();
    }

    std::ranges::sort(result, std::less<>{}, [] (const cql3::description& desc) {
        return std::make_tuple(std::ref(desc.name), std::ref(*desc.create_statement));
    });

    co_return result;
}

} // namespace auth