rublon-ssh/PAM/ssh/include/rublon/method/method_select.hpp

#pragma once

#include <set>
#include <string_view>
#include <variant>
#include <vector>

#include <tl/expected.hpp>

#include <rublon/core_handler.hpp>
#include <rublon/error.hpp>
#include <rublon/memory.hpp>
#include <rublon/pam.hpp>
#include <rublon/pam_action.hpp>

#include <rublon/method/EMAIL.hpp>
#include <rublon/method/PUSH.hpp>
#include <rublon/method/SmsLink.hpp>

#include <rublon/method/OTP.hpp>
#include <rublon/method/SMS.hpp>
#include <rublon/method/YOTP.hpp>

extern std::string g_tid;

namespace rublon {

class MethodProxy {
  public:
    template < typename Method_t >
    MethodProxy(Method_t method) : _impl{std::move(method)} {}

    template < typename Handler_t, typename PamInfo_t = LinuxPam >
    tl::expected< AuthenticationStatus, Error > fire(const CoreHandlerInterface< Handler_t > & coreHandler, const PamInfo_t & pam) const {
        coreHandler.createWSConnection(g_tid);
        return std::visit(
          [&](const auto & method) {
              log(LogLevel::Info, "Using '%s' method", method._name);
              return method.verify(coreHandler, pam);
          },
          _impl);
    }

  private:
    std::variant< method::OTP, method::SMS, method::PUSH, method::EMAIL, method::SmsLink, method::YOTP > _impl;
};

class PostMethod : public AuthenticationStep {
    using base_t     = AuthenticationStep;
    const char * uri = "/api/transaction/methodSSH";
    std::string _method;
    int _prompts;
    bool _autopushPrompt;

    tl::expected< MethodProxy, Error > createMethod(const Document & coreResponse) const {
        const auto & rublonResponse = coreResponse["result"];
        std::string tid             = rublonResponse["tid"].GetString();
        std::string token           = rublonResponse.HasMember("token") ? rublonResponse["token"].GetString() : "";

        if(_method == "totp") {
            return MethodProxy{method::OTP{this->_systemToken, std::move(tid), _prompts}};
        } else if(_method == "sms") {
            return MethodProxy{method::SMS{this->_systemToken, std::move(tid), _prompts}};
        } else if(_method == "push") {
            return MethodProxy{method::PUSH{this->_systemToken, std::move(tid), _autopushPrompt}};
        } else if(_method == "email") {
            return MethodProxy{method::EMAIL{this->_systemToken, std::move(tid)}};
        } else if(_method == "smsLink") {
            return MethodProxy{method::SmsLink{this->_systemToken, std::move(tid)}};
        } else if(_method == "yotp") {
            return MethodProxy{method::YOTP{this->_systemToken, std::move(tid), std::move(token), _prompts}};
        } else
            return tl::unexpected{MethodError{MethodError::BadMethod}};
    }

    void addParams(Document & body) const {
        auto & alloc = body.GetAllocator();
        body.AddMember("method", Value{_method.c_str(), alloc}, alloc);
    }

  public:
    const char * _name = "Confirm Method";

    PostMethod(std::string systemToken, std::string tid, std::string method, int prompts, bool autopushPrompt)
     : base_t(std::move(systemToken), std::move(tid)), _method{method}, _prompts{prompts}, _autopushPrompt{autopushPrompt} {}

    template < typename Hander_t, typename PamInfo_t = LinuxPam >
    tl::expected< MethodProxy, Error > handle(const CoreHandlerInterface< Hander_t > & coreHandler) const {
        auto createMethod = [&](const auto & coreResponse) { return this->createMethod(coreResponse); };

        RapidJSONPMRStackAlloc< 2 * 1024 > alloc{};
        Document body{rapidjson::kObjectType, &alloc};

        this->addSystemToken(body);
        this->addTid(body);
        this->addParams(body);

        return coreHandler
          .request(alloc, uri, body) //
          .and_then(createMethod);
    }
};

class MethodSelect {
    std::string _systemToken;
    std::string _accessToken;
    std::string _tid;
    int _prompts;
    bool _autopushPrompt;

    std::vector< std::string > _methodsAvailable;

  public:
    template < typename Array_t >
    MethodSelect(std::string systemToken, std::string tid, const Array_t & methodsAvailableForUser, int prompts, bool autopushPrompt)
     : _systemToken{std::move(systemToken)}, _tid{std::move(tid)}, _prompts{prompts}, _autopushPrompt{autopushPrompt} {
        using namespace std::string_view_literals;
        memory::MonotonicStackResource< 2024 > stackResource;
        std::pmr::vector< std::string_view > _methods;
        _methodsAvailable.reserve(std::size(methodsAvailableForUser));

        std::pmr::set< std::string_view > methodsSupported{{"totp"sv, "email"sv, "yotp"sv, "sms"sv, "push"sv, "smsLink"sv}, &stackResource};

        transform_if(
          std::begin(methodsAvailableForUser),
          std::end(methodsAvailableForUser),
          std::back_inserter(_methodsAvailable),
          [&](const auto & method) { return method.GetString(); },
          [&](const auto & method) { return methodsSupported.find(method.GetString()) != methodsSupported.end(); });

        rublon::log(LogLevel::Debug, "User has %d methods available", _methodsAvailable.size());
    }

    template < typename Pam_t >
    tl::expected< PostMethod, Error > create(Pam_t & pam) const {
        rublon::log(LogLevel::Debug, "prompting user to select method");
        memory::StrictMonotonic_4k_HeapResource memoryResource;
        std::pmr::map< int, std::string > methods_id{&memoryResource};
        std::pmr::map< int, std::pmr::string > methods_names{&memoryResource};
        int prompts = _prompts;

        if(_methodsAvailable.size() == 0) {
            log(LogLevel::Warning, "None of provided methods are supported by the connector");
            return tl::unexpected(MethodError(MethodError::ErrorClass::NoMethodAvailable));
        }

        pam.print("Select the authentication method to verify your identity: ");

        auto logMethodAvailable = [](auto & method) { //
            rublon::log(LogLevel::Debug, "Method %s found", method.c_str());
        };
        auto logMethodNotAvailable = [](auto & method) {
            rublon::log(LogLevel::Debug, "Method %s found, but has no handler", method.c_str());
        };

        auto printAvailableMethods = [&]() -> tl::expected< int, MethodError > {
            int i{};
            for(const auto & method : _methodsAvailable) {
                if(method == "totp") {
                    logMethodAvailable(method);
                    pam.print("%d: Passcode", i + 1);
                    methods_id[++i]  = method;
                    methods_names[i] = "Passcode";
                    continue;
                }

                if(method == "email") {
                    logMethodAvailable(method);
                    pam.print("%d: Email Link", i + 1);
                    methods_id[++i]  = method;
                    methods_names[i] = "Email Link";
                    continue;
                }

                if(method == "yotp") {
                    logMethodAvailable(method);
                    pam.print("%d: YubiKey OTP Security Key", i + 1);
                    methods_id[++i]  = method;
                    methods_names[i] = "YubiKey OTP Security Key";
                    continue;
                }

                if(method == "sms") {
                    logMethodAvailable(method);
                    pam.print("%d: SMS Passcode", i + 1);
                    methods_id[++i]  = method;
                    methods_names[i] = "SMS Passcode";
                    continue;
                }

                if(method == "push") {
                    logMethodAvailable(method);
                    pam.print("%d: Mobile Push", i + 1);
                    methods_id[++i]  = method;
                    methods_names[i] = "Mobile Push";
                    continue;
                }

                if(method == "smsLink") {
                    logMethodAvailable(method);
                    pam.print("%d: SMS Link", i + 1);
                    methods_id[++i]  = method;
                    methods_names[i] = "SMS Link";
                    continue;
                }

                logMethodNotAvailable(method);
            }
            if(i == 0) {
                log(LogLevel::Warning, "None of provided methods are supported by the connector");
                return tl::unexpected(MethodError(MethodError::ErrorClass::NoMethodAvailable));
            }
            return i;
        };

        const auto toMethodError = [&](conv::Error /*methodid*/) -> MethodError {
            // NaN or out of range
            return MethodError{MethodError::BadUserInput};
        };

        const auto createMethod = [&](std::uint32_t methodid) -> tl::expected< PostMethod, MethodError > {
            auto hasMethod = methods_id.find(methodid) != methods_id.end();
            // pam.print("\t selected: %s", hasMethod ? methods_id.at(methodid).c_str() : "unknown option");
            if(!hasMethod) {
                log(LogLevel::Error, "User selected option %d, which is not correct", methodid);
                return tl::unexpected{MethodError(MethodError::BadMethod)};
            } else {
                log(LogLevel::Info, "User selected option %d{%s}", methodid, methods_names.at(methodid).c_str());
                return PostMethod{_systemToken, _tid, methods_id.at(methodid), _prompts, _autopushPrompt};
            }
        };

        const auto askForMethod = [&](int methods_number) -> tl::expected< uint32_t, MethodError > {
            if(methods_number == 1) {
                pam.print("Automatically selected the only available authentication method: %s", methods_id.at(1).c_str());
                return 1;
            }
            return pam.scan(conv::to_uint32, "\nSelect method [1-%d]: ", methods_number).transform_error(toMethodError);
        };

        auto reducePromptCount = [&](int selected_method) -> tl::expected< uint32_t, MethodError > {
            prompts--;
            return selected_method;
        };

        auto disableAnyNewMethodPrompts = [&]() { prompts = 0; };

        const auto printEnrolementInformation = [&]() {
            log(LogLevel::Warning, "Enrolement send to admin about new user: %s", pam.username());
            pam.print("Please check email");
        };

        const auto handleErrors = [&](const MethodError & e) -> tl::expected< PostMethod, MethodError > {
            switch(e.errorClass) {
                case MethodError::NoMethodAvailable:
                    disableAnyNewMethodPrompts();
                    printEnrolementInformation();
                    break;
                case MethodError::BadMethod: // User provided a number but the number if not found
                    pam.print("Input is not a valid method number. Enter a valid number");
                    break;
                case MethodError::BadUserInput: // User provided id is invalid (NAN)
                    pam.print("Input is not a number. Enter a valid number");
                    break;
            }
            return tl::unexpected(e);
        };

        const auto toGenericError = [](const MethodError & e) -> Error { return Error{e}; };

        return [&]() {
            while(true) {
                auto method = printAvailableMethods() //
                                .and_then(reducePromptCount)
                                .and_then(askForMethod)
                                .and_then(createMethod)
                                .or_else(handleErrors);

                if(not method && prompts > 0)
                    continue;
                return method;
            };
        }()
                 .map_error(toGenericError);
    }
};

} // namespace rublon