Alas/blcrack/updater/server.cpp

#include "server.hpp"

#define CRYPTOPP_ENABLE_NAMESPACE_WEAK 1

#include <fstream>
#include <filesystem>
#include <spdlog/spdlog.h>
#include <cryptopp/md5.h>
#include <cryptopp/hex.h>
#include <cryptopp/filters.h>
#include <json/json.h>

namespace fs = std::filesystem;

UpdateServer::UpdateServer() : m_assets_dir("assets") {
    update_files();

    Post("/files", [this](const httplib::Request& req, httplib::Response& res) {
        Json::Reader reader;
        Json::Value j;
        if (!reader.parse(req.body, j)) {
            SPDLOG_WARN("Invalid JSON request: {}", req.body);
            res.status = 400;
            return;
        }
        if (!j.isObject() || !j.isMember("arch") || !j.isMember("file")) {
            SPDLOG_WARN("Invalid JSON request: {}", req.body);
            res.status = 400;
            return;
        }

        auto arch = j["arch"].asString();
        auto file_name = j["file"].asString();
        SPDLOG_INFO("{} requested download {}/{}", req.remote_addr, arch, file_name);

        if (!is_file_exists(arch, file_name)) {
            SPDLOG_WARN("File {}/{} not exists", arch, file_name);
            res.status = 400;
            return;
        }

        auto& file_data = m_files[arch][file_name].data;
        res.set_content(file_data.data(), file_data.size(), "application/octet-stream");

        res.status = 200;
    });

    Post("/get_hash", [this](const httplib::Request& req, httplib::Response& res) {
        Json::Reader reader;
        Json::Value j;
        if (!reader.parse(req.body, j)) {
            SPDLOG_WARN("Invalid JSON request: {}", req.body);
            res.status = 400;
            return;
        }
        if (!j.isObject() || !j.isMember("arch") || !j.isMember("file")) {
            SPDLOG_WARN("Invalid JSON request: {}", req.body);
            res.status = 400;
            return;
        }

        auto arch = j["arch"].asString();
        auto file_name = j["file"].asString();
        SPDLOG_INFO("{} requested {}/{} hash", req.remote_addr, arch, file_name);

        if (!is_file_exists(arch, file_name)) {
            SPDLOG_WARN("File {}/{} not exists", arch, file_name);
            res.status = 400;
            return;
        }

        auto& hash = m_files[arch][file_name].hash;
        res.set_content(hash, "text/plain");

        res.status = 200;
    });

    Post("/reload", [this](const httplib::Request& req, httplib::Response& res) {
        SPDLOG_INFO("{} requested reload", req.remote_addr);
        if (req.body == m_admin_key) {
            update_files();
            res.status = 200;
            return;
        }
        SPDLOG_WARN("Invalid key: {}", req.body);
        res.status = 400;
    });

    Post("/update", [this](const httplib::Request& req, httplib::Response& res) {
        SPDLOG_INFO("{} requested update", req.remote_addr);
        const auto& key = req.params.equal_range("key").first->second;
        const auto& arch = req.params.equal_range("arch").first->second;
        const auto& file = req.params.equal_range("file").first->second;

        if (key == m_admin_key) {
            FileData data(req.body.begin(), req.body.end());
            auto md5 = hash_file(data);
            SPDLOG_INFO("Update {}/{} hash: {}", arch, file, md5);

            auto arch_path = m_assets_dir / arch;
            if (!exists(arch_path)) {
                create_directories(arch_path);
            }

            std::ofstream out(arch_path / file, std::ios::binary);
            out.write(data.data(), data.size());
            out.close();

            m_files[arch][file] = {
                .data = std::move(data),
                .hash = std::move(md5)
            };

            res.status = 200;
            return;
        }

        res.status = 400;
    });
}

UpdateServer& UpdateServer::start(int port, const std::string& admin_key) {
    m_admin_key = admin_key;
    m_server_thread = std::thread([this, port] {
        SPDLOG_INFO("Update server on port {}, reload key: {}", port, m_admin_key);
        listen("0.0.0.0", port);
    });
    return *this;
}

void UpdateServer::join() {
    m_server_thread.join();
}

UpdateServer& UpdateServer::Instance() {
    static UpdateServer instance;
    return instance;
}

void UpdateServer::update_files() {
    m_files.clear();
    if (!exists(m_assets_dir)) {
        return;
    }
    for (auto& arch : fs::directory_iterator(m_assets_dir)) {
        if (!fs::is_directory(arch)) {
            continue;
        }
        for (auto& file : fs::directory_iterator(arch)) {
            if (!fs::is_regular_file(file)) {
                continue;
            }
            auto arch_name = arch.path().filename().string();
            auto file_name = file.path().filename().string();

            auto size = file.file_size();
            FileData data(size);
            std::ifstream f(file.path().string(), std::ios::binary);
            f.read(data.data(), size);
            auto md5 = hash_file(data);
            SPDLOG_INFO("Load {}/{}, hash: {}", arch_name, file_name, md5);

            m_files[arch_name][file_name] = {
                        .data = std::move(data),
                        .hash = std::move(md5)
                    };
        }
    }
}

UpdateServer::FileHash UpdateServer::hash_file(const FileData& data) {
    FileHash hash;
    CryptoPP::Weak::MD5 md5;

    CryptoPP::StringSource s(
        reinterpret_cast<const unsigned char*>(data.data()), data.size(), true,
        new CryptoPP::HashFilter(md5,
                                 new CryptoPP::HexEncoder(
                                     new CryptoPP::StringSink(hash)
                                 )
        )
    );
    return hash;
}

bool UpdateServer::is_file_exists(const std::string& arch, const std::string& file) const {
    if (!m_files.contains(arch)) {
        return false;
    }
    auto& files = m_files.at(arch);
    return files.contains(file);
}