#include "Message.h"

#include <string.h>
#include <stdexcept>
#include <cassert>

#include <QDebug>

enum FieldType {
    JailId = 1,
    ExecutablePath = 10,
    Argument,
    InitScript,
    JailPassword,          // The jaildir is encfs encrypted. Decrypt password.

    IsTry = 20,             // allow next command to fail
    RBindMountSource,       // mount from an existing directory.
    RBindMountDest,
    UnMountDir,             // unmount (remove) an existing mountpoint.
    CreateTmpFs,            // with the full path as arg

    CopyFrom = 30,
    CopyTo,

    EnvironUnset = 35,      // an env-var we want to filter
    EnvironSet,             // override env-var value.

    DBusProxyFrom = 40,     // start the xdg-dbus-proxy between from and to
    DBusProxyTo,
    DBusProxySystemFrom,    // start the dbus proxy with system rules
    DBusProxySystemTo,

    VpnConfig = 50,      // base path to vpn files (inside jail)
    VpnHasPassFile = 51,    // has a password file in the jail
};

// returns the length of the string, or throw if no closing zero found
static int stringLength(const char *from, const char *end)
{
    auto at = from;
    while (at <= end) {
        if (*at == 0)
            return at - from;
        ++at;
    }
    throw std::runtime_error("unterminated string found");
}

Message::Message(int size)
    : m_reservedSize(size)
{
    assert(size > 0);
    m_buf = std::shared_ptr<char>(new char[size], std::default_delete<char[]>());
    m_writePtr = m_buf.get();
    memset(m_writePtr, 0, m_reservedSize);
}

Message::Message(char *buffer, int bufferSize)
    : m_reservedSize(bufferSize)
{
    m_buf = std::shared_ptr<char>(new char[bufferSize], std::default_delete<char[]>());
    memcpy(m_buf.get(), buffer, bufferSize);
    buffer = m_buf.get();

    const char *end = buffer + bufferSize;
    while (buffer < end) {
        if (*buffer == ExecutablePath) {
            m_path = buffer + 1;
        }
        else if (*buffer == Argument) {
            break;
        }
        const int len = ::stringLength(buffer + 1, end);
        if (buffer + len + 2 > end) {
            throw std::runtime_error("buffer too small");
        }
        buffer += len + 2;
    }
}

char *Message::path() const
{
    return m_path;
}

void Message::setPath(const std::string &path)
{
    m_path = m_writePtr + 1;
    addString(ExecutablePath, path);
}

void Message::addArgument(const char *string)
{
    addString(Argument, string);
}

void Message::printFields() const
{
#ifndef NDEBUG
    qDebug().nospace() << "Message[" << path() << "] " << size() << " bytes";

    const char *buffer  = m_buf.get();
    const char *end = buffer + m_reservedSize;
    while (buffer < end) {
        std::string prefix;
        const char type = *buffer;
        if (type == Argument) {
            prefix = "argument:";
        }
        else if (type == RBindMountSource) {
            prefix = "rbind-src: ";
        }
        else if (type == RBindMountDest) {
            prefix = "rbind-dst: ";
        }
        else if (type == UnMountDir) {
            prefix = "umountDir: ";
        }
        else if (type == CreateTmpFs) {
            prefix = "tmpFs-create: ";
        }
        else if (type == CopyFrom) {
            prefix = "copyFrom: ";
        }
        else if (type == CopyTo) {
            prefix = "copyTo: ";
        }
        else if (type == EnvironSet) {
            prefix = "New Env: ";
        }
        else if (type == EnvironUnset) {
            prefix = "Filter Env: ";
        }
        else if (type == DBusProxyFrom) {
            prefix = "dbus-from: ";
        }
        else if (type == DBusProxySystemFrom) {
            prefix = "dbus-system-from: ";
        }
        else if (type == DBusProxyTo) {
            prefix = "dbus-to: ";
        }
        else if (type == DBusProxySystemTo) {
            prefix = "dbus-system-to: ";
        }
        else if (type == VpnConfig) {
            prefix = "vpn-config: ";
        }
        else if (type == JailId) {
            // This only works if the source and destination have the same endian-ness
            // as the basic design places them in the same exe just forked into two
            // processes, this is not a concern.
            uint32_t data = buffer[4];
            data <<= 8;
            data += buffer[3];
            data <<= 8;
            data += buffer[2];
            data <<= 8;
            data += buffer[1];
            buffer += 5;
            qDebug() << "=> JailId:" << data;
            continue;
        }
        else if (type == VpnHasPassFile) { // bool
            qDebug() << "=> hasPassFile:" << (buffer[1] != 0);
            buffer += 2;
            continue;
        }
        else if (type == 0) {
            qDebug() << Qt::endl;
            return;
        }
        const int len = ::stringLength(buffer + 1, end);
        if (buffer + len + 2 > end) {
            qWarning() << "buffer overrun detected, str len:" << len;
            return;
        }
        if (type == IsTry) {
            qWarning() << "  try";
            buffer += 1;
            continue;
        }
        else if (type != ExecutablePath) {
            qWarning() << "  " << prefix.c_str() << QString::fromLocal8Bit(buffer + 1, len);
        }
        buffer += len + 2;
    }
#endif
}

char *Message::begin() const
{
    return m_buf.get();
}

int Message::size() const
{
    if (m_writePtr)
        return m_writePtr - m_buf.get();
    return m_reservedSize;
}

void Message::addUmountPoint(const std::string &dir)
{
    addString(UnMountDir, dir);
}

void Message::addRemount(const std::string &source, const std::string &destination)
{
    addString(RBindMountSource, source); // always first
    addString(RBindMountDest, destination); // make sure this is last as it 'executes' the whole thing.
}

void Message::setTry(bool isTry)
{
    if (isTry)
        addTag(IsTry); // always first
}

void Message::setJailId(uint32_t jailId)
{
    assert(m_writePtr);
    if (m_buf.get() + m_reservedSize < m_writePtr + 5) {
        throw std::runtime_error("jailId does not fit buffer");
    }
    // This only works if the source and destination have the same endian-ness
    // as the basic design places them in the same exe just forked into two
    // processes, this is not a concern.
    m_writePtr[0] = JailId;
    m_writePtr[1] = jailId & 0xFF;
    m_writePtr[2] = (jailId >> 8) & 0xFF;
    m_writePtr[3] = (jailId >> 16) & 0xFF;
    m_writePtr[4] = (jailId >> 24) & 0xFF;
    m_writePtr += 5;
}

void Message::addMountTmpDir(const std::string &dir)
{
    addString(CreateTmpFs, dir);
}

void Message::addCopy(const std::string &from, const std::string &to)
{
    addString(CopyFrom, from);
    addString(CopyTo, to);
}

void Message::addInitSript(const std::string &text)
{
    addString(InitScript, text);
}

void Message::addEnvToUnset(const std::string &propertyName)
{
    assert(propertyName.find('=') == std::string::npos);
    addString(EnvironUnset, propertyName);
}

void Message::addEnvToSet(const std::string &envVar)
{
    assert(envVar.find('=') != std::string::npos);
    addString(EnvironSet, envVar);
}

void Message::setJailPassword(const std::string &pwd)
{
    addString(JailPassword, pwd);
}

void Message::addDBusProxy(DBusType type, const std::string &from, const std::string &to)
{
    addString(type == UserSessionBus ? DBusProxyFrom : DBusProxySystemFrom, from);
    addString(type == UserSessionBus ? DBusProxyTo : DBusProxySystemTo, to);
}

void Message::setVpnConfig(const std::string &configPath)
{
    addString(VpnConfig, configPath);
}

void Message::setHasVpnPassFile(bool yes)
{
    assert(m_writePtr);
    if (m_buf.get() + m_reservedSize < m_writePtr + 2)
        throw std::runtime_error("Bool does not fit buffer");
    m_writePtr[0] = VpnHasPassFile;
    m_writePtr[1] = yes ? 1 : 0;
    m_writePtr += 2;
}

void Message::addTag(char type)
{
    assert(m_writePtr);
    if (m_buf.get() + m_reservedSize < m_writePtr + 1) {
        throw std::runtime_error("Tag does not fit buffer");
    }
    *m_writePtr = type;
    ++m_writePtr;
}

void Message::addString(char type, const std::string &string)
{
    assert(m_writePtr);
    if (m_buf.get() + m_reservedSize < m_writePtr + string.size() + 2) {
        // won't fit.
        throw std::runtime_error("String does not fit buffer");
    }
    m_writePtr[0] = type;
    memcpy(++m_writePtr, string.c_str(), string.size());
    m_writePtr += string.size();
    m_writePtr[0] = 0; // trailing zero
    ++m_writePtr;
}


// ///////////////////////////////////////////////////////////////

Message::Iterator::Iterator(const Message * const message)
    : m_parent(message),
      m_cur(message->begin()),
      m_recordSize(-1)
{
    next();
}

bool Message::Iterator::isArgument() const
{
    assert(isValid());
    return m_cur[0] == Argument;
}

bool Message::Iterator::isNewEnvVar() const
{
    assert(isValid());
    return m_cur[0] == EnvironSet;
}

bool Message::Iterator::isEnvVarUnset() const
{
    assert(isValid());
    return m_cur[0] == EnvironUnset;
}

bool Message::Iterator::isUnmount() const
{
    assert(isValid());
    return m_cur[0] == UnMountDir;
}

bool Message::Iterator::isRemount() const
{
    assert(isValid());
    return m_cur[0] == RBindMountSource;
}

bool Message::Iterator::isCreateTmp() const
{
    assert(isValid());
    return m_cur[0] == CreateTmpFs;
}

bool Message::Iterator::isCopy() const
{
    assert(isValid());
    return m_cur[0] == CopyFrom;
}

bool Message::Iterator::isJailId() const
{
    assert(isValid());
    return m_cur[0] == JailId;
}

bool Message::Iterator::isInitSript() const
{
    assert(isValid());
    return m_cur[0] == InitScript;
}

bool Message::Iterator::isJailPwd() const
{
    assert(isValid());
    return m_cur[0] == JailPassword;
}

bool Message::Iterator::isVpnConfig() const
{
    assert(isValid());
    return m_cur[0] == VpnConfig;
}

bool Message::Iterator::isVpnPwdBool() const
{
    assert(isValid());
    return m_cur[0] == VpnHasPassFile;
}

bool Message::Iterator::isValid() const
{
    assert(m_parent);
    assert(m_cur);
    assert(m_cur >= m_parent->begin());
    if (m_cur < m_parent->begin() + m_parent->size()) {
        return m_cur[0] != 0;
    }
    return false;
}

bool Message::Iterator::isDBusMapping() const
{
    return m_cur[0] == DBusProxyFrom || m_cur[0] == DBusProxySystemFrom;
}

bool Message::Iterator::boolData() const
{
    return m_cur[0] != 0;
}

CopyMessage Message::Iterator::copyData() const
{
    assert(m_parent);
    CopyMessage rc;
    if (m_cur[0] != CopyFrom)
        throw std::runtime_error("Not copy data");

    const char *end = m_parent->m_buf.get() + m_parent->m_reservedSize;
    auto length = ::stringLength(m_cur + 1, end);
    rc.from = std::string(m_cur + 1, length);
    auto length2 = ::stringLength(m_cur + 1 + length + 2, end);
    rc.to = std::string(m_cur + 1 + length + 2, length2);
    return rc;
}

DBusMapping Message::Iterator::dbusMapping() const
{
    assert(m_parent);
    DBusMapping rc;
    if (m_cur[0] == DBusProxyFrom)
        rc.systemBus = false;
    else if (m_cur[0] == DBusProxySystemFrom)
        rc.systemBus = true;
    else
        throw std::runtime_error("Not dbus-mapping data");

    const char *end = m_parent->m_buf.get() + m_parent->m_reservedSize;
    auto length = ::stringLength(m_cur + 1, end);
    rc.from = std::string(m_cur + 1, length);
    auto length2 = ::stringLength(m_cur + 1 + length + 2, end);
    rc.to = std::string(m_cur + 1 + length + 2, length2);
    return rc;
}

uint32_t Message::Iterator::jailId() const
{
    assert(m_parent);
    const char *end = m_parent->m_buf.get() + m_parent->m_reservedSize;
    if (m_cur + 4 > end)
        throw std::runtime_error("Not enough space in message");
    assert(isJailId()); // check if it actually is a jailid
    uint32_t data = m_cur[4];
    data <<= 8;
    data += m_cur[3];
    data <<= 8;
    data += m_cur[2];
    data <<= 8;
    data += m_cur[1];
    return data;
}

char *Message::Iterator::stringPtr() const
{
    return m_cur + 1;
}

int Message::Iterator::stringLength() const
{
    assert(isValid());
    // not valid for one of the commands that have more than one string.
    assert(!isRemount() && !isCopy());
    return m_recordSize - 2;
}

MountMessage Message::Iterator::mountData() const
{
    assert(m_parent);
    MountMessage rc;
    if (m_cur[0] == RBindMountSource)
        rc.type = MountMessage::Remount;
    else if (m_cur[0] == UnMountDir)
        rc.type = MountMessage::Umount;
    else if (m_cur[0] == CreateTmpFs)
        rc.type = MountMessage::CreateTmpFs;
    else
        throw std::runtime_error("Not mount data");

    const char *end = m_parent->m_buf.get() + m_parent->m_reservedSize;
    int index = 0;
    while (true) {
        char t = m_cur[index];
        const char *str = m_cur + index + 1;
        const int strSize = ::stringLength(m_cur + index + 1, end);
        index += strSize + 2;
        checkAvail(index);
        if (t == RBindMountSource) {
            rc.src = std::string(str, strSize);
        }
        else if (t == RBindMountDest) {
            rc.dst = std::string(str, strSize);
            break; // last in the sequence
        }
        else if (rc.type == MountMessage::Umount) {
            rc.src = std::string(str, strSize);
            // only the src field is expected
            break;
        }
        else if (rc.type == MountMessage::CreateTmpFs) {
            rc.dst = std::string(str, strSize);
            // only the dst field is expected
            break;
        }
    }

    return rc;
}

const char *Message::Iterator::argument() const
{
    assert(m_parent);
    assert(isArgument());
    assert(m_recordSize != -1); // we called 'next'
    return m_cur + 1;
}

bool Message::Iterator::next()
{
    assert(m_cur);
    if (m_recordSize != -1) {
        m_cur = m_cur + m_recordSize;
        if (!isValid())
            return false;
    }

    m_isTry = *m_cur == IsTry;
    if (m_isTry)
        m_cur += 1;

    m_recordSize = 0;
    checkAvail(2);
    // find the string-size(s)
    const char *end = m_parent->m_buf.get() + m_parent->m_reservedSize;
    switch (*m_cur) {
    case ExecutablePath: // fall through
    case Argument: // fall through
    case UnMountDir: // fall through
    case EnvironSet: // fall through
    case EnvironUnset: // fall through
    case InitScript: // fall through
    case JailPassword: // fall through
    case VpnConfig: // fall through
    case CreateTmpFs:
        m_recordSize = ::stringLength(m_cur + 1, end) + 2;
        break;
    case DBusProxyFrom: // fall through
    case DBusProxySystemFrom: // fall through
    case RBindMountSource: // fall-through
    case CopyFrom:
        // these take 2 args
        m_recordSize = ::stringLength(m_cur + 1, end) + 2;
        m_recordSize += ::stringLength(m_cur + m_recordSize + 1, end) + 2;
        break;
    case JailId:
        m_recordSize += 5;
        break;
    case VpnHasPassFile: // bool
        m_recordSize += 2;
        break;
    case DBusProxyTo: // fall through
    case DBusProxySystemTo: // fall through
    case RBindMountDest:
        assert(false);
    default:
        assert(false);
    }
    return true;
}

void Message::Iterator::checkAvail(int bytes) const
{
    if (bytes < 0 || bytes > MAX_SIZE)
        throw std::runtime_error("Serialization failure: impossible size");
    if (m_parent->begin() + m_parent->size() < m_cur + bytes)
        throw std::runtime_error("Message::Iterator: out of bounds");
}