/* This file is part of the KDE libraries
   Copyright (c) 2004 Szombathelyi Gy�gy <gyurco@freemail.hu>

   The implementation is based on the documentation and sample code
   at http://davenport.sourceforge.net/ntlm.html
   The DES encryption functions are from libntlm
   at http://josefsson.org/libntlm/

   This library is free software; you can redistribute it and/or
   modify it under the terms of the GNU Library General Public
   License version 2 as published by the Free Software Foundation.

   This library 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
   Library General Public License for more details.

   You should have received a copy of the GNU Library General Public License
   along with this library; see the file COPYING.LIB.  If not, write to
   the Free Software Foundation, Inc., 51 Franklin Street, Fifth Floor,
   Boston, MA 02110-1301, USA.
*/

#include "kntlm.h"
#include "des.h"

#include <cstring>

#include <QtCore/QDate>
#include <QtCore/QtEndian>
#include <QtCore/QCryptographicHash>

#include <krandom.h>

static const char NTLM_SIGNATURE[] = "NTLMSSP";

static QByteArray QString2UnicodeLE(const QString &target)
{
    QByteArray unicode(target.length() * 2, 0);

    for (int i = 0; i < target.length(); i++) {
        ((quint16 *) unicode.data()) [ i ] = qToLittleEndian(target[i].unicode());
    }

    return unicode;
}

static QString UnicodeLE2QString(const QChar *data, uint len)
{
    QString ret;

    for (uint i = 0; i < len; i++) {
        ret += qFromLittleEndian(data[ i ].unicode());
    }

    return ret;
}

static QByteArray getBuf(const QByteArray &buf, const KNTLM::SecBuf &secbuf)
{
    quint32 offset = qFromLittleEndian(secbuf.offset);
    quint16 len = qFromLittleEndian(secbuf.len);

    //watch for buffer overflows
    if (offset > (quint32) buf.size() || offset + len > (quint32) buf.size()) {
        return QByteArray();
    }

    return QByteArray(buf.data() + offset, buf.size());
}

static void addBuf(QByteArray &buf, KNTLM::SecBuf &secbuf, const QByteArray &data)
{
    quint32 offset = (buf.size() + 1) & 0xfffffffe;
    quint16 len = data.size();
    quint16 maxlen = data.size();

    secbuf.offset = qToLittleEndian((quint32) offset);
    secbuf.len = qToLittleEndian(len);
    secbuf.maxlen = qToLittleEndian(maxlen);
    buf.resize(offset + len);
    memcpy(buf.data() + offset, data.data(), data.size());
}

static QString getString(const QByteArray &buf, const KNTLM::SecBuf &secbuf, bool unicode)
{
    //watch for buffer overflows
    quint32 offset = qFromLittleEndian((quint32) secbuf.offset);
    quint16 len = qFromLittleEndian(secbuf.len);

    if (offset > (quint32) buf.size() || offset + len > (quint32) buf.size()) {
        return QString();
    }

    const char *c = buf.data() + offset;

    if (unicode) {
        return UnicodeLE2QString((QChar *) c, len >> 1);
    }

    return QString::fromLatin1(c, len);
}

static void addString(QByteArray &buf, KNTLM::SecBuf &secbuf, const QString &str, bool unicode = false)
{
    if (unicode) {
        addBuf(buf, secbuf, QString2UnicodeLE(str));
        return;
    }

    addBuf(buf, secbuf, str.toLatin1());
}

/*
* turns a 56 bit key into the 64 bit, odd parity key and sets the key.
* The key schedule ks is also set.
*/
static void convertKey(unsigned char *key_56, void *ks)
{
    unsigned char key[8];

    key[0] = key_56[0];
    key[1] = ((key_56[0] << 7) & 0xFF) | (key_56[1] >> 1);
    key[2] = ((key_56[1] << 6) & 0xFF) | (key_56[2] >> 2);
    key[3] = ((key_56[2] << 5) & 0xFF) | (key_56[3] >> 3);
    key[4] = ((key_56[3] << 4) & 0xFF) | (key_56[4] >> 4);
    key[5] = ((key_56[4] << 3) & 0xFF) | (key_56[5] >> 5);
    key[6] = ((key_56[5] << 2) & 0xFF) | (key_56[6] >> 6);
    key[7] = (key_56[6] << 1) & 0xFF;

    for (uint i = 0; i < 8; i++) {
        unsigned char b = key[i];
        bool needsParity = ((((b >> 7) ^ (b >> 6) ^ (b >> 5) ^ (b >> 4) ^ (b >> 3) ^ (b >> 2) ^ (b >> 1)) & 0x01) == 0);

        if (needsParity) {
            key[i] |= 0x01;
        } else {
            key[i] &= 0xfe;
        }
    }

    ntlm_des_set_key((DES_KEY *) ks, (char *) &key, sizeof(key));
    memset(&key, 0, sizeof(key));
}

static QByteArray createBlob(const QByteArray &targetinfo)
{
    QByteArray blob(sizeof(KNTLM::Blob) + 4 + targetinfo.size(), 0);

    KNTLM::Blob *bl = (KNTLM::Blob *) blob.data();
    bl->signature = qToBigEndian((quint32) 0x01010000);
    quint64 now = QDateTime::currentDateTime().toTime_t();
    now += (quint64) 3600 * (quint64) 24 * (quint64) 134774;
    now *= (quint64) 10000000;
    bl->timestamp = qToLittleEndian(now);

    for (uint i = 0; i < 8; i++) {
        bl->challenge[i] = KRandom::random() % 0xff;
    }

    memcpy(blob.data() + sizeof(KNTLM::Blob), targetinfo.data(), targetinfo.size());
    return blob;
}

static QByteArray hmacMD5(const QByteArray &data, const QByteArray &key)
{
    quint8 ipad[64], opad[64];
    QByteArray ret;

    memset(ipad, 0x36, sizeof(ipad));
    memset(opad, 0x5c, sizeof(opad));

    for (int i = key.size() - 1; i >= 0; i--) {
        ipad[i] ^= key[i];
        opad[i] ^= key[i];
    }

    QByteArray content(data.size() + 64, 0);
    memcpy(content.data(), ipad, 64);
    memcpy(content.data() + 64, data.data(), data.size());

    QCryptographicHash md5(QCryptographicHash::Md5);
    md5.addData(content);
    content.resize(64);
    memcpy(content.data(), opad, 64);
    content += md5.result();

    md5.reset();
    md5.addData(content);

    return md5.result();
}

/*************************************** KNTLM implementation ***************************************/

bool KNTLM::getNegotiate(QByteArray &negotiate, const QString &domain, const QString &workstation, quint32 flags)
{
    QByteArray rbuf(sizeof(Negotiate), 0);

    memcpy(rbuf.data(), NTLM_SIGNATURE, sizeof(NTLM_SIGNATURE));
    ((Negotiate *) rbuf.data())->msgType = qToLittleEndian((quint32) 1);

    if (!domain.isEmpty()) {
        flags |= Negotiate_Domain_Supplied;
        addString(rbuf, ((Negotiate *) rbuf.data())->domain, domain);
    }

    if (!workstation.isEmpty()) {
        flags |= Negotiate_WS_Supplied;
        addString(rbuf, ((Negotiate *) rbuf.data())->workstation, workstation);
    }

    ((Negotiate *) rbuf.data())->flags = qToLittleEndian(flags);
    negotiate = rbuf;
    return true;
}

bool KNTLM::getAuth(QByteArray &auth, const QByteArray &challenge,
                    const QString &user, const QString &password, const QString &domain,
                    const QString &workstation, AuthFlags authflags)
{
    QByteArray rbuf(sizeof(Auth), 0);
    Challenge *ch = (Challenge *) challenge.data();
    QByteArray response;
    const uint chsize = challenge.size();
    bool unicode = false;
    QString dom;

    //challenge structure too small
    if (chsize < 32) {
        return false;
    }

    unicode = qFromLittleEndian(ch->flags) & Negotiate_Unicode;

    if (domain.isEmpty()) {
        dom = getString(challenge, ch->targetName, unicode);
    } else {
        dom = domain;
    }

    memcpy(rbuf.data(), NTLM_SIGNATURE, sizeof(NTLM_SIGNATURE));
    ((Auth *) rbuf.data())->msgType = qToLittleEndian((quint32) 3);
    ((Auth *) rbuf.data())->flags = ch->flags;
    QByteArray targetInfo;
    if (chsize >= sizeof(Challenge)) {
        targetInfo = getBuf(challenge, ch->targetInfo);
    }

    if (!(authflags & Force_V1) &&
            ((authflags & Force_V2) ||
             (!targetInfo.isEmpty() && (qFromLittleEndian(ch->flags) & Negotiate_Target_Info))) /* may support NTLMv2 */) {
        bool ret = false;

        if (qFromLittleEndian(ch->flags) & Negotiate_NTLM) {
            if (targetInfo.isEmpty()) {
                return false;
            }

            response = getNTLMv2Response(dom, user, password, targetInfo, ch->challengeData);
            addBuf(rbuf, ((Auth *) rbuf.data())->ntResponse, response);
            ret = true;
        }

        if (authflags & Add_LM) {
            response = getLMv2Response(dom, user, password, ch->challengeData);
            addBuf(rbuf, ((Auth *) rbuf.data())->lmResponse, response);
            ret = true;
        }

        if (!ret) {
            return false;
        }
    } else { //if no targetinfo structure and NTLMv2 or LMv2 not forced, or v1 forced, try the older methods
        bool ret = false;

        if (qFromLittleEndian(ch->flags) & Negotiate_NTLM) {
            response = getNTLMResponse(password, ch->challengeData);
            addBuf(rbuf, ((Auth *) rbuf.data())->ntResponse, response);
            ret = true;
        }

        if (authflags & Add_LM) {
            response = getLMResponse(password, ch->challengeData);
            addBuf(rbuf, ((Auth *) rbuf.data())->lmResponse, response);
            ret = true;
        }

        if (!ret) {
            return false;
        }
    }

    if (!dom.isEmpty()) {
        addString(rbuf, ((Auth *) rbuf.data())->domain, dom, unicode);
    }

    addString(rbuf, ((Auth *) rbuf.data())->user, user, unicode);

    if (!workstation.isEmpty()) {
        addString(rbuf, ((Auth *) rbuf.data())->workstation, workstation, unicode);
    }

    auth = rbuf;
    return true;
}

QByteArray KNTLM::getLMResponse(const QString &password, const unsigned char *challenge)
{
    QByteArray hash, answer;

    hash = lmHash(password);
    hash.resize(21);
    memset(hash.data() + 16, 0, 5);
    answer = lmResponse(hash, challenge);
    hash.fill(0);
    return answer;
}

QByteArray KNTLM::lmHash(const QString &password)
{
    QByteArray keyBytes(14, 0);
    QByteArray hash(16, 0);
    DES_KEY ks;
    const char *magic = "KGS!@#$%";

    strncpy(keyBytes.data(), password.toUpper().toLocal8Bit().constData(), 14);

    convertKey((unsigned char *) keyBytes.data(), &ks);
    ntlm_des_ecb_encrypt(magic, 8, &ks, (unsigned char *) hash.data());

    convertKey((unsigned char *) keyBytes.data() + 7, &ks);
    ntlm_des_ecb_encrypt(magic, 8, &ks, (unsigned char *) hash.data() + 8);

    keyBytes.fill(0);
    memset(&ks, 0, sizeof(ks));

    return hash;
}

QByteArray KNTLM::lmResponse(const QByteArray &hash, const unsigned char *challenge)
{
    DES_KEY ks;
    QByteArray answer(24, 0);

    convertKey((unsigned char *) hash.data(), &ks);
    ntlm_des_ecb_encrypt(challenge, 8, &ks, (unsigned char *) answer.data());

    convertKey((unsigned char *) hash.data() + 7, &ks);
    ntlm_des_ecb_encrypt(challenge, 8, &ks, (unsigned char *) answer.data() + 8);

    convertKey((unsigned char *) hash.data() + 14, &ks);
    ntlm_des_ecb_encrypt(challenge, 8, &ks, (unsigned char *) answer.data() + 16);

    memset(&ks, 0, sizeof(ks));
    return answer;
}

QByteArray KNTLM::getNTLMResponse(const QString &password, const unsigned char *challenge)
{
    QByteArray hash = ntlmHash(password);
    hash.resize(21);
    memset(hash.data() + 16, 0, 5);
    QByteArray answer = lmResponse(hash, challenge);
    hash.fill(0);
    return answer;
}

QByteArray KNTLM::ntlmHash(const QString &password)
{
    QByteArray unicode;
    unicode = QString2UnicodeLE(password);

    return QCryptographicHash::hash(unicode, QCryptographicHash::Md4);
}

QByteArray KNTLM::getNTLMv2Response(const QString &target, const QString &user,
                                    const QString &password, const QByteArray &targetInformation,
                                    const unsigned char *challenge)
{
    QByteArray hash = ntlmv2Hash(target, user, password);
    QByteArray blob = createBlob(targetInformation);
    return lmv2Response(hash, blob, challenge);
}

QByteArray KNTLM::getLMv2Response(const QString &target, const QString &user,
                                  const QString &password, const unsigned char *challenge)
{
    QByteArray hash = ntlmv2Hash(target, user, password);
    QByteArray clientChallenge(8, 0);

    for (uint i = 0; i < 8; i++) {
        clientChallenge.data() [i] = KRandom::random() % 0xff;
    }

    return lmv2Response(hash, clientChallenge, challenge);
}

QByteArray KNTLM::ntlmv2Hash(const QString &target, const QString &user, const QString &password)
{
    const QByteArray hash = ntlmHash(password);
    const QByteArray key = QString2UnicodeLE(user.toUpper() + target);
    return hmacMD5(key, hash);
}

QByteArray KNTLM::lmv2Response(const QByteArray &hash,
                               const QByteArray &clientData, const unsigned char *challenge)
{
    QByteArray data(8 + clientData.size(), 0);
    memcpy(data.data(), challenge, 8);
    memcpy(data.data() + 8, clientData.data(), clientData.size());

    QByteArray mac = hmacMD5(data, hash);
    mac.resize(16 + clientData.size());
    memcpy(mac.data() + 16, clientData.data(), clientData.size());
    return mac;
}
