/*
 * Copyright (C) 2003-2007  Justin Karneges <justin@affinix.com>
 * Copyright (C) 2004,2005  Brad Hards <bradh@frogmouth.net>
 *
 * 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 WITHOUT ANY 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, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
 * 02110-1301  USA
 *
 */

#include "qca_tools.h"

#include "qdebug.h"

#ifdef Q_OS_UNIX
#include <cstdlib>
#include <sys/mman.h>
#endif
#include "botantools/botantools.h"

namespace QCA {

static bool can_lock()
{
#ifdef Q_OS_UNIX
    bool ok = false;
#ifdef MLOCK_NOT_VOID_PTR
#define MLOCK_TYPE char *
#define MLOCK_TYPE_CAST (MLOCK_TYPE)
#else
#define MLOCK_TYPE void *
#define MLOCK_TYPE_CAST
#endif

    MLOCK_TYPE d = MLOCK_TYPE_CAST malloc(256);
    if (mlock(d, 256) == 0) {
        munlock(d, 256);
        ok = true;
    }
    free(d);
    return ok;
#else
    return true;
#endif
}

// Botan shouldn't throw any exceptions in our init/deinit.

static Botan::Allocator *alloc = nullptr;

void botan_throw_abort()
{
    fprintf(stderr, "QCA: Exception from internal Botan\n");
    abort();
}

bool botan_init(int prealloc, bool mmap)
{
    // 64k minimum
    if (prealloc < 64)
        prealloc = 64;

    bool secmem = false;

    try {
        Botan::Builtin_Modules modules;
        Botan::Library_State  *libstate = new Botan::Library_State(modules.mutex_factory());
        libstate->prealloc_size         = prealloc * 1024;
        Botan::set_global_state(libstate);
        Botan::global_state().load(modules);

        if (can_lock()) {
            Botan::global_state().set_default_allocator("locking");
            secmem = true;
        } else if (mmap) {
            Botan::global_state().set_default_allocator("mmap");
            secmem = true;
        }
        alloc = Botan::Allocator::get(true);
    } catch (std::exception &) {
        fprintf(stderr, "QCA: Error initializing internal Botan\n");
        abort();
    }

    return secmem;
}

void botan_deinit()
{
    try {
        alloc = nullptr;
        Botan::set_global_state(nullptr);
    } catch (std::exception &) {
        botan_throw_abort();
    }
}

void *botan_secure_alloc(int bytes)
{
    try {
        return alloc->allocate((Botan::u32bit)bytes);
    } catch (std::exception &) {
        botan_throw_abort();
    }
    return nullptr; // never get here
}

void botan_secure_free(void *p, int bytes)
{
    try {
        alloc->deallocate(p, (Botan::u32bit)bytes);
    } catch (std::exception &) {
        botan_throw_abort();
    }
}

} // end namespace QCA

void *qca_secure_alloc(int bytes)
{
    // allocate enough room to store a size value in front, return a pointer after it
    char *c                       = (char *)QCA::botan_secure_alloc(bytes + sizeof(int));
    reinterpret_cast<int *>(c)[0] = bytes + sizeof(int);
    return c + sizeof(int);
}

void qca_secure_free(void *p)
{
    // backtrack to read the size value
    char *c = (char *)p;
    c -= sizeof(int);
    const int bytes = reinterpret_cast<int *>(c)[0];
    QCA::botan_secure_free(c, bytes);
}

void *qca_secure_realloc(void *p, int bytes)
{
    // if null, do a plain alloc (just like how realloc() works)
    if (!p)
        return qca_secure_alloc(bytes);

    // backtrack to read the size value
    char *c = (char *)p;
    c -= sizeof(int);
    const int oldsize = reinterpret_cast<int *>(c)[0] - sizeof(int);

    // alloc the new chunk
    char *new_p = (char *)qca_secure_alloc(bytes);
    if (!new_p)
        return nullptr;

    // move over the memory from the original block
    memmove(new_p, p, qMin(oldsize, bytes));

    // free the original
    qca_secure_free(p);

    // done
    return new_p;
}

namespace QCA {

// secure or non-secure buffer, with trailing 0-byte.
// buffer size of 0 is okay (sbuf/qbuf will be 0).
struct alloc_info
{
    bool  sec;
    char *data;
    int   size;

    // internal
    Botan::SecureVector<Botan::byte> *sbuf;
    QByteArray                       *qbuf;
};

// note: these functions don't return error if memory allocation/resizing
//   fails..  maybe fix this someday?

// ai: uninitialized
// size: >= 0
// note: memory will be initially zero'd out
static bool ai_new(alloc_info *ai, int size, bool sec);

// ai: uninitialized
// from: initialized
static bool ai_copy(alloc_info *ai, const alloc_info *from);

// ai: initialized
// new_size: >= 0
static bool ai_resize(alloc_info *ai, int new_size);

// ai: initialized
static void ai_delete(alloc_info *ai);

bool ai_new(alloc_info *ai, int size, bool sec)
{
    if (size < 0)
        return false;

    ai->size = size;
    ai->sec  = sec;

    if (size == 0) {
        ai->sbuf = nullptr;
        ai->qbuf = nullptr;
        ai->data = nullptr;
        return true;
    }

    if (sec) {
        try {
            ai->sbuf = new Botan::SecureVector<Botan::byte>((Botan::u32bit)size + 1);
        } catch (std::exception &) {
            botan_throw_abort();
            return false; // never get here
        }

        (*(ai->sbuf))[size] = 0;
        ai->qbuf            = nullptr;
        Botan::byte *bp     = (Botan::byte *)(*(ai->sbuf));
        ai->data            = (char *)bp;
    } else {
        ai->sbuf = nullptr;
        ai->qbuf = new QByteArray(size, 0);
        ai->data = ai->qbuf->data();
    }

    return true;
}

bool ai_copy(alloc_info *ai, const alloc_info *from)
{
    ai->size = from->size;
    ai->sec  = from->sec;

    if (ai->size == 0) {
        ai->sbuf = nullptr;
        ai->qbuf = nullptr;
        ai->data = nullptr;
        return true;
    }

    if (ai->sec) {
        try {
            ai->sbuf = new Botan::SecureVector<Botan::byte>(*(from->sbuf));
        } catch (std::exception &) {
            botan_throw_abort();
            return false; // never get here
        }

        ai->qbuf        = nullptr;
        Botan::byte *bp = (Botan::byte *)(*(ai->sbuf));
        ai->data        = (char *)bp;
    } else {
        ai->sbuf = nullptr;
        ai->qbuf = new QByteArray(*(from->qbuf));
        ai->data = ai->qbuf->data();
    }

    return true;
}

bool ai_resize(alloc_info *ai, int new_size)
{
    if (new_size < 0)
        return false;

    // new size is empty
    if (new_size == 0) {
        // we currently aren't empty
        if (ai->size > 0) {
            if (ai->sec) {
                delete ai->sbuf;
                ai->sbuf = nullptr;
            } else {
                delete ai->qbuf;
                ai->qbuf = nullptr;
            }

            ai->size = 0;
            ai->data = nullptr;
        }

        return true;
    }

    if (ai->sec) {
        Botan::SecureVector<Botan::byte> *new_buf;
        try {
            new_buf = new Botan::SecureVector<Botan::byte>((Botan::u32bit)new_size + 1);
        } catch (std::exception &) {
            botan_throw_abort();
            return false; // never get here
        }

        Botan::byte *new_p = (Botan::byte *)(*new_buf);
        if (ai->size > 0) {
            const Botan::byte *old_p = (const Botan::byte *)(*(ai->sbuf));
            memcpy(new_p, old_p, qMin(new_size, ai->size));
            delete ai->sbuf;
        }
        ai->sbuf                = new_buf;
        ai->size                = new_size;
        (*(ai->sbuf))[new_size] = 0;
        ai->data                = (char *)new_p;
    } else {
        if (ai->size > 0)
            ai->qbuf->resize(new_size);
        else
            ai->qbuf = new QByteArray(new_size, 0);

        ai->size = new_size;
        ai->data = ai->qbuf->data();
    }

    return true;
}

void ai_delete(alloc_info *ai)
{
    if (ai->size > 0) {
        if (ai->sec)
            delete ai->sbuf;
        else
            delete ai->qbuf;
    }
}

//----------------------------------------------------------------------------
// MemoryRegion
//----------------------------------------------------------------------------
static char blank[] = "";

class MemoryRegion::Private : public QSharedData
{
public:
    alloc_info ai;

    Private(int size, bool sec)
    {
        ai_new(&ai, size, sec);
    }

    Private(const QByteArray &from, bool sec)
    {
        ai_new(&ai, from.size(), sec);
        memcpy(ai.data, from.data(), ai.size);
    }

    Private(const Private &from)
        : QSharedData(from)
    {
        ai_copy(&ai, &from.ai);
    }

    ~Private()
    {
        ai_delete(&ai);
    }

    bool resize(int new_size)
    {
        return ai_resize(&ai, new_size);
    }

    void setSecure(bool sec)
    {
        // if same mode, do nothing
        if (ai.sec == sec)
            return;

        alloc_info other;
        ai_new(&other, ai.size, sec);
        memcpy(other.data, ai.data, ai.size);
        ai_delete(&ai);
        ai = other;
    }
};

MemoryRegion::MemoryRegion()
    : _secure(false)
    , d(nullptr)
{
}

MemoryRegion::MemoryRegion(const char *str)
    : _secure(false)
    , d(new Private(QByteArray::fromRawData(str, strlen(str)), false))
{
}

MemoryRegion::MemoryRegion(const QByteArray &from)
    : _secure(false)
    , d(new Private(from, false))
{
}

MemoryRegion::MemoryRegion(const MemoryRegion &from)
    : _secure(from._secure)
    , d(from.d)
{
}

MemoryRegion::~MemoryRegion()
{
}

MemoryRegion &MemoryRegion::operator=(const MemoryRegion &from)
{
    _secure = from._secure;
    d       = from.d;
    return *this;
}

MemoryRegion &MemoryRegion::operator=(const QByteArray &from)
{
    set(from, false);
    return *this;
}

bool MemoryRegion::isNull() const
{
    return (d ? false : true);
}

bool MemoryRegion::isSecure() const
{
    return _secure;
}

QByteArray MemoryRegion::toByteArray() const
{
    if (!d)
        return QByteArray();

    if (d->ai.sec) {
        QByteArray buf(d->ai.size, 0);
        memcpy(buf.data(), d->ai.data, d->ai.size);
        return buf;
    } else {
        if (d->ai.size > 0)
            return *(d->ai.qbuf);
        else
            return QByteArray((int)0, (char)0);
    }
}

MemoryRegion::MemoryRegion(bool secure)
    : _secure(secure)
    , d(nullptr)
{
}

MemoryRegion::MemoryRegion(int size, bool secure)
    : _secure(secure)
    , d(new Private(size, secure))
{
}

MemoryRegion::MemoryRegion(const QByteArray &from, bool secure)
    : _secure(secure)
    , d(new Private(from, secure))
{
}

char *MemoryRegion::data()
{
    if (!d)
        return blank;
    return d->ai.data;
}

const char *MemoryRegion::data() const
{
    if (!d)
        return blank;
    return d->ai.data;
}

const char *MemoryRegion::constData() const
{
    if (!d)
        return blank;
    return d->ai.data;
}

char &MemoryRegion::at(int index)
{
    return *(d->ai.data + index);
}

const char &MemoryRegion::at(int index) const
{
    return *(d->ai.data + index);
}

int MemoryRegion::size() const
{
    if (!d)
        return 0;
    return d->ai.size;
}

bool MemoryRegion::isEmpty() const
{
    if (!d)
        return true;
    return (d->ai.size > 0 ? false : true);
}

bool MemoryRegion::resize(int size)
{
    if (!d) {
        d = new Private(size, _secure);
        return true;
    }

    if (d->ai.size == size)
        return true;

    return d->resize(size);
}

void MemoryRegion::set(const QByteArray &from, bool secure)
{
    _secure = secure;

    if (!from.isEmpty())
        d = new Private(from, secure);
    else
        d = new Private(0, secure);
}

void MemoryRegion::setSecure(bool secure)
{
    _secure = secure;

    if (!d) {
        d = new Private(0, secure);
        return;
    }

    d->setSecure(secure);
}

//----------------------------------------------------------------------------
// SecureArray
//----------------------------------------------------------------------------
SecureArray::SecureArray()
    : MemoryRegion(true)
{
}

SecureArray::SecureArray(int size, char ch)
    : MemoryRegion(size, true)
{
    // ai_new fills with zeros for us
    if (ch != 0)
        fill(ch, size);
}

SecureArray::SecureArray(const char *str)
    : MemoryRegion(QByteArray::fromRawData(str, strlen(str)), true)
{
}

SecureArray::SecureArray(const QByteArray &a)
    : MemoryRegion(a, true)
{
}

SecureArray::SecureArray(const MemoryRegion &a)
    : MemoryRegion(a)
{
    setSecure(true);
}

SecureArray::SecureArray(const SecureArray &from)
    : MemoryRegion(from)
{
}

SecureArray::~SecureArray()
{
}

SecureArray &SecureArray::operator=(const SecureArray &from)
{
    MemoryRegion::operator=(from);
    return *this;
}

SecureArray &SecureArray::operator=(const QByteArray &from)
{
    MemoryRegion::set(from, true);
    return *this;
}

void SecureArray::clear()
{
    MemoryRegion::resize(0);
}

bool SecureArray::resize(int size)
{
    return MemoryRegion::resize(size);
}

char &SecureArray::operator[](int index)
{
    return at(index);
}

const char &SecureArray::operator[](int index) const
{
    return at(index);
}

char &SecureArray::at(int index)
{
    return MemoryRegion::at(index);
}

const char &SecureArray::at(int index) const
{
    return MemoryRegion::at(index);
}

char *SecureArray::data()
{
    return MemoryRegion::data();
}

const char *SecureArray::data() const
{
    return MemoryRegion::data();
}

const char *SecureArray::constData() const
{
    return MemoryRegion::constData();
}

int SecureArray::size() const
{
    return MemoryRegion::size();
}

bool SecureArray::isEmpty() const
{
    return MemoryRegion::isEmpty();
}

QByteArray SecureArray::toByteArray() const
{
    return MemoryRegion::toByteArray();
}

SecureArray &SecureArray::append(const SecureArray &a)
{
    const int oldsize = size();
    resize(oldsize + a.size());
    memcpy(data() + oldsize, a.data(), a.size());
    return *this;
}

bool SecureArray::operator==(const MemoryRegion &other) const
{
    if (this == &other)
        return true;
    if (size() == other.size() && memcmp(data(), other.data(), size()) == 0)
        return true;
    return false;
}

SecureArray &SecureArray::operator+=(const SecureArray &a)
{
    return append(a);
}

void SecureArray::fill(char fillChar, int fillToPosition)
{
    const int len = (fillToPosition == -1) ? size() : qMin(fillToPosition, size());
    if (len > 0)
        memset(data(), (int)fillChar, len);
}

void SecureArray::set(const SecureArray &from)
{
    *this = from;
}

void SecureArray::set(const QByteArray &from)
{
    *this = from;
}

const SecureArray operator+(const SecureArray &a, const SecureArray &b)
{
    SecureArray c = a;
    return c.append(b);
}

//----------------------------------------------------------------------------
// BigInteger
//----------------------------------------------------------------------------
static void negate_binary(char *a, int size)
{
    // negate = two's compliment + 1
    bool done = false;
    for (int n = size - 1; n >= 0; --n) {
        a[n] = ~a[n];
        if (!done) {
            if ((unsigned char)a[n] < 0xff) {
                ++a[n];
                done = true;
            } else
                a[n] = 0;
        }
    }
}

class BigInteger::Private : public QSharedData
{
public:
    Botan::BigInt n;
};

BigInteger::BigInteger()
{
    d = new Private;
}

BigInteger::BigInteger(int i)
{
    d = new Private;
    if (i < 0) {
        d->n = Botan::BigInt(i * (-1));
        d->n.set_sign(Botan::BigInt::Negative);
    } else {
        d->n = Botan::BigInt(i);
        d->n.set_sign(Botan::BigInt::Positive);
    }
}

BigInteger::BigInteger(const char *c)
{
    d = new Private;
    fromString(QString::fromLatin1(c));
}

BigInteger::BigInteger(const QString &s)
{
    d = new Private;
    fromString(s);
}

BigInteger::BigInteger(const SecureArray &a)
{
    d = new Private;
    fromArray(a);
}

BigInteger::BigInteger(const BigInteger &from)
{
    *this = from;
}

BigInteger::~BigInteger()
{
}

BigInteger &BigInteger::operator=(const BigInteger &from)
{
    d = from.d;
    return *this;
}

BigInteger &BigInteger::operator+=(const BigInteger &i)
{
    d->n += i.d->n;
    return *this;
}

BigInteger &BigInteger::operator-=(const BigInteger &i)
{
    d->n -= i.d->n;
    return *this;
}

BigInteger &BigInteger::operator*=(const BigInteger &i)
{
    d->n *= i.d->n;
    return *this;
}

BigInteger &BigInteger::operator/=(const BigInteger &i)
{
    try {
        d->n /= i.d->n;
    } catch (std::exception &) {
        fprintf(stderr, "QCA: Botan integer division error\n");
        abort();
    }
    return *this;
}

BigInteger &BigInteger::operator%=(const BigInteger &i)
{
    try {
        d->n %= i.d->n;
    } catch (std::exception &) {
        fprintf(stderr, "QCA: Botan integer division error\n");
        abort();
    }
    return *this;
}

BigInteger &BigInteger::operator=(const QString &s)
{
    fromString(s);
    return *this;
}

int BigInteger::compare(const BigInteger &n) const
{
    return ((d->n).cmp(n.d->n, true));
}

QTextStream &operator<<(QTextStream &stream, const BigInteger &b)
{
    stream << b.toString();
    return stream;
}

SecureArray BigInteger::toArray() const
{
    int size = d->n.encoded_size(Botan::BigInt::Binary);

    // return at least 8 bits
    if (size == 0) {
        SecureArray a(1);
        a[0] = 0;
        return a;
    }

    int         offset = 0;
    SecureArray a;

    // make room for a sign bit if needed
    if (d->n.get_bit((size * 8) - 1)) {
        ++size;
        a.resize(size);
        a[0] = 0;
        ++offset;
    } else
        a.resize(size);

    Botan::BigInt::encode((Botan::byte *)a.data() + offset, d->n, Botan::BigInt::Binary);

    if (d->n.is_negative())
        negate_binary(a.data(), a.size());

    return a;
}

void BigInteger::fromArray(const SecureArray &_a)
{
    if (_a.isEmpty()) {
        d->n = Botan::BigInt(0);
        return;
    }
    SecureArray a = _a;

    Botan::BigInt::Sign sign = Botan::BigInt::Positive;
    if (a[0] & 0x80)
        sign = Botan::BigInt::Negative;

    if (sign == Botan::BigInt::Negative)
        negate_binary(a.data(), a.size());

    d->n = Botan::BigInt::decode((const Botan::byte *)a.data(), a.size(), Botan::BigInt::Binary);
    d->n.set_sign(sign);
}

QString BigInteger::toString() const
{
    QByteArray cs;
    try {
        cs.resize(d->n.encoded_size(Botan::BigInt::Decimal));
        Botan::BigInt::encode((Botan::byte *)cs.data(), d->n, Botan::BigInt::Decimal);
    } catch (std::exception &) {
        return QString();
    }

    QString str;
    if (d->n.is_negative())
        str += QLatin1Char('-');
    str += QString::fromLatin1(cs);
    str.remove(QChar::Null);
    return str;
}

bool BigInteger::fromString(const QString &s)
{
    if (s.isEmpty())
        return false;
    const QByteArray cs = s.toLatin1();

    bool neg = false;
    if (s[0] == QLatin1Char('-'))
        neg = true;

    try {
        d->n = Botan::BigInt::decode(
            (const Botan::byte *)cs.data() + (neg ? 1 : 0), cs.length() - (neg ? 1 : 0), Botan::BigInt::Decimal);
    } catch (std::exception &) {
        return false;
    }

    if (neg)
        d->n.set_sign(Botan::BigInt::Negative);
    else
        d->n.set_sign(Botan::BigInt::Positive);
    return true;
}

}