smsub/SMSubServer.cpp

/*****************************************************************************
* smsub Server Manager Subprocess
* Copyright (C) 2020-2024 Syping
*
* Redistribution and use in source and binary forms, with or without modification,
* are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice,
* this list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* This software is provided as-is, no warranties are given to you, we are not
* responsible for anything with use of the software, you are self responsible.
*****************************************************************************/

#include <QCoreApplication>
#include <QTimer>
#include <QUuid>
#include "SMSubServer.h"
#include "smsub.h"
#ifdef BOOST_JSON
#include <boost/json.hpp>
#else
#include <QJsonDocument>
#include <QJsonObject>
#endif

SMSubServer::SMSubServer(SMSubServerSettings *serverSettings, const QString &socket) : serverSettings(serverSettings)
{
    QLocalServer *localServer = new QLocalServer(this);
    localServer->setSocketOptions(QLocalServer::UserAccessOption | QLocalServer::GroupAccessOption);
    localServer->listen(socket);

    QObject::connect(localServer, &QLocalServer::newConnection, this, &SMSubServer::newConnection);

    type = ServerType::Local;
    server = localServer;
    status = false;
    startTime = QDateTime::currentDateTimeUtc().toSecsSinceEpoch();
    stopTime = startTime;
}

SMSubServer::SMSubServer(SMSubServerSettings *serverSettings, const QString &serverName, const quint16 &port) : serverSettings(serverSettings)
{
    QWebSocketServer *webSocketServer = new QWebSocketServer(serverName, QWebSocketServer::NonSecureMode, this);
    webSocketServer->setSupportedSubprotocols(QStringList() << "smsub" << "smsub_json");
    webSocketServer->listen(QHostAddress::LocalHost, port);

    QObject::connect(webSocketServer, &QWebSocketServer::newConnection, this, &SMSubServer::newConnection);

    type = ServerType::WebSocket;
    server = webSocketServer;
    status = false;
    startTime = QDateTime::currentDateTimeUtc().toSecsSinceEpoch();
    stopTime = startTime;
}

bool SMSubServer::isListening()
{
    if (type == ServerType::Local) {
        return static_cast<QLocalServer*>(server)->isListening();
    }
    else if (type == ServerType::WebSocket) {
        return static_cast<QWebSocketServer*>(server)->isListening();
    }
    return false;
}

void SMSubServer::newConnection()
{
    QObject *socket;
    if (Q_LIKELY(type == ServerType::Local)) {
        QLocalSocket *localSocket = static_cast<QLocalServer*>(server)->nextPendingConnection();
        QObject::connect(localSocket, &QLocalSocket::readyRead, this, &SMSubServer::lsReadyRead);
        QObject::connect(localSocket, &QLocalSocket::disconnected, this, &SMSubServer::deleteSocket);
        localSocket->write(QString("SMSub Version %1\n").arg(QCoreApplication::applicationVersion()).toUtf8());
        socket = localSocket;
        QTextStream(stderr) << "LocalSocket connected!" << smsub_endl;
    }
    else if (type == ServerType::WebSocket) {
        QWebSocket *webSocket = static_cast<QWebSocketServer*>(server)->nextPendingConnection();
        QObject::connect(webSocket, &QWebSocket::binaryMessageReceived, this, &SMSubServer::wsMessageReceived);
        QObject::connect(webSocket, &QWebSocket::disconnected, this, &SMSubServer::deleteSocket);
        webSocket->sendBinaryMessage(QString("SMSub Version %1\n").arg(QCoreApplication::applicationVersion()).toUtf8());
        socket = webSocket;
        QTextStream(stderr) << QString("WebSocket %1:%2 connected!").arg(webSocket->peerAddress().toString(), QString::number(webSocket->peerPort())) << smsub_endl;

        if (webSocket->subprotocol() == "smsub_json")
            socket->setProperty("ReceiveJson", true);
    }
    else {
        // Just for being sure
        return;
    }

    // Set authentication state
    if (serverSettings->isLocal) {
        socket->setProperty("Authenticated", true);
        sockets << socket;
    }
    else {
        socket->setProperty("Authenticated", false);
    }
}

bool SMSubServer::messageReceived(QObject *socket, const QByteArray &message)
{
    // Only allow commands being sent if authenticated
    const bool isAuthenticated = socket->property("Authenticated").toBool();
    if (isAuthenticated) {
        if (message == "+dbg") {
            socket->setProperty("ReceiveDbgMsg", true);
            sendMessage(socket, "Debug messages enabled!\n");
        }
        else if (message == "-dbg") {
            socket->setProperty("ReceiveDbgMsg", false);
            sendMessage(socket, "Debug messages disabled!\n");
        }
        else if (message == "+json") {
            socket->setProperty("ReceiveJson", true);
            debugOutput(socket, "JSON output enabled!");
        }
        else if (message == "-json") {
            socket->setProperty("ReceiveJson", false);
            debugOutput(socket, "JSON output disabled!");
        }
        else if (message == "+log") {
            socket->setProperty("ReceiveLog", true);
            debugOutput(socket, "Log output enabled!");
        }
        else if (message == "-log") {
            socket->setProperty("ReceiveLog", false);
            debugOutput(socket, "Log output disabled!");
        }
        else if (message == "+status") {
            socket->setProperty("ReceiveStatus", true);
            debugOutput(socket, "Status updates enabled!");
        }
        else if (message == "-status") {
            socket->setProperty("ReceiveStatus", false);
            debugOutput(socket, "Status updates disabled!");
        }
        else if (message == "+reg") {
            if (serverSettings->canRegister) {
                QByteArray authUuid = QUuid::createUuid().toByteArray(QUuid::Id128) +
                                      QUuid::createUuid().toByteArray(QUuid::Id128);
                authUuid = QByteArray::fromHex(authUuid).toBase64(QByteArray::OmitTrailingEquals);
                emit tokenRegistered(QString::fromUtf8(authUuid));
                sendMessage(socket, "Token: " + authUuid + '\n');
            }
            else {
                sendMessage(socket, "Permission denied!\n");
            }
        }
        else if (message == "status") {
            if (isVariantTrue(socket->property("ReceiveJson"))) {
#ifdef BOOST_JSON
                boost::json::object object;
                object["type"] = "status";
                object["status"] = status ? "on" : "off";
                object["start"] = startTime;
                object["stop"] = stopTime;
                const std::string json = boost::json::serialize(object) + '\n';
                sendMessage(socket, QByteArray::fromRawData(json.data(), json.size()));
#else
                QJsonObject object;
                object["type"] = "status";
                object["status"] = status ? "on" : "off";
                object["start"] = startTime;
                object["stop"] = stopTime;
                sendMessage(socket, QJsonDocument(object).toJson(QJsonDocument::Compact) + '\n');
#endif
            }
            else {
                sendMessage(socket, QString("Status: %1\nLast Start: %2\nLast Stop: %3\n").arg(status ? "on" : "off", QString::number(startTime), QString::number(stopTime)).toUtf8());
            }
        }
        else if (message == "start") {
            emit startRequested();
            debugOutput(socket, "Starting server!");
        }
        else if (message == "stop") {
            emit stopRequested();
            debugOutput(socket, "Stopping server!");
        }
        else if (message == "kill") {
            emit killRequested();
            debugOutput(socket, "Killing server!");
        }
        else if (message == "quit") {
            QTimer::singleShot(0, qApp, &QCoreApplication::quit);
            debugOutput(socket, "Qutting smsub!");
        }
        else if (message == "wl") {
            emit inputWritten("\n");
            debugOutput(socket, "Writing line!");
        }
        else if (message.startsWith("wl ")) {
            const QByteArray writeData = message.mid(3);
            emit inputWritten(writeData + '\n');
            debugOutput(socket, "Writing line: " + writeData);
        }
        else if (message.startsWith("w ")) {
            const QByteArray writeData = message.mid(2);
            emit inputWritten(writeData);
            debugOutput(socket, "Writing: " + writeData);
        }
    }
    else {
        // Authenticate when token is valid, otherwise disconnect
        if (tokens.contains(QString::fromUtf8(message))) {
            // Set client as authenticated and add it to vector
            socket->setProperty("Authenticated", true);
            sendMessage(socket, "Login successful!\n");
            sockets << socket;
        }
        else {
            // Stop receiving data and disconnect socket
            if (type == ServerType::Local) {
                QLocalSocket *localSocket = static_cast<QLocalSocket*>(socket);
                QObject::disconnect(localSocket, &QLocalSocket::readyRead, this, &SMSubServer::lsReadyRead);
                localSocket->write("Incorrect token!\n");
                localSocket->disconnectFromServer();
                return false;
            }
            else if (type == ServerType::WebSocket) {
                QWebSocket *webSocket = static_cast<QWebSocket*>(socket);
                QObject::disconnect(webSocket, &QWebSocket::binaryMessageReceived, this, &SMSubServer::wsMessageReceived);
                webSocket->sendBinaryMessage("Incorrect token!\n");
                webSocket->close(QWebSocketProtocol::CloseCodeNormal);
                return false;
            }
        }
    }
    return true;
}

bool SMSubServer::isVariantTrue(const QVariant &variant)
{
#if QT_VERSION >= 0x060000
    if (variant.typeId() == QMetaType::Bool) {
#else
    if (variant.type() == QVariant::Bool) {
#endif
        return variant.toBool();
    }
    return false;
}

void SMSubServer::wsMessageReceived(const QByteArray &message)
{
    QWebSocket *socket = static_cast<QWebSocket*>(sender());
    messageReceived(socket, message.trimmed());
}

void SMSubServer::lsReadyRead()
{
    QLocalSocket *socket = static_cast<QLocalSocket*>(sender());
#ifdef SMSUB_IODEBUG
    QTextStream(stderr) << "LocalSocket I/O RR!" << smsub_endl;
#endif
    while (socket->canReadLine()) {
#ifdef SMSUB_IODEBUG
        QTextStream(stderr) << "LocalSocket I/O WL!" << smsub_endl;
#endif
        const QByteArray message = socket->readLine().trimmed();
        if (!messageReceived(socket, message))
            return;
    }
}

void SMSubServer::deleteSocket()
{
    // Delete socket and remove from index
    QObject *socket = sender();
    sockets.removeAll(socket);
    socket->deleteLater();
}

void SMSubServer::debugOutput(QObject *socket, const QByteArray &message)
{
    // Only send debug messages when the client opted-in
    if (isVariantTrue(socket->property("ReceiveDbgMsg"))) {
        if (isVariantTrue(socket->property("ReceiveJson"))) {
#ifdef BOOST_JSON
            boost::json::object object;
            object["type"] = "debug";
            object["message"] = message.toBase64(QByteArray::OmitTrailingEquals).constData();
            const std::string json = boost::json::serialize(object) + '\n';
            sendMessage(socket, QByteArray::fromRawData(json.data(), json.size()));
#else
            QJsonObject object;
            object["type"] = "debug";
            object["message"] = QString::fromUtf8(message.toBase64(QByteArray::OmitTrailingEquals));
            sendMessage(socket, QJsonDocument(object).toJson(QJsonDocument::Compact) + '\n');
#endif
        }
        else {
            sendMessage(socket, message + '\n');
        }
    }
}

void SMSubServer::writeOutput(const QByteArray &output)
{
    // Read process output when client opted-in for log
    QByteArray json_output;
    for (auto it = sockets.constBegin(); it != sockets.constEnd(); it++) {
        if (isVariantTrue((*it)->property("ReceiveLog"))) {
            if (isVariantTrue((*it)->property("ReceiveJson"))) {
                if (json_output.isEmpty()) {
#ifdef BOOST_JSON
                    boost::json::object object;
                    object["type"] = "log";
                    object["message"] = output.toBase64(QByteArray::OmitTrailingEquals).constData();
                    json_output = QByteArray::fromStdString(boost::json::serialize(object) + '\n');
#else
                    QJsonObject object;
                    object["type"] = "log";
                    object["message"] = QString::fromUtf8(output.toBase64(QByteArray::OmitTrailingEquals));
                    json_output = QJsonDocument(object).toJson(QJsonDocument::Compact) + '\n';
#endif
                }
                sendMessage(*it, json_output);
            }
            else {
                sendMessage(*it, output);
            }
        }
    }
}

void SMSubServer::sendMessage(QObject *socket, const QByteArray &message)
{
    if (type == ServerType::Local) {
        QLocalSocket *localSocket = static_cast<QLocalSocket*>(socket);
        localSocket->write(message);
    }
    else if (type == ServerType::WebSocket) {
        QWebSocket *webSocket = static_cast<QWebSocket*>(socket);
        webSocket->sendBinaryMessage(message);
    }
}

void SMSubServer::registerToken(const QString &token)
{
    // Register temporary token for a secure remote connection
    tokens << token;
    QTimer::singleShot(30000, this, [=]() {
        tokens.removeAll(token);
    });
}

void SMSubServer::statusUpdated(const bool status_, const qint64 time)
{
    status = status_;
    status ? (startTime = time) : (stopTime = time);

    QByteArray output, json_output;
    for (auto it = sockets.constBegin(); it != sockets.constEnd(); it++) {
        if (isVariantTrue((*it)->property("ReceiveStatus"))) {
            if (isVariantTrue((*it)->property("ReceiveJson"))) {
                if (json_output.isEmpty()) {
#ifdef BOOST_JSON
                    boost::json::object object;
                    object["type"] = "status";
                    object["status"] = status ? "on" : "off";
                    object["start"] = startTime;
                    object["stop"] = stopTime;
                    json_output = QByteArray::fromStdString(boost::json::serialize(object) + '\n');
#else
                    QJsonObject object;
                    object["type"] = "status";
                    object["status"] = status ? "on" : "off";
                    object["start"] = startTime;
                    object["stop"] = stopTime;
                    json_output = QJsonDocument(object).toJson(QJsonDocument::Compact) + '\n';
#endif
                }
                sendMessage(*it, json_output);
            }
            else {
                if (output.isEmpty()) {
                    output = QString("Status: %1\nLast Start: %2\nLast Stop: %3\n").arg(status ? "on" : "off", QString::number(startTime), QString::number(stopTime)).toUtf8();
                }
                sendMessage(*it, json_output);
            }
        }
    }
}