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

#pragma once

#include <set>
#include <tl/expected.hpp>

#include <variant>

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

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

template < class F >
struct return_type;

template < class R, class... A >
struct return_type< R (*)(A...) > {
    typedef R type;
};

template < typename T >
using return_type_t = typename return_type< T >::type;

namespace std {
inline ::rublon::Value::ConstValueIterator begin(const rublon::Value & __ils) noexcept {
    return __ils.Begin();
}
inline ::rublon::Value::ConstValueIterator end(const rublon::Value & __ils) noexcept {
    return __ils.End();
}

[[nodiscard]] inline std::size_t size(const rublon::Value & __cont) {
    return __cont.Size();
}

} // namespace std

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 {
        return std::visit(
          [&](const auto & method) {
              rublon::log(LogLevel::Info, "Using '%s' method", method.name);
              return method.fire(coreHandler, pam);
          },
          _impl);
    }

  private:
    std::variant< method::OTP, method::SMS > _impl;
};

class PostMethod : public rublon::AuthenticationStep< PostMethod > {
    using base_t = rublon::AuthenticationStep< PostMethod >;

    const char * uri = "/api/transaction/methodSSH";
    std::string _method;

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

        if(_method == "totp") {
            return MethodProxy{method::OTP{this->_systemToken, std::move(tid)}};
        } else if(_method == "sms") {
            return MethodProxy{method::SMS{this->_systemToken, std::move(tid)}};
        }

        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);
        body.AddMember("GDPRAccepted", "true", alloc);
        body.AddMember("tosAccepted", "true", alloc);
    }

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

    PostMethod(std::string systemToken, std::string tid, std::string method)
     : base_t(std::move(systemToken), std::move(tid)), _method{method} {}

    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 _tid;

    std::vector< std::string > _methods;

  public:
    template < typename Array_t >
    MethodSelect(std::string systemToken, std::string tid, const Array_t & methodsAvailableForUser)
     : _systemToken{std::move(systemToken)}, _tid{std::move(tid)} {
        _methods.reserve(std::size(methodsAvailableForUser));
        std::transform(
          std::begin(methodsAvailableForUser), std::end(methodsAvailableForUser), std::back_inserter(_methods), [](const auto & method) {
              return method.GetString();
          });
        rublon::log(LogLevel::Debug, "User has %d methods available", _methods.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_2k_HeapResource memoryResource;
        std::pmr::map< int, std::string > methods_id{&memoryResource};
        pam.print("select method: ");
        
        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 print_methods = [&]() {
            int i{};
            for(const auto & method : _methods) {
                if(method == "totp") {
                    logMethodAvailable(method);
                    pam.print("%d: Mobile TOTP", i + 1);
                    methods_id[++i] = "totp";
                    continue;
                }

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

        const auto toMethodError = [&](conv::Error /*methodid*/) -> Error {
            pam.print("Input is not an number, please correct");
            return MethodError{MethodError::BadMethod};
        };

        const auto createMethod = [&](std::uint32_t methodid) -> tl::expected< PostMethod, Error > {
            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) {
                return tl::unexpected{Error{Critical{}}}; /// TODO change to some more meaningfull error
            } else {
                log(LogLevel::Info, "User selected option %d{%s}", methodid, methods_id.at(methodid).c_str());
                return PostMethod{_systemToken, _tid, methods_id.at(methodid)};
            }
        };

        const auto askForMethodAgain = [&](const Error &) {
            print_methods();
            return pam
              .scan(conv::to_uint32, "\nSelect method [1-%d]: ", methods_id.size()) //
              .transform_error(toMethodError)
              .and_then(createMethod);
            /// TODO or_else(printErrorAndDenyAccess); ??
        };
        
        const auto askForMethod = [&]() -> tl::expected<uint32_t,conv::Error>{
            if(methods_id.size() == 1){
                pam.print("Only one method available");
                return 1;
            }
            return pam.scan(conv::to_uint32, "\nSelect method [1-%d]: ", methods_id.size());
        };
        
        print_methods();
        
        return askForMethod() //
          .transform_error(toMethodError)
          .and_then(createMethod)
          .or_else(askForMethodAgain);
    }
};

} // namespace rublon