UsenetSearch/src/UsenetClient.cpp

/*
    Copyright© 2021 John Sennesael

    UsenetSearch is Free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    UsenetSearch is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with UsenetSearch.  If not, see <https://www.gnu.org/licenses/>.
*/

#include "usenetsearch/UsenetClient.h"

#include "usenetsearch/Application.h"
#include "usenetsearch/Except.h"
#include "usenetsearch/Logger.h"
#include "usenetsearch/StringUtils.h"

#include <codecvt>
#include <fstream>
#include <locale>
#include <memory>
#include <mutex>
#include <string>

namespace usenetsearch {

// UsenetClient class ----------------------------------------------------------

void UsenetClient::Authenticate(
    const std::wstring& user,
    const std::wstring& password)
{
    // Send user name
    Write(L"AUTHINFO USER " + user + L"\r\n");
    auto response = ReadLine();
    if (IsError(response))
    {
        Logger::Get().Fatal<UsenetClientException>(
            LOGID("UsenetClient"),
            "Error authenticating with NNTP server: "
            + response.message
        );
    }
    // Send password
    Write(L"AUTHINFO PASS " + password + L"\r\n");
    response = ReadLine();
    if (IsError(response))
    {
        Logger::Get().Fatal<UsenetClientException>(
            LOGID("UsenetClient"),
            "Error authenticating with NNTP server: "
            + response.message
        );
    }
}

void UsenetClient::Connect(
    const std::string& host,
    std::uint16_t port,
    bool useSSL)
{
    // Establish connection.
    m_useSSL = useSSL;
    try
    {
        m_tcp = std::make_unique<usenetsearch::TcpConnection>();
        m_tcp->Connect(host, port);
        if (useSSL)
        {
            m_ssl = std::make_unique<SSLConnection>(std::move(m_tcp));
            m_ssl->Connect();
        }
    }
    catch (const UsenetSearchException& e)
    {
        Logger::Get().Fatal<UsenetClientException>(
            LOGID("UsenetClient"),
            "Error while trying to connect to host: " + host + ":"
            + std::to_string(port) + " - " + e.what()
        );
    }
    // Read server banner.
    const auto serverHello = ReadLine();
    if (IsError(serverHello))
    {
        Logger::Get().Fatal<UsenetClientException>(
            LOGID("UsenetClient"),
            "Error received from NNTP server: "
            + serverHello.message
        );
    }
}

void UsenetClient::Group(const std::wstring& newsgroup)
{
    // Send user name
    Write(L"GROUP " + newsgroup + L"\r\n");
    auto response = ReadLine();
    if (IsError(response))
    {
        Logger::Get().Fatal<UsenetClientException>(
            LOGID("UsenetClient"),
            "Error changing group to " + StringFromWideString(newsgroup) + " : "
            + response.message
        );
    }
}

NntpHeader UsenetClient::Head(std::uint64_t articleID)
{
    Write(L"HEAD " + std::to_wstring(articleID) + L"\r\n");
    /* Typical response is a code 221 followed by the headers ending with a
       line containing a period by itself. */
    auto response = ReadLine();
    if (IsError(response))
    {
        Logger::Get().Fatal<UsenetClientException>(
            LOGID("UsenetClient"),
            "Error getting headers for article id "
            + std::to_string(articleID)
            + " : " + response.message
        );
    }
    const auto headerLinesStr = ReadUntil("\r\n.\r\n");
    // parse the headers.
    const auto lines = StringSplit(headerLinesStr, std::string{"\r\n"});
    NntpHeader result;
    result.articleID = articleID;
    for (const auto& line: lines)
    {
        const auto kvp = StringSplit(line, std::string{":"}, 2);
        if (kvp.size() != 2) continue; // Bad header line?
        const auto key = StringToLower(StringTrim(kvp[0]));
        const auto value = StringTrim(kvp[1]);
        if (key == "subject")
        {
            result.subject = value;
        }
    }
    return result;
}

bool UsenetClient::IsError(const NntpMessage& msg) const
{
    if (msg.code >= 400) return true;
    return false;
}

std::unique_ptr<std::vector<NntpListEntry>> UsenetClient::List()
{
    Write(L"LIST COUNTS\r\n");
    /* In response, we should get a 215 response followed by the list of news
       groups ending in a period on it's own line. */
    const auto response = ReadLine();
    if (IsError(response))
    {
        Logger::Get().Fatal<UsenetClientException>(
            LOGID("UsenetClient"),
            "Failed to fetch newsgroup list from server, "
            + std::string{"server responded with: "} 
            + response.message
        );
    }
    const auto listStr = ReadUntil("\r\n.\r\n");
    // parse the list.
    auto lines = StringSplit(listStr, std::string{"\r\n"});
    auto result = std::make_unique<std::vector<NntpListEntry>>();
    if (lines.empty()) return result;
    for (const auto& line: lines)
    {
        NntpListEntry entry;
        const auto fields = StringSplit(line, std::string{" "});
        if (fields.size() == 5)
        {
            entry.name = fields[0];
            entry.high = std::stoul(fields[1]);
            entry.low = std::stoul(fields[2]);
            entry.count = std::stoul(fields[3]);
            entry.status = fields[4];
            entry.id = 0; // incremented by db when saving.
            entry.lastIndexedArticle = NntpListEntry::NOT_INDEXED;
            if (Application::Get().GetFilter().ProcessNewsgroup(entry.name))
            {
                result->emplace_back(entry);
            }
        }
    }
    return result;
}

std::unique_ptr<std::vector<std::uint64_t>> UsenetClient::ListGroup(
    const std::wstring& newsGroup)
{
    auto result = std::make_unique<std::vector<std::uint64_t>>();
    if (!Application::Get().GetFilter().ProcessNewsgroup(
        StringFromWideString(newsGroup))
    )
    {
        return result;
    }
    Write(L"LISTGROUP " + newsGroup + L"\r\n");
    /* In response, we should get a 211 response followed by the list of
       article ID's ending in a period on it's own line. */
    const auto response = ReadLine();
    if (IsError(response))
    {
        Logger::Get().Fatal<UsenetClientException>(
            LOGID("UsenetClient"),
            "Failed to fetch newsgroup list from server, "
            + std::string{"server responded with: "} 
            + response.message
        );
    }
    const auto listStr = ReadUntil("\r\n.\r\n");
    // parse the list.
    auto lines = StringSplit(listStr, std::string{"\r\n"});
    if (lines.empty()) return result;
    for (const auto& line: lines)
    {
        result->emplace_back(stoul(StringTrim(line)));
    }
    return result;
}

void UsenetClient::ProcessHeaders(
    std::uint64_t startMessage,
    std::function<void(std::shared_ptr<NntpHeaders>)> processFn,
    std::uint64_t batchSize)
{
    Write(L"XHDR subject " + std::to_wstring(startMessage) + L"-\r\n");
    /* Typical response is a code 221 followed by the headers ending with a
       line containing a period by itself. */
    const auto response = ReadLine();
    if (IsError(response))
    {
        Logger::Get().Fatal<UsenetClientException>(
            LOGID("UsenetClient"),
            "Error getting headers: "
            + response.message
        );
    }
    NntpHeaders headers;
    std::mutex headersLock;
    while (true)
    {
        const auto line = StringTrim(ReadUntil("\r\n"));
        if (line == ".") break;
        // parse the headers.
        const auto kvp = StringSplit(line, std::string{" "}, 2);
        if (kvp.size() != 2) continue; // Bad header line?
        const std::uint64_t articleID = std::stoul(StringTrim(kvp[0]));
        const auto subject = StringTrim(kvp[1]);
        NntpHeader hdr;
        hdr.articleID = articleID;
        hdr.subject = subject;
        {
            std::lock_guard<std::mutex> lock(headersLock);
            headers.emplace_back(hdr);
        }
        if (headers.size() >= batchSize)
        {
            auto headersToPass = std::make_shared<NntpHeaders>(
                headers.begin(),
                headers.end()
            );
            processFn(std::move(headersToPass));
            headers.clear();
        }
    }
    // Left over headers.
    if (!headers.empty())
    {
        auto headersToPass = std::make_shared<NntpHeaders>(
            headers.begin(),
            headers.end()
        );
        processFn(std::move(headersToPass));
    }
}

NntpMessage UsenetClient::ReadLine()
{
    NntpMessage result{};
    std::string line;
    line = ReadUntil("\r\n");
    if (line.length() < 2)
    {
        Logger::Get().Fatal<UsenetClientException>(
            LOGID("UsenetClient"),
            "NNTP protocol error - invalid response from server: "
            + line
        );
    }
    std::string codeStr = line.substr(0, 3);
    result.code = std::stoi(codeStr);
    result.message = line.substr(4, line.length());
    return result;
}

std::string UsenetClient::ReadUntil(const std::string& deliminator)
{
    std::string result;
    if (m_useSSL)
    {
        result = m_ssl->ReadUntil(deliminator);
    }
    else
    {
        result = m_tcp->ReadUntil(deliminator);
    }
    result.erase(result.size() - deliminator.size(), deliminator.size());
    return result;
}

void UsenetClient::Write(const std::wstring& message)
{
    const std::string toSend = StringFromWideString(message);
    if (m_useSSL)
    {
        m_ssl->Write(toSend);
    }
    else
    {
        m_tcp->Write(toSend);
    }
}

} // namespace usenetsearch