/*
 * This file is part of the Flowee project
 * Copyright (c) 2008 Damien Bergamini <damien.bergamini@free.fr>
 * Copyright (C) 2021 Tom Zander <tom@flowee.org>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program 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 General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
#include "pbkdf512.h"

#include <string.h>
#include "hmac_sha512.h"

int pbkdf512(const uint8_t* passphrase, size_t passphrase_length, const uint8_t* salt, size_t salt_length, uint8_t* key, size_t key_length, size_t iterations)
{
    if (salt_length > SIZE_MAX - 4)
        return -1;

    uint8_t buffer[CHMAC_SHA512::OUTPUT_SIZE];
    memset(buffer, 0, sizeof(buffer));
    uint8_t digest1[CHMAC_SHA512::OUTPUT_SIZE];
    memset(digest1, 0, sizeof(digest1));
    uint8_t digest2[CHMAC_SHA512::OUTPUT_SIZE];
    memset(digest2, 0, sizeof(digest2));

    /* An iteration count of 0 is equivalent to a count of 1. */
    /* A key_length of 0 is a no-op. */
    /* A salt_length of 0 is perfectly valid. */

    const size_t asalt_size = salt_length + 4;
    uint8_t *asalt = (unsigned char*) malloc(asalt_size);
    if (asalt == nullptr)
        return -1;

    memcpy(asalt, salt, salt_length);
    for (size_t count = 1; key_length > 0; count++) {
        asalt[salt_length + 0] = (count >> 24) & 0xff;
        asalt[salt_length + 1] = (count >> 16) & 0xff;
        asalt[salt_length + 2] = (count >> 8) & 0xff;
        asalt[salt_length + 3] = (count >> 0) & 0xff;

        CHMAC_SHA512 context(passphrase, passphrase_length);
        context.Write(asalt, asalt_size);
        context.Finalize(digest1);
        memcpy(buffer, digest1, sizeof(buffer));

        for (size_t iteration = 1; iteration < iterations; iteration++) {
            CHMAC_SHA512 hash2(passphrase, passphrase_length);
            hash2.Write(digest1, sizeof(digest1));
            hash2.Finalize(digest2);
            memcpy(digest1, digest2, sizeof(digest1));
            for (size_t index = 0; index < sizeof(buffer); index++)
                buffer[index] ^= digest1[index];
        }

        size_t length = (key_length < sizeof(buffer) ? key_length : sizeof(buffer));
        memcpy(key, buffer, length);
        key += length;
        key_length -= length;
    }

    memset(asalt, 0, asalt_size);
    free(asalt);
    return 0;
}