/*
 * Python interface to libgfshare()
 * Copyright (C) 2017 Chris Lamb <chris@chris-lamb.co.uk>
 *
 * 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.
 * 3. Neither the name of the University nor the names of its contributors
 *    may be used to endorse or promote products derived from this software
 *    without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE REGENTS 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 REGENTS 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 "Python.h"
#include "libgfshare.h"


struct module_state {
    PyObject *pRange;
    PyObject *pSampleFunc;
};

#define BUFFER_SIZE 4096
#define MAX_SHARECOUNT 253

#ifndef MIN
#define MIN(a,b) ((a)<(b))?(a):(b)
#endif


PyDoc_STRVAR(gfshare__doc__,
    "This Python library implements `Shamir's method for secret sharing\n"
    "<https://en.wikipedia.org/wiki/Shamir%27s_Secret_Sharing>`__ which can be used\n"
    "to split a \"secret\" into multiple parts. An arbitrary number of those parts\n"
    "are then needed to recover the original file but any smaller combination of\n"
    "parts are useless to an attacker.\n"
    "\n"
    "For instance, you might split a GPG key into a \"3-of-5\" share, putting one share\n"
    "on each of three computers and two shares on a USB memory stick. You can then\n"
    "use the GPG key on any of those three computers using the memory stick. If\n"
    "the memory stick is lost you can ultimately recover the key by bringing the three\n"
    "computers back together again.\n"
    "\n"
    "Under the hood it uses Daniel Silverstone's `libgfshare\n"
    "<http://www.digital-scurf.org/software/libgfshare>`__ library.\n"
    "\n"
    "Example\n"
    "=======\n"
    "\n"
    "    >>> import gfshare\n"
    "    >>> shares = gfshare.split(3, 5, b\"secret\")\n"
    "    >>> shares\n"
    "    {104: b'1\\x9cQ\\xd8\\xd3\\xaf',\n"
    "     164: b'\\x15\\xa4\\xcf7R\\xd2',\n"
    "     171: b'>\\xf5*\\xce\\xa2\\xe2',\n"
    "     173: b'd\\xd1\\xaaR\\xa5\\x1d',\n"
    "     183: b'\\x0c\\xb4Y\\x8apC'}\n"
    "    >>> gfshare.combine(shares)\n"
    "    b\"secret\"\n"
    "\n"
    "We remove 2 shares. We can still reconstruct the secret as we have 3 out\n"
    "of the 5 originals.\n"
    "\n"
    "    >>> del shares['104']\n"
    "    >>> del shares['171']\n"
    "    >>> gfshare.combine(shares)\n"
    "    b\"secret\"\n"
    "\n"
    "After removing another share we now cannot reconstruct the original secret.\n"
    "\n"
    "    >>> del shares['164']\n"
    "    >>> gfshare.combine(shares)\n"
    "    b'\\xea\\x87\\x99\\x06)\\x86'\n"
);

PyDoc_STRVAR(gfshare_split__doc__,
    "split(threshold, sharecount, secret)\n"
    "--\n"
    "\n"
    "Generate an \"n-of-m\" share of a given secret into a number of shares.\n"
    "\n"
    "The threshold number of shares cannot exceed the total share count.\n"
    "\n"
    ":param threshold:  The number of shares required to reconstruct the\n"
    "                   secret (min: 1).\n"
    ":param sharecount: The total number of shares to make (min: 2, max: 253)\n"
    ":param secret:     The secret to split.\n"
    ":type threshold:   int\n"
    ":type sharecount:  int\n"
    ":type secret:      bytes\n"
    ":returns:          The generated shares associated with their share number.\n"
    ":rtype:            dict\n"
    ":raises:           ValueError, TypeError\n"
);

PyDoc_STRVAR(gfshare_combine__doc__,
    "combine(shares)\n"
    "--\n"
    "\n"
    "Combine a number of shares generated by :func:`split` to form the original secret.\n"
    "\n"
    "If the share cannot be reconstructed, the result is undefined.\n"
    "\n"
    ":param shares: A dict of ``int``, ``bytes`` pairs.\n"
    ":type shares:  dict\n"
    ":returns:      The original secret (if available)\n"
    ":rtype:        bytes\n"
    ":raises:       ValueError, TypeError, RuntimeError\n"
);


int populate_sharenrs(PyObject *module, unsigned char *sharenrs, int sharecount) {
    int result = 1;
    PyObject *pArgs = NULL;
    PyObject *pResult = NULL;

    struct module_state *st = PyModule_GetState(module);
    if (!st)
        goto cleanup;

    pArgs = Py_BuildValue("(Oi)", st->pRange, sharecount);
    if (!pArgs)
        goto cleanup;

    pResult = PyObject_CallObject(st->pSampleFunc, pArgs);
    if (!pResult || !PyList_Check(pResult))
        goto cleanup;

    Py_ssize_t len = PySequence_Size(pResult);
    PyObject **items = PySequence_Fast_ITEMS(pResult);

    for (int i = 0; i < len; i++) {
        sharenrs[i] = PyLong_AsLong(items[i]);
    }

    Py_DECREF(items);
    result = 0;

cleanup:
    if (PyErr_Occurred())
        PyErr_Print();
    Py_XDECREF(pArgs);
    Py_XDECREF(pResult);
    return result;
}

static PyObject *
split(PyObject *m, PyObject *args, PyObject *keywds) {
    int threshold;
    int sharecount;
    char** shares = NULL;
    PyObject *result = NULL;
    gfshare_ctx *ctx = NULL;
    const char* secret;
    unsigned char* buffer = NULL;
    unsigned char* sharenrs = NULL;

    char *kwlist[] = {"threshold", "sharecount", "secret", NULL};
    if (!PyArg_ParseTupleAndKeywords(args, keywds, "iiy", kwlist,
                                     &threshold, &sharecount, &secret))
        goto cleanup;

    if (threshold < 1) {
        PyErr_SetString(PyExc_ValueError, "threshold must be >= 1");
        goto cleanup;
    }

    if (sharecount < 2) {
        PyErr_SetString(PyExc_ValueError, "sharecount must be >= 2");
        goto cleanup;
    }

    if (sharecount > MAX_SHARECOUNT) {
        PyErr_Format(PyExc_ValueError, "sharecount must be < %d", MAX_SHARECOUNT);
        goto cleanup;
    }

    if (threshold > sharecount) {
        PyErr_SetString(PyExc_ValueError, "sharecount must be >= threshold");
        goto cleanup;
    }

    const unsigned int sharesize = strlen(secret);

    buffer = malloc(BUFFER_SIZE);
    sharenrs = malloc(sharecount);
    if (buffer == NULL || sharenrs == NULL) {
        result = PyErr_NoMemory();
        goto cleanup;
    }

    if (populate_sharenrs(m, sharenrs, sharecount)) {
        PyErr_SetString(PyExc_RuntimeError, "could not populate sharenrs");
        goto cleanup;
    }

    ctx = gfshare_ctx_init_enc(sharenrs, sharecount, threshold,
                               MIN(BUFFER_SIZE, sharesize));

    if (!ctx) {
        PyErr_SetString(PyExc_RuntimeError, "gfshare_ctx_init_enc");
        goto cleanup;
    }

    result = PyDict_New();
    if (!result) {
        result = PyErr_NoMemory();
        goto cleanup;
    }

    shares = malloc(sharecount * sizeof(char*));
    if (shares == NULL) {
        result = PyErr_NoMemory();
        goto cleanup;
    }

    for (int i = 0; i < sharecount; ++i) {
        shares[i] = malloc(sharesize);
        if (shares[i] == NULL) {
            result = PyErr_NoMemory();
            goto cleanup;
        }
    }

    for (unsigned int offset = 0; offset < sharesize; offset += BUFFER_SIZE) {
        const unsigned int chunk = MIN(BUFFER_SIZE, (sharesize - offset));

        memcpy(buffer, secret + offset, chunk);
        gfshare_ctx_enc_setsecret(ctx, buffer);

        for (int i = 0; i < sharecount; ++i) {
            gfshare_ctx_enc_getshare(ctx, i, buffer);
            memcpy(shares[i] + offset, buffer, chunk);
        }
    }

    for (int i = 0; i < sharecount; ++i) {
        PyObject *idx = PyLong_FromLong(sharenrs[i]);
        PyObject *bytes = PyBytes_FromStringAndSize(shares[i], sharesize);
        PyDict_SetItem(result, idx, bytes);
        Py_DECREF(idx);
        Py_DECREF(bytes);
        free(shares[i]);
    }

cleanup:
    free(shares);
    free(sharenrs);
    free(buffer);
    if (ctx)
        gfshare_ctx_free(ctx);
    return result;
}

static PyObject *
combine(PyObject *m, PyObject *args, PyObject *keywds) {
    int i;
    int sharesize = -1;
    char* secret = NULL;
    PyObject *key, *value;
    PyObject *result = NULL;
    Py_ssize_t pos;
    gfshare_ctx *ctx = NULL;
    unsigned char* buffer = NULL;
    unsigned char* sharenrs = NULL;

    PyObject *shares;
    char *kwlist[] = {"shares", NULL};
    if (!PyArg_ParseTupleAndKeywords(args, keywds, "O", kwlist, &shares))
        goto fail;

    int sharecount = PyDict_Size(shares);
    if (sharecount < 2) {
        PyErr_SetString(PyExc_ValueError, "size of shares must be > 1");
        goto fail;
    }

    buffer = malloc(BUFFER_SIZE);
    sharenrs = malloc(sharecount);
    if (buffer == NULL || sharenrs == NULL) {
        result = PyErr_NoMemory();
        goto fail;
    }

    for (i = pos = 0; PyDict_Next(shares, &pos, &key, &value); ++i) {
        sharenrs[i] = PyLong_AsLong(key);

        if (sharenrs[i] < 1) {
            PyErr_SetString(PyExc_TypeError,
                            "shares keys not comprised entirely of positive numbers");
            goto fail;
        }

        if (!PyBytes_Check(value)) {
            PyErr_SetString(PyExc_TypeError,
                            "shares values not comprised entirely of byte objects");
            goto fail;
        }

        // Check consistent sizes
        int current_size = PyBytes_Size(value);
        if (sharesize >= 0 && sharesize != current_size) {
            PyErr_SetString(PyExc_TypeError,
                            "shares values not all bytes of the same length");
            goto fail;
        }
        sharesize = current_size;
    }

    ctx = gfshare_ctx_init_dec(sharenrs, sharecount, BUFFER_SIZE);
    if (!ctx) {
        PyErr_SetString(PyExc_RuntimeError, "gfshare_ctx_init_dec");
        goto fail;
    }

    secret = malloc(sharesize);
    if (secret == NULL) {
        result = PyErr_NoMemory();
        goto fail;
    }

    for (unsigned int offset = 0; offset < sharesize; offset += BUFFER_SIZE) {
        const unsigned int chunk = MIN(BUFFER_SIZE, (sharesize - offset));

        for (i = pos = 0; PyDict_Next(shares, &pos, &key, &value); ++i) {
            const char* valuevalue = PyBytes_AsString(value);
            memcpy(buffer, valuevalue + offset, chunk);
            gfshare_ctx_dec_giveshare(ctx, i, buffer);
        }

        gfshare_ctx_dec_extract(ctx, buffer);
        memcpy(secret + offset, buffer, chunk);
    }

    result = PyBytes_FromStringAndSize(secret, sharesize);
    if (!result) {
        PyErr_SetString(PyExc_RuntimeError, "could not construct result");
        result = NULL;
    }

fail:
    if (ctx)
        gfshare_ctx_free(ctx);
    free(secret);
    free(sharenrs);
    free(buffer);
    return result;
}

static PyMethodDef gfshare_methods[] = {
    {
        "split",
        (PyCFunction)split,
        METH_VARARGS | METH_KEYWORDS,
        gfshare_split__doc__
    },
    {
        "combine",
        (PyCFunction)combine,
        METH_VARARGS | METH_KEYWORDS,
        gfshare_combine__doc__
    },
    {NULL, NULL}
};

static struct PyModuleDef moduledef = {
        PyModuleDef_HEAD_INIT,
        "gfshare",
        gfshare__doc__,
        sizeof(struct module_state),
        gfshare_methods,
        NULL,
        NULL,
        NULL,
        NULL
};

PyMODINIT_FUNC
PyInit_gfshare(void)
{
    PyObject *module = PyModule_Create(&moduledef);
    PyObject *pRandomModule = NULL;

    if (module == NULL)
        return NULL;

    gfshare_fill_rand = (gfshare_rand_func_t)_PyOS_URandom;

    struct module_state *st = PyModule_GetState(module);
    st->pRange = PyObject_CallFunction((PyObject*)&PyRange_Type, "lll",
                                       1, 254, 1);
    if (!st->pRange || !PyRange_Check(st->pRange))
        goto fail;

    pRandomModule = PyImport_ImportModule("random");
    if (!pRandomModule)
        goto fail;

    st->pSampleFunc = PyObject_GetAttrString(pRandomModule, "sample");
    if (!st->pSampleFunc || !PyCallable_Check(st->pSampleFunc))
        goto fail;

    PyModule_AddObject(module, "_BUFFER_SIZE", PyLong_FromLong(BUFFER_SIZE));
    PyModule_AddObject(module, "_MAX_SHARECOUNT", PyLong_FromLong(MAX_SHARECOUNT));

    return module;

fail:
    Py_XDECREF(pRandomModule);
    Py_XDECREF(st->pRange);
    Py_XDECREF(st->pSampleFunc);
    Py_DECREF(module);
    return NULL;
}
