/*
  fortuna.c
  Fortuna cryptographic PRNG
  As seen in Ferguson and Schneier "Practical Cryptography"
  Greg Parker     gparker@cs.stanford.edu     2003-11-16
*/
/**********
 * Copyright (c) 2003-2004 Greg Parker.  All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY GREG PARKER ``AS IS'' AND ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
 * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
 * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 **********/

#include "includes.h"
#include "assert.h"
#include "fortuna.h"
#include "fortuna_locl.h"


static void ctr_inc(uint8_t *ctr) FORTUNA_SEGMENT;
static void generator_set_key(generator_t *g, uint8_t *newkey) FORTUNA_SEGMENT;
static void generator_emit_blocks(generator_t *g, uint8_t *outData, unsigned int blockCount) FORTUNA_SEGMENT;


static void ctr_inc(uint8_t *ctr)
{
    int carry = 0;
    uint32_t *cp = (uint32_t *)ctr;

    // Process in 4 32-byte chunks
#define inc(c) \
    if (c == 0xffffffffUL) { \
        c = carry; \
        carry = 1; \
    } else { \
        c += 1; \
    } \

    inc(cp[0]);
    inc(cp[1]);
    inc(cp[2]);
    inc(cp[3]);
}


void generator_init(generator_t *g)
{
    memset(g, 0, sizeof(*g));
}


static void generator_set_key(generator_t *g, uint8_t *newkey)
{
    fortuna_log("gen_set_key");
    memcpy(g->key, newkey, KEY_BYTES);
    AES_set_encrypt_key(g->key, 128, &g->aes_key);
}

void generator_reseed(generator_t *g, uint8_t *data, unsigned int len)
{
    SHA256_CTX ctx;
    uint8_t newkey[KEY_BYTES];

    // Reseed: 
    // g->key = SHA256(SHA256(g->key || data))
    // g->ctr++

    // First round: newkey = SHA256(g->key || data)
    SHA256_Init(&ctx);
    SHA256_Update(&ctx, g->key, KEY_BYTES);
    SHA256_Update(&ctx, data, len);
    SHA256_Final(newkey, &ctx);

    // Second round: newkey = SHA256(newkey)
    SHA256_Init(&ctx);
    SHA256_Update(&ctx, newkey, KEY_BYTES);
    SHA256_Final(newkey, &ctx);

    // Install the new key
    generator_set_key(g, newkey);

    // Increment counter
    ctr_inc(g->ctr);
}


// outData may overlap g->key
// outData may NOT overlap g->ctr
static void generator_emit_blocks(generator_t *g, 
                                  uint8_t *outData, unsigned int blockCount)
{
    unsigned int i;
    uint8_t *dst = outData;

    // Repeatedly encrypt plaintext g->ctr with key g->key (aka g->aes_key), 
    // incrementing the counter for each block

    fortuna_log("gen_emit_blocks (%d)", blockCount);

    for (i = 0; i < blockCount; i++) {
        AES_encrypt(g->ctr, dst, &g->aes_key);
        ctr_inc(g->ctr);
        dst += BLOCK_BYTES;
    }
}


void generator_emit_bytes(generator_t *g, uint8_t *outBytes, unsigned int byteCount)
{
    // round up to block boundary
    unsigned int wholeBlockCount;
    unsigned int incompleteBlockBytes;
    uint8_t blockbuf[KEY_BYTES];
    uint8_t *dst = outBytes;

    fortuna_log("gen_emit_bytes (%d)", byteCount);

    if (byteCount == 0) return;
    
    // F&S use 2^20 upper bound for cryptographic reasons, 
    // though anything over a few KB is probably insane for pssh
    assert(byteCount < 1048576); // 2^20

    incompleteBlockBytes = byteCount % BLOCK_BYTES;
    wholeBlockCount = byteCount / BLOCK_BYTES;

    // Generate all whole blocks (i.e. everything but maybe the last)
    generator_emit_blocks(g, dst, wholeBlockCount);
    dst += wholeBlockCount * BLOCK_BYTES;

    // Generate another block if we need a partial block at the end
    if (incompleteBlockBytes > 0) {
        generator_emit_blocks(g, blockbuf, 1);
        memcpy(dst, blockbuf, incompleteBlockBytes);
    }

    // Generate two blocks for the new key
    generator_emit_blocks(g, blockbuf, 2);
    generator_set_key(g, blockbuf);
}






