/*
 * SPDX-License-Identifier: GPL-3.0-or-later
 * SPDX-FileCopyrightText: 2020 Johan Ouwerkerk <jm.ouwerkerk@gmail.com>
 */
#include "oath.h"

#include "../hmac/hmac.h"
#include "../logging_p.h"

#include <limits>

KEYSMITH_LOGGER(logger, ".oath")

static qint64 maxMSecsOffset = std::numeric_limits<qint64>::max();

static QString encodeDefaults(quint32 value, uint tokenLength)
{
    Q_ASSERT_X(tokenLength >= 6, Q_FUNC_INFO, "token length should be at least 6 characters long");

    QString base;
    base.setNum(value, 10);

    QString prefix(QLatin1String(""));
    for (uint i = base.size(); i < tokenLength; ++i) {
        prefix += QLatin1Char('0');
    }

    return prefix + base;
}

static QString encodeDefaultsWithChecksum(quint32 value, uint tokenLength)
{
    QString prefix = encodeDefaults(value, tokenLength);

    QString check;
    check.setNum(oath::luhnChecksum(value, tokenLength), 10);

    return prefix + check;
}

static quint32 truncate(const QByteArray &hash, uint offset)
{
    Q_ASSERT_X(hash.size() >= 4, Q_FUNC_INFO, "hash output is too small");
    Q_ASSERT_X(offset <= (((uint) hash.size()) - 4UL), Q_FUNC_INFO, "truncation offset is too large for the hash output");

    return ((((quint32) hash[offset]) & 0x7FUL) << 24)
        | ((((quint32) hash[offset + 1]) & 0xFFUL) << 16)
        | ((((quint32) hash[offset + 2]) & 0xFFUL) << 8)
        | (((quint32) hash[offset + 3]) & 0xFFUL);
}

static quint32 truncateDynamically(const QByteArray &hash)
{
    Q_ASSERT_X(hash.size() >= 20, Q_FUNC_INFO, "hash output is too small");
    return truncate(hash, ((uint) hash[hash.size() - 1]) & 0x0FUL);
}

namespace oath
{
    Encoder::Encoder(uint tokenLength, bool addChecksum) : m_tokenLength(tokenLength), m_addChecksum(addChecksum)
    {
    }

    Encoder::~Encoder()
    {
    }

    quint32 Encoder::reduceMod10(quint32 value, uint tokenLength)
    {
        /*
         * Skip modulo 10 reduction for tokens of 10 or more characters:
         * the value is already guaranteed to be in its modulo 10 reduced form, because 2^32 is less than 10^10.
         * This check also takes care of possible integer overflow, for the same reason.
         */
        return tokenLength <= 9 ? value % powerTable[tokenLength] : value;
    }

    QString Encoder::encode(quint32 value) const
    {
        value = reduceMod10(value, m_tokenLength);
        return m_addChecksum ? encodeDefaultsWithChecksum(value, m_tokenLength) : encodeDefaults(value, m_tokenLength);
    }

    uint Encoder::tokenLength(void) const
    {
        return m_tokenLength;
    }

    bool Encoder::checksum(void) const
    {
        return m_addChecksum;
    }

    bool Algorithm::validate(const Encoder *encoder)
    {
        // HOTP spec mandates a minimum token length of 6 digits
        return encoder && encoder->tokenLength() >= 6;
    }

    bool Algorithm::validate(QCryptographicHash::Algorithm algorithm, const std::optional<uint> offset)
    {
        /*
         * An nullopt offset indicates dynamic truncation.
         * Dynamic truncation works by taking the last nible and interpreting it as offset for truncation, i.e. it will always be <= 15.
         * Accounting for the last nibble (therefore last byte) assume a max truncation offset of 16 if dynamic truncation is used.
         */
        uint truncateAt = offset ? *offset : 16U;

        /*
         * The given algorithm must be supported/have a known digest size.
         * There must be at least 4 bytes available at the given truncation offset/limit.
         */
        std::optional<uint> digestSize = hmac::outputSize(algorithm);
        return digestSize && *digestSize >= 4U && (*digestSize - 4U) >= truncateAt;
    }

    std::optional<Algorithm> Algorithm::create(QCryptographicHash::Algorithm algorithm, const std::optional<uint> offset, const QSharedPointer<const Encoder> &encoder, bool requireSaneKeyLength)
    {
        if(!validate(algorithm, offset)) {
            qCDebug(logger) << "Invalid algorithm:" << algorithm << "or incompatible with truncation offset:" << (offset ? *offset : 16U);
            return std::nullopt;
        }

        if (!encoder || !validate(encoder.data())) {
            qCDebug(logger) << "Invalid token encoder";
            return std::nullopt;
        }

        std::function<quint32(const QByteArray &)> truncation(truncateDynamically);
        if (offset) {
            uint at = *offset;
            truncation = [at](const QByteArray &bytes) -> quint32
            {
                return truncate(bytes, at);
            };
        }

        return std::optional<Algorithm>(Algorithm(encoder, truncation, algorithm, requireSaneKeyLength));
    }

    std::optional<Algorithm> Algorithm::totp(QCryptographicHash::Algorithm algorithm, uint tokenLength, bool requireSaneKeyLength)
    {
        const QSharedPointer<const Encoder> encoder(new Encoder(tokenLength, false));
        return create(algorithm, std::nullopt, encoder, requireSaneKeyLength);
    }

    std::optional<Algorithm> Algorithm::hotp(const std::optional<uint> offset, uint tokenLength, bool checksum, bool requireSaneKeyLength)
    {
        const QSharedPointer<const Encoder> encoder(new Encoder(tokenLength, checksum));
        return create(QCryptographicHash::Sha1, offset, encoder, requireSaneKeyLength);
    }

    Algorithm::Algorithm(const QSharedPointer<const Encoder> &encoder, const std::function<quint32(const QByteArray &)> &truncation, QCryptographicHash::Algorithm algorithm, bool requireSaneKeyLength) :
        m_encoder(encoder), m_truncation(truncation), m_enforceKeyLength(requireSaneKeyLength), m_algorithm(algorithm)
    {
    }

    std::optional<QString> Algorithm::compute(quint64 counter, char * secretBuffer, int length) const
    {
        if (!secretBuffer) {
            return std::nullopt;
        }

        if (!hmac::validateKeySize(m_algorithm, length, m_enforceKeyLength)) {
            qCDebug(logger)
                << "Invalid key size:" << length << "for algorithm:" << m_algorithm
                << "Sane key length requirements apply:" << m_enforceKeyLength;
            return std::nullopt;
        }

        QByteArray message;
        message.resize(8);

        for (int i = 0; i < 8; ++i) {
            message[i] = (char) ((counter >> (56 - i * 8)) & 0xFFULL);
        }

        std::optional<QByteArray> digest = hmac::compute(m_algorithm, secretBuffer, length, message, m_enforceKeyLength);
        if (digest) {
            quint32 result = m_truncation(*digest);
            result = Encoder::reduceMod10(result, m_encoder->tokenLength());
            return std::optional<QString>(m_encoder->encode(result));
        }

        qCDebug(logger) << "Failed to compute token";
        return std::nullopt;
    }

    uint luhnChecksum(quint32 value, uint digits)
    {
        static const uint lookupTable[10] = {
            0, // 0 * 2
            2, // 1 * 2
            4, // 2 * 2
            6, // 3 * 2
            8, // 4 * 2
            1, // 5 * 2 - 9
            3, // 6 * 2 - 9
            5, // 7 * 2 - 9
            7, // 8 * 2 - 9
            9, // 9 * 2 - 9
        };

        Q_ASSERT_X(digits > 0UL, Q_FUNC_INFO, "checksum cannot be computed over less than 1 digit");
        uint sum = 0UL;
        bool doubledMinus9 = true;
        for (uint d = 0UL; d < digits && value != 0UL; ++d) {
            uint position = value % 10UL;

            sum += doubledMinus9 ? lookupTable[position] : position;

            value /= 10UL;
            doubledMinus9 = !doubledMinus9;
        }

        sum = sum % 10ULL;
        return sum == 0UL ? 0UL : 10UL - sum;
    }

    std::optional<quint64> count(const QDateTime &epoch, uint timeStep, const std::function<qint64(void)> &clock)
    {
        qint64 epochMillis = epoch.toMSecsSinceEpoch();
        qint64 now = clock();

        if (now < epochMillis) {
            qCDebug(logger) << "Unable to count time steps: epoch is in the future";
            return std::nullopt;
        }

        if (timeStep == 0UL) {
            qCDebug(logger) << "Unable to count time steps: invalid step size:" << timeStep;
            return std::nullopt;
        }

        quint64 msecs = ((quint64) (now - epochMillis));
        quint64 stepInMsecs = ((quint64) timeStep) * 1000ULL;
        return std::optional<quint64>(msecs / stepInMsecs);
    }

    /*
     * Converts a negative qint64 value to its absolute value equivalent in quint64.
     */
    static quint64 flipSign(qint64 value)
    {
        static const quint64 max = std::numeric_limits<quint64>::max();
        // take advantage of two's complement to simplify this
        return max - ((quint64) value) + 1ULL;
    }

    std::optional<QDateTime> fromCounter(quint64 count, const QDateTime &epoch, uint timeStep)
    {
        qint64 epochMillis = epoch.toMSecsSinceEpoch();

        /*
         * Calculate the number of milliseconds that would be available for the given token.
         */
        quint64 max = epochMillis >= 0
            ? (quint64) (maxMSecsOffset - epochMillis)
            : ((quint64) maxMSecsOffset) + flipSign(epochMillis);

        quint64 step = timeStep * 1000ULL;

        // see if the requested count of time steps 'fits' inside the number of available milliseconds
        if ((max / step) < count) {
            qCDebug(logger)
                << "Unable to compute datetime matching the given count of time steps:"
                << "Storage type not wide enough, not enough milliseconds available";
            return std::nullopt;
        }

        quint64 ms = count * step;
        qint64 offset = epochMillis;
        if (ms <= ((quint64) maxMSecsOffset)) {
            offset += (qint64) ms;
        } else {
            /*
             * This is safe to do because:
             * - it has been verified that the number of requested steps 'fits' within the number of available ms
             * - therefore the epoch must have been negative
             */
            offset += maxMSecsOffset;
            ms -= (quint64) maxMSecsOffset;
            offset += (qint64) ms;
        }

        /*
         * QDateTime::fromMSecsSinceEpoch() is documented that it cannot handle the full qint64 width, but it is not
         * documented what exactly the restrictions are. Implement a sanity check to detect and recover from confused
         * 'nonsense' answers.
         */
        auto v = QDateTime::fromMSecsSinceEpoch(offset);
        if (v.toMSecsSinceEpoch() != offset) {
            qCDebug(logger)
                << "Unable to compute datetime matching the given count of time steps:"
                << "Internal confusion in QDateTime detected, number of milliseconds is probably out of range";
            return std::nullopt;
        }

        return std::optional<QDateTime>(v);
    }
}
