xbot/main.cpp

#include <boost/asio.hpp>
#include <boost/asio/ssl.hpp>

#include "linebuffer.hpp"
#include "ircmsg.hpp"
#include "settings.hpp"
#include "irc_thread.hpp"
#include "connection.hpp"

#include <algorithm>
#include <chrono>
#include <fstream>
#include <coroutine>
#include <iostream>
#include <limits>
#include <list>
#include <memory>
#include <string>
#include <string_view>
#include <tuple>
#include <utility>
#include <variant>
#include <vector>
#include <unordered_map>
#include <unordered_set>

using namespace std::chrono_literals;


struct ChatThread : public IrcThread
{
    auto priority() const -> priority_type override
    {
        return 100;
    }
    auto on_msg(ircmsg const& irc) -> std::pair<ThreadOutcome, EventOutcome> override
    {
        if (irc.command == "PRIVMSG" && 2 == irc.args.size())
        {
            std::cout << "Chat from " << irc.source << ": " << irc.args[1] << std::endl;
            return {ThreadOutcome::Continue, EventOutcome::Pass};
        }
        else
        {
            return {ThreadOutcome::Continue, EventOutcome::Pass};
        }
    }
};

struct UnhandledThread : public IrcThread
{
    auto priority() const -> priority_type override
    {
        return std::numeric_limits<IrcThread::priority_type>::max();
    }

    auto on_msg(ircmsg const& irc) -> std::pair<ThreadOutcome, EventOutcome> override
    {
        std::cout << "Unhandled message " << irc.command;
        for (auto const arg : irc.args)
        {
            std::cout << " " << arg;
        }
        std::cout << "\n";
        return {ThreadOutcome::Continue, EventOutcome::Pass};
    }
};

class PingThread : public IrcThread
{
    Connection * connection_;

public:
    PingThread(Connection * connection) noexcept : connection_{connection} {}

    auto priority() const -> priority_type override
    {
        return 0;
    }

    auto on_msg(ircmsg const& irc) -> std::pair<ThreadOutcome, EventOutcome> override
    {
        if (irc.command == "PING" && 1 == irc.args.size())
        {
            connection_->write("PONG", irc.args[0]);
            return {ThreadOutcome::Continue, EventOutcome::Consume};
        }
        else
        {
            return {};
        }
    }
};

struct RegistrationThread : IrcThread
{
    Connection * connection_;
    std::string password_;
    std::string username_;
    std::string realname_;
    std::string nickname_;

    std::unordered_map<std::string, std::string> caps;
    std::unordered_set<std::string> outstanding;

    enum class Stage
    {
        LsReply,
        AckReply,
    };

    Stage stage_;

    RegistrationThread(
        Connection * connection_,
        std::string password,
        std::string username,
        std::string realname,
        std::string nickname
    )
    : connection_{connection_}
    , password_{password}
    , username_{username}
    , realname_{realname}
    , nickname_{nickname}
    , stage_{Stage::LsReply}
    {}

    auto priority() const -> priority_type override { return 1; }
    auto on_connect() -> IrcThread::callback_result override
    {
        connection_->write("CAP", "LS", "302");
        connection_->write("PASS", password_);
        connection_->write("USER", username_, "*", "*", realname_);
        connection_->write("NICK", nickname_);
        return {};
    }

    auto send_req() -> IrcThread::callback_result
    {
        std::string request;
        char const* want[] = { "extended-join", "account-notify", "draft/chathistory", "batch", "soju.im/no-implicit-names", "chghost", "setname", "account-tag", "solanum.chat/oper", "solanum.chat/identify-msg", "solanum.chat/realhost", "server-time", "invite-notify", "extended-join" };
        for (auto cap : want)
        {
            if (caps.contains(cap))
            {
                request.append(cap);
                request.push_back(' ');
                outstanding.insert(cap);
            }
        }
        if (not outstanding.empty())
        {
            request.pop_back();
            connection_->write("CAP", "REQ", request);
            stage_ = Stage::AckReply;
            return {ThreadOutcome::Continue, EventOutcome::Consume};
        }
        else
        {
            connection_->write("CAP", "END");
            return {ThreadOutcome::Finish, EventOutcome::Consume};
        }
    }

    auto capack(ircmsg const& msg) -> IrcThread::callback_result
    {
        auto const n = msg.args.size();
        if ("CAP" == msg.command && n >= 2 && "*" == msg.args[0] && "ACK" == msg.args[1])
        {
            auto in = std::istringstream{std::string{msg.args[2]}};
            std::for_each(
                std::istream_iterator<std::string>{in},
                std::istream_iterator<std::string>{},
                [this](std::string x) {
                    outstanding.erase(x);
                }
            );
            if (outstanding.empty())
            {
                connection_->write("CAP","END");
                return {ThreadOutcome::Finish, EventOutcome::Consume};
            }
            else
            {
                return {ThreadOutcome::Continue, EventOutcome::Consume};
            }
        }
        else
        {
            return {};
        }
    }

    auto capls(ircmsg const& msg) -> IrcThread::callback_result
    {
        auto const n = msg.args.size();
        if ("CAP" == msg.command && n >= 2 && "*" == msg.args[0] && "LS" == msg.args[1])
        {
            std::string_view const* kvs;
            bool last;

            if (3 == n)
            {
                kvs = &msg.args[2];
                last = true;
            }
            else if (4 == n && "*" == msg.args[2])
            {
                kvs = &msg.args[3];
                last = false;
            }
            else
            {
                return {};
            }

            auto in = std::istringstream{std::string{*kvs}};

            std::for_each(
                std::istream_iterator<std::string>{in},
                std::istream_iterator<std::string>{},
                [this](std::string x) {
                    auto const eq = x.find('=');
                    if (eq == x.npos)
                    {
                        caps.emplace(x, std::string{});
                    }
                    else
                    {
                        caps.emplace(std::string{x, 0, eq}, std::string{x, eq+1, x.npos});
                    }
                }
            );

            if (last)
            {
                return send_req();
            }

            return {ThreadOutcome::Continue, EventOutcome::Consume};
        }
        else
        {
            return {};
        }

    }

    auto on_msg(ircmsg const& msg) -> IrcThread::callback_result override
    {
        switch (stage_)
        {
            case Stage::LsReply: return capls(msg);
            case Stage::AckReply: return capack(msg);
            default: return {};
        }
    }
};

auto start(boost::asio::io_context & io, Settings const& settings) -> void
{
    auto connection = std::make_shared<Connection>(io);

    connection->listen(std::make_unique<PingThread>(connection.get()));
    connection->listen(std::make_unique<RegistrationThread>(connection.get(), settings.password, settings.username, settings.realname, settings.nickname));
    connection->listen(std::make_unique<ChatThread>());
    connection->listen(std::make_unique<UnhandledThread>());

    boost::asio::co_spawn(
        io,
        connection->connect(io, settings),
        [&io, &settings](std::exception_ptr e)
        {
            auto timer = boost::asio::steady_timer{io};
            timer.expires_from_now(5s);
            timer.async_wait([&io, &settings](auto) {
                start(io, settings);
            });
        });
}

auto get_settings() -> Settings
{
    if (auto config_stream = std::ifstream {"config.toml"})
    {
        return Settings::from_stream(config_stream);
    }
    else
    {
        std::cerr << "Unable to open config.toml\n";
        std::exit(1);
    }
}

auto main() -> int
{
    auto const settings = get_settings();
    auto io = boost::asio::io_context{};
    start(io, settings);
    io.run();
}