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/PUSH.hpp>
#include <rublon/method/SMS.hpp>
#include <rublon/method/SmsLink.hpp>
#include <rublon/method/YOTP.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, method::PUSH, method::SmsLink, method::YOTP > _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 if(_method == "push"){
            return MethodProxy{method::PUSH{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)}};
        }

        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)
     : 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 printAvailableMethods = [&]() -> tl::expected< int, MethodError > {
            int i{};
            for(const auto & method : _methods) {
                if(method == "totp") {
                    logMethodAvailable(method);
                    pam.print("%d: Mobile TOTP", i + 1);
                    methods_id[++i] = method;
                    continue;
                }
                
                       /// Needs changes in Core to work
                // if(method == "yotp") {
                //     logMethodAvailable(method);
                //     pam.print("%d: Yubikey", i + 1);
                //     methods_id[++i] = method;
                //     continue;
                // }
                
                if(method == "sms") {
                    logMethodAvailable(method);
                    pam.print("%d: SMS code", i + 1);
                    methods_id[++i] = method;
                    continue;
                }

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


                logMethodNotAvailable(method);
            }
            if(i == 0) {
                return tl::unexpected(MethodError(MethodError::ErrorClass::NoMethodAvailable));
            }
            return i;
        };

        const auto toMethodError = [&](conv::Error /*methodid*/) -> MethodError {
            pam.print("Input is not an number, please correct");
            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) {
                return tl::unexpected{MethodError(MethodError::BadMethod)};
            } 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 = [&]() -> tl::expected< PostMethod, MethodError > {
            printAvailableMethods();
            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 = [&](int methods_number) -> tl::expected< uint32_t, MethodError > {
            if(methods_number == 1) {
                pam.print("Only one method available");
                return 1;
            }
            return pam.scan(conv::to_uint32, "\nSelect method [1-%d]: ", methods_number).transform_error(toMethodError);
        };

        const auto handleErrors = [&](const MethodError & e) -> tl::expected< PostMethod, MethodError > {
            switch(e.errorClass) {
                case MethodError::BadMethod: // User provided a number but the number if not found
                    return askForMethodAgain();
                case MethodError::BadUserInput: // User provided id is invalid (NAN)
                    return askForMethodAgain();
                case MethodError::NoMethodAvailable:
                    /// TODO print enrolement
                default:
                    return tl::unexpected(e);
            }
        };

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

        return printAvailableMethods() //
          .and_then(askForMethod)
          .and_then(createMethod)
          .or_else(handleErrors)
          .map_error(toGenericError);
    }
};

} // namespace rublon