/*
 * Copyright (c) 2011 Alex Hornung <alex@alexhornung.com>.
 * 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 THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * ``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
 * COPYRIGHT HOLDERS OR CONTRIBUTORS 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 <sys/types.h>

#if defined(__DragonFly__)
#include <sys/endian.h>
#elif defined(__linux__)
#include <endian.h>
#endif
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <inttypes.h>
#include <string.h>

#include "crc32.h"
#include "tcplay.h"

/* Endianess macros */
#define BE_TO_HOST(n, v) v = be ## n ## toh(v)
#define LE_TO_HOST(n, v) v = le ## n ## toh(v)
#define HOST_TO_BE(n, v) v = htobe ## n (v)
#define HOST_TO_LE(n, v) v = htole ## n (v)

struct sig_hdr_cfg {
        const char      *sig;
        uint16_t        min_ver;
};

struct sig_hdr_cfg sig_hdr_cfgs[] = {
        { TC_SIG,       0x0007 },
        { VC_SIG,       0x0b01 },
        { NULL,         0x0000 }
};

static
const
struct sig_hdr_cfg *hdr_cfg_from_sig(const char *sig)
{
        const struct sig_hdr_cfg *cfg;

        for (cfg = &sig_hdr_cfgs[0]; cfg->sig != NULL; cfg++) {
                if (strcmp(cfg->sig, sig) == 0)
                        return cfg;
        }

        return NULL;
}

struct tchdr_dec *
decrypt_hdr(struct tchdr_enc *ehdr, struct tc_cipher_chain *cipher_chain,
    unsigned char *key)
{
        struct tchdr_dec *dhdr;
        unsigned char iv[128];
        int error;

        if ((dhdr = alloc_safe_mem(sizeof(struct tchdr_dec))) == NULL) {
                tc_log(1, "Error allocating safe tchdr_dec memory\n");
                return NULL;
        }

        memset(iv, 0, sizeof(iv));

        error = tc_decrypt(cipher_chain, key, iv, ehdr->enc,
            sizeof(struct tchdr_dec), (unsigned char *)dhdr);
        if (error) {
                tc_log(1, "Header decryption failed\n");
                free_safe_mem(dhdr);
                return NULL;
        }

        BE_TO_HOST(16, dhdr->tc_ver);
        LE_TO_HOST(16, dhdr->tc_min_ver);
        BE_TO_HOST(32, dhdr->crc_keys);
        BE_TO_HOST(64, dhdr->vol_ctime);
        BE_TO_HOST(64, dhdr->hdr_ctime);
        BE_TO_HOST(64, dhdr->sz_hidvol);
        BE_TO_HOST(64, dhdr->sz_vol);
        BE_TO_HOST(64, dhdr->off_mk_scope);
        BE_TO_HOST(64, dhdr->sz_mk_scope);
        BE_TO_HOST(32, dhdr->flags);
        BE_TO_HOST(32, dhdr->sec_sz);
        BE_TO_HOST(32, dhdr->crc_dhdr);

        return dhdr;
}

int
verify_hdr(struct tchdr_dec *hdr, struct pbkdf_prf_algo *prf_algo)
{
        uint32_t crc;

        if (memcmp(hdr->tc_str, prf_algo->sig, sizeof(hdr->tc_str)) != 0) {
#ifdef DEBUG
                fprintf(stderr, "Signature mismatch\n");
#endif
                return 0;
        }

        crc = crc32((void *)&hdr->keys, 256);
        if (crc != hdr->crc_keys) {
#ifdef DEBUG
                fprintf(stderr, "CRC32 mismatch (crc_keys)\n");
#endif
                return 0;
        }

        switch(hdr->tc_ver) {
        case 1:
        case 2:
                /* Unsupported header version */
                tc_log(1, "Header version %d unsupported\n", hdr->tc_ver);
                return 0;

        case 3:
        case 4:
                hdr->sec_sz = 512;
                break;
        }

        return 1;
}

struct tchdr_enc *
create_hdr(unsigned char *pass, int passlen, struct pbkdf_prf_algo *prf_algo,
    struct tc_cipher_chain *cipher_chain, size_t sec_sz,
    disksz_t total_blocks __unused,
    off_t offset, disksz_t blocks, int hidden, int weak, struct tchdr_enc **backup_hdr)
{
        struct tchdr_enc *ehdr, *ehdr_backup;
        struct tchdr_dec *dhdr;
        unsigned char *key, *key_backup;
        unsigned char iv[128];
        const struct sig_hdr_cfg *hdr_cfg;
        int error;

        key = key_backup = NULL;
        dhdr = NULL;
        ehdr = ehdr_backup = NULL;

        if (backup_hdr != NULL)
                *backup_hdr = NULL;

        if ((dhdr = (struct tchdr_dec *)alloc_safe_mem(sizeof(*dhdr))) == NULL) {
                tc_log(1, "could not allocate safe dhdr memory\n");
                goto error;
        }

        if ((ehdr = (struct tchdr_enc *)alloc_safe_mem(sizeof(*ehdr))) == NULL) {
                tc_log(1, "could not allocate safe ehdr memory\n");
                goto error;
        }

        if ((ehdr_backup = (struct tchdr_enc *)alloc_safe_mem
            (sizeof(*ehdr_backup))) == NULL) {
                tc_log(1, "could not allocate safe ehdr_backup memory\n");
                goto error;
        }

        if ((key = alloc_safe_mem(MAX_KEYSZ)) == NULL) {
                tc_log(1, "could not allocate safe key memory\n");
                goto error;
        }

        if ((key_backup = alloc_safe_mem(MAX_KEYSZ)) == NULL) {
                tc_log(1, "could not allocate safe backup key memory\n");
                goto error;
        }

        if ((error = get_random(ehdr->salt, sizeof(ehdr->salt), weak)) != 0) {
                tc_log(1, "could not get salt\n");
                goto error;
        }

        if ((error = get_random(ehdr_backup->salt, sizeof(ehdr_backup->salt), weak))
            != 0) {
                tc_log(1, "could not get salt for backup header\n");
                goto error;
        }

        error = pbkdf2(prf_algo, (char *)pass, passlen,
            ehdr->salt, sizeof(ehdr->salt),
            MAX_KEYSZ, key);
        if (error) {
                tc_log(1, "could not derive key\n");
                goto error;
        }

        error = pbkdf2(prf_algo, (char *)pass, passlen,
            ehdr_backup->salt, sizeof(ehdr_backup->salt),
            MAX_KEYSZ, key_backup);
        if (error) {
                tc_log(1, "could not derive backup key\n");
                goto error;
        }

        memset(dhdr, 0, sizeof(*dhdr));

        if ((error = get_random(dhdr->keys, sizeof(dhdr->keys), weak)) != 0) {
                tc_log(1, "could not get key random bits\n");
                goto error;
        }

        if ((hdr_cfg = hdr_cfg_from_sig(prf_algo->sig)) == NULL) {
                tc_log(1, "could not find internal header configuration\n");
                goto error;
        }

        memcpy(dhdr->tc_str, prf_algo->sig, 4);
        dhdr->tc_ver = 5;
        dhdr->tc_min_ver = hdr_cfg->min_ver;
        dhdr->crc_keys = crc32((void *)&dhdr->keys, 256);
        dhdr->sz_vol = blocks * sec_sz;
        if (hidden)
                dhdr->sz_hidvol = dhdr->sz_vol;
        dhdr->off_mk_scope = offset * sec_sz;
        dhdr->sz_mk_scope = blocks * sec_sz;
        dhdr->sec_sz = sec_sz;
        dhdr->flags = 0;

        HOST_TO_BE(16, dhdr->tc_ver);
        HOST_TO_LE(16, dhdr->tc_min_ver);
        HOST_TO_BE(32, dhdr->crc_keys);
        HOST_TO_BE(64, dhdr->sz_vol);
        HOST_TO_BE(64, dhdr->sz_hidvol);
        HOST_TO_BE(64, dhdr->off_mk_scope);
        HOST_TO_BE(64, dhdr->sz_mk_scope);
        HOST_TO_BE(32, dhdr->sec_sz);
        HOST_TO_BE(32, dhdr->flags);

        dhdr->crc_dhdr = crc32((void *)dhdr, 188);
        HOST_TO_BE(32, dhdr->crc_dhdr);

        memset(iv, 0, sizeof(iv));
        error = tc_encrypt(cipher_chain, key, iv, (unsigned char *)dhdr,
            sizeof(struct tchdr_dec), ehdr->enc);
        if (error) {
                tc_log(1, "Header encryption failed\n");
                goto error;
        }

        memset(iv, 0, sizeof(iv));
        error = tc_encrypt(cipher_chain, key_backup, iv,
            (unsigned char *)dhdr,
            sizeof(struct tchdr_dec), ehdr_backup->enc);
        if (error) {
                tc_log(1, "Backup header encryption failed\n");
                goto error;
        }

        free_safe_mem(key);
        free_safe_mem(key_backup);
        free_safe_mem(dhdr);

        if (backup_hdr != NULL)
                *backup_hdr = ehdr_backup;
        else
                free_safe_mem(ehdr_backup);

        return ehdr;
        /* NOT REACHED */

error:
        if (key)
                free_safe_mem(key);
        if (key_backup)
                free_safe_mem(key_backup);
        if (dhdr)
                free_safe_mem(dhdr);
        if (ehdr)
                free_safe_mem(ehdr);
        if (ehdr_backup)
                free_safe_mem(ehdr_backup);

        return NULL;
}

struct tchdr_enc *copy_reencrypt_hdr(unsigned char *pass, int passlen,
    struct pbkdf_prf_algo *prf_algo, int weak, struct tcplay_info *info,
    struct tchdr_enc **backup_hdr)
{
        struct tchdr_enc *ehdr, *ehdr_backup;
        unsigned char *key, *key_backup;
        unsigned char iv[128];
        const struct sig_hdr_cfg *hdr_cfg;
        int error;

        key = key_backup = NULL;
        ehdr = ehdr_backup = NULL;

        /* By default stick to current PRF algo */
        if (prf_algo == NULL)
                prf_algo = info->pbkdf_prf;

        if ((ehdr = (struct tchdr_enc *)alloc_safe_mem(sizeof(*ehdr))) == NULL) {
                tc_log(1, "could not allocate safe ehdr memory\n");
                goto error;
        }

        if ((ehdr_backup = (struct tchdr_enc *)alloc_safe_mem
            (sizeof(*ehdr_backup))) == NULL) {
                tc_log(1, "could not allocate safe ehdr_backup memory\n");
                goto error;
        }

        if ((key = alloc_safe_mem(MAX_KEYSZ)) == NULL) {
                tc_log(1, "could not allocate safe key memory\n");
                goto error;
        }

        if ((key_backup = alloc_safe_mem(MAX_KEYSZ)) == NULL) {
                tc_log(1, "could not allocate safe backup key memory\n");
                goto error;
        }

        if ((error = get_random(ehdr->salt, sizeof(ehdr->salt), weak)) != 0) {
                tc_log(1, "could not get salt\n");
                goto error;
        }

        if ((error = get_random(ehdr_backup->salt, sizeof(ehdr_backup->salt), weak))
            != 0) {
                tc_log(1, "could not get salt for backup header\n");
                goto error;
        }

        error = pbkdf2(prf_algo, (char *)pass, passlen,
            ehdr->salt, sizeof(ehdr->salt),
            MAX_KEYSZ, key);
        if (error) {
                tc_log(1, "could not derive key\n");
                goto error;
        }

        error = pbkdf2(prf_algo, (char *)pass, passlen,
            ehdr_backup->salt, sizeof(ehdr_backup->salt),
            MAX_KEYSZ, key_backup);
        if (error) {
                tc_log(1, "could not derive backup key\n");
                goto error;
        }

        if ((hdr_cfg = hdr_cfg_from_sig(prf_algo->sig)) == NULL) {
                tc_log(1, "could not find internal header configuration\n");
                goto error;
        }

        /* Update signature and min_ver depending on selected PBKDF2 PRF algo */
        memcpy(info->hdr->tc_str, prf_algo->sig, 4);
        info->hdr->tc_min_ver = hdr_cfg->min_ver;

        HOST_TO_BE(16, info->hdr->tc_ver);
        HOST_TO_LE(16, info->hdr->tc_min_ver);
        HOST_TO_BE(32, info->hdr->crc_keys);
        HOST_TO_BE(64, info->hdr->vol_ctime);
        HOST_TO_BE(64, info->hdr->hdr_ctime);
        HOST_TO_BE(64, info->hdr->sz_vol);
        HOST_TO_BE(64, info->hdr->sz_hidvol);
        HOST_TO_BE(64, info->hdr->off_mk_scope);
        HOST_TO_BE(64, info->hdr->sz_mk_scope);
        HOST_TO_BE(32, info->hdr->sec_sz);
        HOST_TO_BE(32, info->hdr->flags);
        HOST_TO_BE(32, info->hdr->crc_dhdr);

        memset(iv, 0, sizeof(iv));
        error = tc_encrypt(info->cipher_chain, key, iv,
            (unsigned char *)info->hdr, sizeof(struct tchdr_dec), ehdr->enc);
        if (error) {
                tc_log(1, "Header encryption failed\n");
                goto error;
        }

        memset(iv, 0, sizeof(iv));
        error = tc_encrypt(info->cipher_chain, key_backup, iv,
            (unsigned char *)info->hdr,
            sizeof(struct tchdr_dec), ehdr_backup->enc);
        if (error) {
                tc_log(1, "Backup header encryption failed\n");
                goto error;
        }

        free_safe_mem(key);
        free_safe_mem(key_backup);

        if (backup_hdr != NULL)
                *backup_hdr = ehdr_backup;
        else
                free_safe_mem(ehdr_backup);

        return ehdr;
        /* NOT REACHED */

error:
        if (key)
                free_safe_mem(key);
        if (key_backup)
                free_safe_mem(key_backup);
        if (ehdr)
                free_safe_mem(ehdr);
        if (ehdr_backup)
                free_safe_mem(ehdr_backup);

        return NULL;
}