/*
 * Copyright (C) 2009  Barracuda Networks, Inc.
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this library.  If not, see <https://www.gnu.org/licenses/>.
 *
 */

#include "stuntypes.h"

#include "stunutil.h"
#include "transportaddress.h"

#include <QtCrypto>
#include <stdio.h>

#define STRING_MAX_CHARS 127
#define STRING_MAX_BYTES 763

namespace XMPP {
using namespace StunUtil;
namespace StunTypes {
    static void xorIPv4(QByteArray *in, const quint8 *magic)
    {
        quint8 *p = (quint8 *)in->data();
        p[2] ^= magic[0];
        p[3] ^= magic[1];
        for (int n = 0; n < 4; ++n)
            p[n + 4] ^= magic[n];
    }

    static void xorIPv6(QByteArray *in, const quint8 *magic, const quint8 *id)
    {
        quint8 *p = (quint8 *)in->data();
        p[2] ^= magic[0];
        p[3] ^= magic[1];
        for (int n = 0; n < 4; ++n)
            p[n + 4] ^= magic[n];
        for (int n = 0; n < 12; ++n)
            p[n + 8] ^= id[n];
    }

    static bool validateString(const QByteArray &in, QString *out)
    {
        if (in.size() <= STRING_MAX_BYTES) {
            QString s = QString::fromUtf8(in);
            if (s.length() <= STRING_MAX_CHARS) {
                *out = s;
                return true;
            }
        }

        return false;
    }

    QByteArray createMappedAddress(const TransportAddress &addr)
    {
        QByteArray out;

        if (addr.addr.protocol() == QAbstractSocket::IPv6Protocol) {
            out              = QByteArray(20, 0);
            out[1]           = 0x02; // IPv6
            Q_IPV6ADDR addr6 = addr.addr.toIPv6Address();
            memcpy(out.data() + 4, addr6.c, 16);
        } else if (addr.addr.protocol() == QAbstractSocket::IPv4Protocol) {
            out    = QByteArray(8, 0);
            out[1] = 0x01; // IPv4
            write32((quint8 *)out.data() + 4, addr.addr.toIPv4Address());
        } else
            Q_ASSERT(0);

        write16((quint8 *)out.data() + 2, addr.port);
        return out;
    }

    QByteArray createUsername(const QString &username) { return username.left(STRING_MAX_CHARS).toUtf8(); }

    QByteArray createErrorCode(int code, const QString &reason)
    {
        QByteArray out(4, 0);

        int ih = code / 100;
        int il = code % 100;
        ih &= 0x07; // keep only lower 3 bits

        unsigned char ch = (unsigned char)ih;
        unsigned char cl = (unsigned char)il;
        out[2]           = ch;
        out[3]           = cl;
        out += reason.left(STRING_MAX_CHARS).toUtf8();
        return out;
    }

    QByteArray createUnknownAttributes(const QList<quint16> &typeList)
    {
        if (typeList.isEmpty())
            return QByteArray();

        QByteArray out(typeList.count() * 2, 0);
        for (int n = 0; n < typeList.count(); ++n)
            write16((quint8 *)out.data() + (n * 2), typeList[n]);
        return out;
    }

    QByteArray createRealm(const QString &realm) { return realm.left(STRING_MAX_CHARS).toUtf8(); }

    QByteArray createNonce(const QString &nonce) { return nonce.left(STRING_MAX_CHARS).toUtf8(); }

    QByteArray createXorMappedAddress(const TransportAddress &addr, const quint8 *magic, const quint8 *id)
    {
        QByteArray out = createMappedAddress(addr);
        if (addr.addr.protocol() == QAbstractSocket::IPv6Protocol)
            xorIPv6(&out, magic, id);
        else // IPv4
            xorIPv4(&out, magic);
        return out;
    }

    QByteArray createChannelNumber(quint16 i)
    {
        QByteArray val(4, 0);
        write16((quint8 *)val.data(), i);
        // bytes 2-3 are zeroed out
        return val;
    }

    QByteArray createLifetime(quint32 i)
    {
        QByteArray val(4, 0);
        write32((quint8 *)val.data(), i);
        return val;
    }

    QByteArray createXorPeerAddress(const TransportAddress &addr, const quint8 *magic, const quint8 *id)
    {
        return createXorMappedAddress(addr, magic, id);
    }

    QByteArray createXorRelayedAddress(const TransportAddress &addr, const quint8 *magic, const quint8 *id)
    {
        return createXorMappedAddress(addr, magic, id);
    }

    QByteArray createEvenPort(bool reserve)
    {
        QByteArray    val(1, 0);
        unsigned char c = 0;
        if (reserve)
            c |= 0x80; // set high bit
        val[0] = c;
        return val;
    }

    QByteArray createRequestedTransport(quint8 proto)
    {
        QByteArray val(4, 0);
        val[0] = proto;
        // bytes 1-3 are zeroed out
        return val;
    }

    QByteArray createReservationToken(const QByteArray &token)
    {
        Q_ASSERT(token.size() == 8);
        return token;
    }

    QByteArray createPriority(quint32 i)
    {
        QByteArray val(4, 0);
        write32((quint8 *)val.data(), i);
        return val;
    }

    QByteArray createSoftware(const QString &str) { return str.left(STRING_MAX_CHARS).toUtf8(); }

    QByteArray createAlternateServer(const TransportAddress &addr) { return createMappedAddress(addr); }

    QByteArray createIceControlled(quint64 i)
    {
        QByteArray val(8, 0);
        write64((quint8 *)val.data(), i);
        return val;
    }

    QByteArray createIceControlling(quint64 i)
    {
        QByteArray val(8, 0);
        write64((quint8 *)val.data(), i);
        return val;
    }

    bool parseMappedAddress(const QByteArray &val, TransportAddress &addr)
    {
        if (val[1] == 0x02 && val.size() == 20) // IPv6
        {
            addr.port      = read16((const quint8 *)val.data() + 2);
            QByteArray buf = val.mid(4);
            addr.addr      = QHostAddress((quint8 *)buf.data());
            return true;
        } else if (val[1] == 0x01 && val.size() == 8) // IPv4
        {
            addr.port = read16((const quint8 *)val.data() + 2);
            addr.addr = QHostAddress(read32((const quint8 *)val.data() + 4));
            return true;
        } else
            return false;
    }

    bool parseUsername(const QByteArray &val, QString *username) { return validateString(val, username); }

    bool parseErrorCode(const QByteArray &val, int *code, QString *reason)
    {
        if (val.size() < 4)
            return false;

        unsigned char ch = (unsigned char)val[2];
        unsigned char cl = (unsigned char)val[3];
        int           ih = ch & 0x07; // lower 3 bits
        int           x  = ih * 100 + (int)cl;

        QString str;
        if (validateString(val.mid(4), &str)) {
            *code   = x;
            *reason = str;
            return true;
        }

        return false;
    }

    bool parseUnknownAttributes(const QByteArray &val, QList<quint16> *typeList)
    {
        if (val.size() % 2 != 0)
            return false;

        typeList->clear();
        int count = val.size() / 2;
        for (int n = 0; n < count; ++n)
            typeList->append(read16((const quint8 *)val.data() + (n * 2)));
        return true;
    }

    bool parseRealm(const QByteArray &val, QString *realm) { return validateString(val, realm); }

    bool parseNonce(const QByteArray &val, QString *nonce) { return validateString(val, nonce); }

    bool parseXorMappedAddress(const QByteArray &val, const quint8 *magic, const quint8 *id, TransportAddress &addr)
    {
        if (val.size() < 4)
            return false;

        QByteArray buf;
        if (val[1] == 0x02 && val.size() == 20) // IPv6
        {
            buf = val;
            xorIPv6(&buf, magic, id);
        } else if (val[1] == 0x01 && val.size() == 8) // IPv4
        {
            buf = val;
            xorIPv4(&buf, magic);
        } else
            return false;

        return parseMappedAddress(buf, addr);
    }

    bool parseChannelNumber(const QByteArray &val, quint16 *i)
    {
        if (val.size() != 4)
            return false;

        const quint8 *p = (const quint8 *)val.data();
        *i              = read16(p);
        return true;
    }

    bool parseLifetime(const QByteArray &val, quint32 *i)
    {
        if (val.size() != 4)
            return false;

        const quint8 *p = (const quint8 *)val.data();
        *i              = read32(p);
        return true;
    }

    bool parseXorPeerAddress(const QByteArray &val, const quint8 *magic, const quint8 *id, TransportAddress &addr)
    {
        return parseXorMappedAddress(val, magic, id, addr);
    }

    bool parseXorRelayedAddress(const QByteArray &val, const quint8 *magic, const quint8 *id, TransportAddress &addr)
    {
        return parseXorMappedAddress(val, magic, id, addr);
    }

    bool parseEvenPort(const QByteArray &val, bool *reserve)
    {
        if (val.size() != 1)
            return false;

        unsigned char c = static_cast<unsigned char>(val[0]);
        *reserve        = (c & 0x80) != 0;

        return true;
    }

    bool parseRequestedTransport(const QByteArray &val, quint8 *proto)
    {
        if (val.size() != 4)
            return false;

        *proto = quint8(val[0]);
        return true;
    }

    bool parseReservationToken(const QByteArray &val, QByteArray *token)
    {
        if (val.size() != 8)
            return false;

        *token = val;
        return true;
    }

    bool parsePriority(const QByteArray &val, quint32 *i)
    {
        if (val.size() != 4)
            return false;

        const quint8 *p = (const quint8 *)val.data();
        *i              = read32(p);
        return true;
    }

    bool parseSoftware(const QByteArray &val, QString *str)
    {
        *str = QString::fromUtf8(val);
        return true;
    }

    bool parseAlternateServer(const QByteArray &val, TransportAddress &addr) { return parseMappedAddress(val, addr); }

    bool parseIceControlled(const QByteArray &val, quint64 *i)
    {
        if (val.size() != 8)
            return false;

        const quint8 *p = (const quint8 *)val.data();
        *i              = read64(p);
        return true;
    }

    bool parseIceControlling(const QByteArray &val, quint64 *i)
    {
        if (val.size() != 8)
            return false;

        const quint8 *p = (const quint8 *)val.data();
        *i              = read64(p);
        return true;
    }

#define METHOD_ENTRY(x)                                                                                                \
    {                                                                                                                  \
        x, #x                                                                                                          \
    }

    struct MethodEntry {
        Method      method;
        const char *str;
    } method_table[]
        = { METHOD_ENTRY(Binding), METHOD_ENTRY(Allocate),         METHOD_ENTRY(Refresh),     METHOD_ENTRY(Send),
            METHOD_ENTRY(Data),    METHOD_ENTRY(CreatePermission), METHOD_ENTRY(ChannelBind), { (Method)-1, nullptr } };

    QString methodToString(int method)
    {
        for (int n = 0; method_table[n].str; ++n) {
            if (method_table[n].method == (Method)method)
                return QString::fromLatin1(method_table[n].str);
        }

        return QString();
    }

#define ATTRIB_ENTRY(x)                                                                                                \
    {                                                                                                                  \
        x, #x                                                                                                          \
    }

    struct AttribEntry {
        Attribute   type;
        const char *str;
    } attrib_table[] = { ATTRIB_ENTRY(MAPPED_ADDRESS),
                         ATTRIB_ENTRY(USERNAME),
                         ATTRIB_ENTRY(MESSAGE_INTEGRITY),
                         ATTRIB_ENTRY(ERROR_CODE),
                         ATTRIB_ENTRY(UNKNOWN_ATTRIBUTES),
                         ATTRIB_ENTRY(REALM),
                         ATTRIB_ENTRY(NONCE),
                         ATTRIB_ENTRY(XOR_MAPPED_ADDRESS),
                         ATTRIB_ENTRY(CHANNEL_NUMBER),
                         ATTRIB_ENTRY(LIFETIME),
                         ATTRIB_ENTRY(XOR_PEER_ADDRESS),
                         ATTRIB_ENTRY(DATA),
                         ATTRIB_ENTRY(XOR_RELAYED_ADDRESS),
                         ATTRIB_ENTRY(EVEN_PORT),
                         ATTRIB_ENTRY(REQUESTED_TRANSPORT),
                         ATTRIB_ENTRY(DONT_FRAGMENT),
                         ATTRIB_ENTRY(RESERVATION_TOKEN),
                         ATTRIB_ENTRY(PRIORITY),
                         ATTRIB_ENTRY(USE_CANDIDATE),
                         ATTRIB_ENTRY(SOFTWARE),
                         ATTRIB_ENTRY(ALTERNATE_SERVER),
                         ATTRIB_ENTRY(FINGERPRINT),
                         ATTRIB_ENTRY(ICE_CONTROLLED),
                         ATTRIB_ENTRY(ICE_CONTROLLING),
                         { (Attribute)-1, nullptr } };

    QString attributeTypeToString(int type)
    {
        for (int n = 0; attrib_table[n].str; ++n) {
            if (attrib_table[n].type == (Attribute)type) {
                QString name = QString::fromLatin1(attrib_table[n].str);
                name.replace('_', '-');
                return name;
            }
        }

        return QString();
    }

    static QString quoted(const QString &in) { return QString("\"") + in + '\"'; }

    QString attributeValueToString(int type, const QByteArray &val, const quint8 *magic, const quint8 *id)
    {
        switch ((Attribute)type) {
        case MAPPED_ADDRESS: {
            TransportAddress addr;
            if (parseMappedAddress(val, addr))
                return QString(addr);
            break;
        }
        case USERNAME: {
            QString str;
            if (parseUsername(val, &str))
                return quoted(str);
            break;
        }
        case MESSAGE_INTEGRITY: {
            return QCA::arrayToHex(val);
        }
        case ERROR_CODE: {
            int     code;
            QString reason;
            if (parseErrorCode(val, &code, &reason)) {
                QString out = QString::number(code);
                if (!reason.isEmpty())
                    out += QString(", ") + quoted(reason);
                return out;
            }
            break;
        }
        case UNKNOWN_ATTRIBUTES: {
            QList<quint16> typeList;
            if (parseUnknownAttributes(val, &typeList)) {
                if (!typeList.isEmpty()) {
                    QStringList strList;
                    for (quint16 i : qAsConst(typeList))
                        strList += QString::asprintf("0x%04x", i);
                    return strList.join(", ");
                } else
                    return "(None)";
            }
            break;
        }
        case REALM: {
            QString str;
            if (parseRealm(val, &str))
                return quoted(str);
            break;
        }
        case NONCE: {
            QString str;
            if (parseNonce(val, &str))
                return quoted(str);
            break;
        }
        case XOR_MAPPED_ADDRESS: {
            TransportAddress addr;
            if (parseXorMappedAddress(val, magic, id, addr))
                return QString(addr);
            break;
        }
        case CHANNEL_NUMBER: {
            quint16 i;
            if (parseChannelNumber(val, &i))
                return QString::asprintf("0x%04x", (int)i);
            break;
        }
        case LIFETIME: {
            quint32 i;
            if (parseLifetime(val, &i))
                return QString::number(i);
            break;
        }
        case XOR_PEER_ADDRESS: {
            return attributeValueToString(XOR_MAPPED_ADDRESS, val, magic, id);
        }
        case DATA: {
            return QString("len=%1, ").arg(val.size()) + QCA::arrayToHex(val);
        }
        case XOR_RELAYED_ADDRESS: {
            return attributeValueToString(XOR_MAPPED_ADDRESS, val, magic, id);
        }
        case EVEN_PORT: {
            bool reserve;
            if (parseEvenPort(val, &reserve))
                return QString("reserve=") + (reserve ? "true" : "false");
            break;
        }
        case REQUESTED_TRANSPORT: {
            quint8 proto;
            if (parseRequestedTransport(val, &proto)) {
                QString str = QString::number((int)proto);
                if (proto == 17) // UDP
                    str += " (UDP)";
                else
                    str += " (Unknown)";
                return str;
            }
            break;
        }
        case DONT_FRAGMENT: {
            return QString("");
        }
        case RESERVATION_TOKEN: {
            QByteArray token;
            if (parseReservationToken(val, &token))
                return QCA::arrayToHex(token);
            break;
        }
        case PRIORITY: {
            quint32 i;
            if (parsePriority(val, &i))
                return QString::number(i);
            break;
        }
        case USE_CANDIDATE: {
            return QString("");
        }
        case SOFTWARE: {
            QString out;
            if (parseSoftware(val, &out))
                return quoted(out);
            break;
        }
        case ALTERNATE_SERVER: {
            return attributeValueToString(MAPPED_ADDRESS, val, magic, id);
        }
        case FINGERPRINT: {
            return QCA::arrayToHex(val);
        }
        case ICE_CONTROLLED: {
            quint64 i;
            if (parseIceControlled(val, &i))
                return QString::number(i);
            break;
        }
        case ICE_CONTROLLING: {
            quint64 i;
            if (parseIceControlling(val, &i))
                return QString::number(i);
            break;
        }
        }

        return QString();
    }

    QString print_packet_str(const StunMessage &message)
    {
        QString out;

        QString mclass;
        if (message.mclass() == StunMessage::Request)
            mclass = "Request";
        else if (message.mclass() == StunMessage::SuccessResponse)
            mclass = "Response (Success)";
        else if (message.mclass() == StunMessage::ErrorResponse)
            mclass = "Response (Error)";
        else if (message.mclass() == StunMessage::Indication)
            mclass = "Indication";
        else
            Q_ASSERT(0);

        out += QString("Class: %1\n").arg(mclass);
        out += QString("Method: %1\n").arg(methodToString(message.method()));
        out += QString("Transaction id: %1\n").arg(QCA::arrayToHex(QByteArray((const char *)message.id(), 12)));
        out += "Attributes:";
        QList<StunMessage::Attribute> attribs = message.attributes();
        if (!attribs.isEmpty()) {
            for (const StunMessage::Attribute &a : attribs) {
                out += '\n';

                QString name = attributeTypeToString(a.type);
                if (!name.isNull()) {
                    QString val = attributeValueToString(a.type, a.value, message.magic(), message.id());
                    if (val.isNull())
                        val = QString("Unable to parse %1 bytes").arg(a.value.size());

                    out += QString("  %1").arg(name);
                    if (!val.isEmpty())
                        out += QString(" = %1").arg(val);
                } else
                    out += QString::asprintf("  Unknown attribute (0x%04x) of %d bytes", a.type, a.value.size());
            }
        } else
            out += "\n  (None)";

        return out;
    }

    void print_packet(const StunMessage &message) { printf("%s\n", qPrintable(print_packet_str(message))); }
} // namespace StunTypes
} // namespace XMPP
