/*-
 * Copyright (c) 2015 Allan Jude <allanjude@FreeBSD.org>
 * Copyright (c) 2005-2011 Pawel Jakub Dawidek <pawel@dawidek.net>
 * 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 AUTHORS 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 AUTHORS 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 <stand.h>
#include <stdarg.h>
#include "geliboot.h"
#include "geliboot_internal.h"

struct known_dev {
        char                    name[GELIDEV_NAMELEN];
        struct geli_dev         *gdev;
        SLIST_ENTRY(known_dev)  entries;
};

SLIST_HEAD(known_dev_list, known_dev) known_devs_head = 
    SLIST_HEAD_INITIALIZER(known_devs_head);

static geli_ukey saved_keys[GELI_MAX_KEYS];
static unsigned int nsaved_keys = 0;

/*
 * Copy keys from local storage to the keybuf struct.
 * Destroy the local storage when finished.
 */
void
geli_export_key_buffer(struct keybuf *fkeybuf)
{
        unsigned int i;

        for (i = 0; i < nsaved_keys; i++) {
                fkeybuf->kb_ents[i].ke_type = KEYBUF_TYPE_GELI;
                memcpy(fkeybuf->kb_ents[i].ke_data, saved_keys[i],
                    G_ELI_USERKEYLEN);
        }
        fkeybuf->kb_nents = nsaved_keys;
        explicit_bzero(saved_keys, sizeof(saved_keys));
}

/*
 * Copy keys from a keybuf struct into local storage.
 * Zero out the keybuf.
 */
void
geli_import_key_buffer(struct keybuf *skeybuf)
{
        unsigned int i;

        for (i = 0; i < skeybuf->kb_nents && i < GELI_MAX_KEYS; i++) {
                memcpy(saved_keys[i], skeybuf->kb_ents[i].ke_data,
                    G_ELI_USERKEYLEN);
                explicit_bzero(skeybuf->kb_ents[i].ke_data,
                    G_ELI_USERKEYLEN);
                skeybuf->kb_ents[i].ke_type = KEYBUF_TYPE_NONE;
        }
        nsaved_keys = skeybuf->kb_nents;
        skeybuf->kb_nents = 0;
}

void
geli_add_key(geli_ukey key)
{

        /*
         * If we run out of key space, the worst that will happen is
         * it will ask the user for the password again.
         */
        if (nsaved_keys < GELI_MAX_KEYS) {
                memcpy(saved_keys[nsaved_keys], key, G_ELI_USERKEYLEN);
                nsaved_keys++;
        }
}

static int
geli_findkey(struct geli_dev *gdev, u_char *mkey)
{
        u_int keynum;
        int i;

        if (gdev->keybuf_slot >= 0) {
                if (g_eli_mkey_decrypt_any(&gdev->md, saved_keys[gdev->keybuf_slot],
                    mkey, &keynum) == 0) {
                        return (0);
                }
        }

        for (i = 0; i < nsaved_keys; i++) {
                if (g_eli_mkey_decrypt_any(&gdev->md, saved_keys[i], mkey,
                    &keynum) == 0) {
                        gdev->keybuf_slot = i;
                        return (0);
                }
        }

        return (1);
}

/*
 * Read the last sector of a drive or partition and see if it is GELI encrypted.
 */
struct geli_dev *
geli_taste(geli_readfunc readfunc, void *readpriv, daddr_t lastsector,
    const char *namefmt, ...)
{
        va_list args;
        struct g_eli_metadata md;
        struct known_dev *kdev;
        struct geli_dev *gdev;
        u_char *buf;
        char devname[GELIDEV_NAMELEN];
        int error;
        off_t alignsector;

        /*
         * Format the name into a temp buffer and use that to search for an
         * existing known_dev instance.  If not found, this has the side effect
         * of initializing kdev to NULL.
         */
        va_start(args, namefmt);
        vsnprintf(devname, sizeof(devname), namefmt, args);
        va_end(args);
        SLIST_FOREACH(kdev, &known_devs_head, entries) {
                if (strcmp(kdev->name, devname) == 0)
                        return (kdev->gdev);
        }

        /* Determine whether the new device is geli-encrypted... */
        if ((buf = malloc(DEV_GELIBOOT_BSIZE)) == NULL)
                goto out;
        alignsector = rounddown2(lastsector * DEV_BSIZE, DEV_GELIBOOT_BSIZE);
        if (alignsector + DEV_GELIBOOT_BSIZE > ((lastsector + 1) * DEV_BSIZE)) {
                /* Don't read past the end of the disk */
                alignsector = (lastsector * DEV_BSIZE) + DEV_BSIZE -
                    DEV_GELIBOOT_BSIZE;
        }
        error = readfunc(NULL, readpriv, alignsector, buf, DEV_GELIBOOT_BSIZE);
        if (error != 0) {
                goto out;
        }

        /*
         * We have a new known_device.  Whether it's geli-encrypted or not,
         * record its existence so we can avoid doing IO to probe it next time.
         */
        if ((kdev = malloc(sizeof(*kdev))) == NULL)
                goto out;
        strlcpy(kdev->name, devname, sizeof(kdev->name));
        kdev->gdev = NULL;
        SLIST_INSERT_HEAD(&known_devs_head, kdev, entries);

        /* Extract the last 4k sector of the disk. */
        error = eli_metadata_decode(buf, &md);
        if (error != 0) {
                /* Try the last 512 byte sector instead. */
                error = eli_metadata_decode(buf +
                    (DEV_GELIBOOT_BSIZE - DEV_BSIZE), &md);
                if (error != 0) {
                        goto out;
                }
        }

        if (!(md.md_flags & G_ELI_FLAG_GELIBOOT)) {
                /* The GELIBOOT feature is not activated */
                goto out;
        }
        if ((md.md_flags & G_ELI_FLAG_ONETIME)) {
                /* Swap device, skip it. */
                goto out;
        }

        /*
         * It's geli-encrypted, create a geli_dev for it and link it into the
         * known_dev instance.
         */
        gdev = malloc(sizeof(struct geli_dev));
        if (gdev == NULL)
                goto out;
        gdev->part_end = lastsector;
        gdev->keybuf_slot = -1;
        gdev->md = md;
        gdev->name = kdev->name;
        eli_metadata_softc(&gdev->sc, &md, DEV_BSIZE,
            (lastsector + DEV_BSIZE) * DEV_BSIZE);
        kdev->gdev = gdev;
out:
        free(buf);
        if (kdev == NULL)
                return (NULL);
        return (kdev->gdev);
}

/*
 * Attempt to decrypt the device.  This will try existing keys first, then will
 * prompt for a passphrase if there are no existing keys that work.
 */
static int
geli_probe(struct geli_dev *gdev, const char *passphrase, u_char *mkeyp)
{
        u_char key[G_ELI_USERKEYLEN], mkey[G_ELI_DATAIVKEYLEN], *mkp;
        u_int keynum;
        struct hmac_ctx ctx;
        int error;

        if (mkeyp != NULL) {
                memcpy(&mkey, mkeyp, G_ELI_DATAIVKEYLEN);
                explicit_bzero(mkeyp, G_ELI_DATAIVKEYLEN);
                goto found_key;
        }

        if (geli_findkey(gdev, mkey) == 0) {
                goto found_key;
        }

        g_eli_crypto_hmac_init(&ctx, NULL, 0);
        /*
         * Prepare Derived-Key from the user passphrase.
         */
        if (gdev->md.md_iterations < 0) {
                /* XXX TODO: Support loading key files. */
                return (1);
        } else if (gdev->md.md_iterations == 0) {
                g_eli_crypto_hmac_update(&ctx, gdev->md.md_salt,
                    sizeof(gdev->md.md_salt));
                g_eli_crypto_hmac_update(&ctx, (const uint8_t *)passphrase,
                    strlen(passphrase));
        } else if (gdev->md.md_iterations > 0) {
                printf("Calculating GELI Decryption Key for %s %d"
                    " iterations...\n", gdev->name, gdev->md.md_iterations);
                u_char dkey[G_ELI_USERKEYLEN];

                pkcs5v2_genkey(dkey, sizeof(dkey), gdev->md.md_salt,
                    sizeof(gdev->md.md_salt), passphrase,
                    gdev->md.md_iterations);
                g_eli_crypto_hmac_update(&ctx, dkey, sizeof(dkey));
                explicit_bzero(dkey, sizeof(dkey));
        }

        g_eli_crypto_hmac_final(&ctx, key, 0);

        error = g_eli_mkey_decrypt_any(&gdev->md, key, mkey, &keynum);
        if (error == -1) {
                explicit_bzero(mkey, sizeof(mkey));
                explicit_bzero(key, sizeof(key));
                printf("Bad GELI key: bad password?\n");
                return (error);
        } else if (error != 0) {
                explicit_bzero(mkey, sizeof(mkey));
                explicit_bzero(key, sizeof(key));
                printf("Failed to decrypt GELI master key: %d\n", error);
                return (error);
        } else {
                /* Add key to keychain */
                geli_add_key(key);
                explicit_bzero(&key, sizeof(key));
        }

found_key:
        /* Store the keys */
        bcopy(mkey, gdev->sc.sc_mkey, sizeof(gdev->sc.sc_mkey));
        bcopy(mkey, gdev->sc.sc_ivkey, sizeof(gdev->sc.sc_ivkey));
        mkp = mkey + sizeof(gdev->sc.sc_ivkey);
        if ((gdev->sc.sc_flags & G_ELI_FLAG_AUTH) == 0) {
                bcopy(mkp, gdev->sc.sc_ekey, G_ELI_DATAKEYLEN);
        } else {
                /*
                 * The encryption key is: ekey = HMAC_SHA512(Data-Key, 0x10)
                 */
                g_eli_crypto_hmac(mkp, G_ELI_MAXKEYLEN, (const uint8_t *)"\x10", 1,
                    gdev->sc.sc_ekey, 0);
        }
        explicit_bzero(mkey, sizeof(mkey));

        /* Initialize the per-sector IV. */
        switch (gdev->sc.sc_ealgo) {
        case CRYPTO_AES_XTS:
                break;
        default:
                SHA256_Init(&gdev->sc.sc_ivctx);
                SHA256_Update(&gdev->sc.sc_ivctx, gdev->sc.sc_ivkey,
                    sizeof(gdev->sc.sc_ivkey));
                break;
        }

        return (0);
}

int
geli_io(struct geli_dev *gdev, geli_op_t enc, off_t offset, u_char *buf,
    size_t bytes)
{
        u_char iv[G_ELI_IVKEYLEN];
        u_char *pbuf;
        int error;
        off_t dstoff;
        uint64_t keyno;
        size_t n, nsec, secsize;
        struct g_eli_key gkey;

        pbuf = buf;

        secsize = gdev->sc.sc_sectorsize;
        nsec = bytes / secsize;
        if (nsec == 0) {
                /*
                 * A read of less than the GELI sector size has been
                 * requested. The caller provided destination buffer may
                 * not be big enough to boost the read to a full sector,
                 * so just attempt to decrypt the truncated sector.
                 */
                secsize = bytes;
                nsec = 1;
        }

        for (n = 0, dstoff = offset; n < nsec; n++, dstoff += secsize) {

                g_eli_crypto_ivgen(&gdev->sc, dstoff, iv, G_ELI_IVKEYLEN);

                /* Get the key that corresponds to this offset. */
                keyno = (dstoff >> G_ELI_KEY_SHIFT) / secsize;
                g_eli_key_fill(&gdev->sc, &gkey, keyno);

                error = geliboot_crypt(gdev->sc.sc_ealgo, enc, pbuf, secsize,
                    gkey.gek_key, gdev->sc.sc_ekeylen, iv);

                if (error != 0) {
                        explicit_bzero(&gkey, sizeof(gkey));
                        printf("%s: Failed to %s!", __func__,
                            enc ? "encrypt" : "decrypt");
                        return (error);
                }
                pbuf += secsize;
        }
        explicit_bzero(&gkey, sizeof(gkey));
        return (0);
}

int
geli_havekey(struct geli_dev *gdev)
{
        u_char mkey[G_ELI_DATAIVKEYLEN];
        int err;

        err = ENOENT;
        if (geli_findkey(gdev, mkey) == 0) {
                if (geli_probe(gdev, NULL, mkey) == 0)
                        err = 0;
                explicit_bzero(mkey, sizeof(mkey));
        }
        return (err);
}

int
geli_passphrase(struct geli_dev *gdev, char *pw)
{
        int i;

        /* TODO: Implement GELI keyfile(s) support */
        for (i = 0; i < 3; i++) {
                /* Try cached passphrase */
                if (i == 0 && pw[0] != '\0') {
                        if (geli_probe(gdev, pw, NULL) == 0) {
                                return (0);
                        }
                }
                printf("GELI Passphrase for %s ", gdev->name);
                pwgets(pw, GELI_PW_MAXLEN,
                    (gdev->md.md_flags & G_ELI_FLAG_GELIDISPLAYPASS) == 0);
                printf("\n");
                if (geli_probe(gdev, pw, NULL) == 0) {
                        return (0);
                }
        }

        return (1);
}