#include <Python.h>

#include <openssl/asn1.h>
#include <openssl/rsa.h>
#include <openssl/err.h>
#include <openssl/ssl.h>
#include <openssl/pem.h>

static void delcsr(void *ptr)
{
    X509_REQ_free((X509_REQ *)ptr);
}

static X509_NAME *makeSubjectFromDict(PyObject *dict, unsigned long chtype)
{
    X509_NAME *subj = X509_NAME_new();
    int nid;
    PyObject *pKey = NULL, *pValue = NULL;
    int pos = 0;
    
    while (PyDict_Next(dict, &pos, &pKey, &pValue)) {
        char *key = PyString_AsString(pKey);
        char *val = PyString_AsString(pValue);

        if (strstr(key, "Email")) {
            nid = NID_pkcs9_emailAddress;
        } else {
            if ((nid = OBJ_txt2nid(key)) == NID_undef &&
                (nid = OBJ_sn2nid(key)) == NID_undef) {
                printf("Failed to find key %s\n", key);
                continue;
            }
        }

        X509_NAME_add_entry_by_NID(subj, nid, chtype, 
                                   (unsigned char *)val,
                                   -1, -1, 0);
    }
    return subj;
}

static PyObject *
readCSRFromFile(PyObject *self, PyObject *args)
{
    X509_REQ *req=NULL;
    BIO *in = NULL;
    char *fn = NULL;

    if (! PyArg_ParseTuple(args, "s", &fn))
        return NULL;

    if (! fn)
        return NULL;

    in = BIO_new(BIO_s_file());
    if (!in) {
        PyErr_SetString(PyExc_MemoryError, "Unable to create a BIO object");
        return NULL;
    }
    
    if (BIO_read_filename(in, fn) <= 0) {
        BIO_free_all(in);
        PyErr_SetString(PyExc_IOError, "Unable to read CSR from filename");
        return NULL;
    }

    /* We expect the CSR to be in PEM format, so try that first... */
    req=PEM_read_bio_X509_REQ(in, NULL, NULL, NULL);
    /* If that fails, see if it was in ASN1 format */
    if (!req)
        req=d2i_X509_REQ_bio(in, NULL);

    BIO_free_all(in);

    if (!req) {
        PyErr_SetString(PyExc_IOError, "Unable to get REQ object from file");
        return NULL;
    }
    return PyCObject_FromVoidPtr(req, delcsr);
}

static PyObject *
readCSRFromMemory(PyObject *self, PyObject *args)
{
    X509_REQ *req=NULL;
    BIO *in = NULL;
    char *ptr = NULL;
    long len = 0;

    if (! PyArg_ParseTuple(args, "s#", &ptr, &len))
        return NULL;

    if (len == 0) {
        PyErr_SetString(PyExc_IOError, "Zero length string passed");
        return NULL;
    }

    in = BIO_new_mem_buf(ptr, len);
    if (!in) {
        PyErr_SetString(PyExc_MemoryError, "Unable to create a BIO object");
        return NULL;
    }
    /* We expect the CSR to be in PEM format, so try that first... */
    req=PEM_read_bio_X509_REQ(in, NULL, NULL, NULL);
    /* If that fails, see if it was in ASN1 format */
    if (!req)
        req=d2i_X509_REQ_bio(in, NULL);

    BIO_free_all(in);

    if (!req) {
        PyErr_SetString(PyExc_IOError, "Unable to get REQ object from memory");
        return NULL;
    }
    return PyCObject_FromVoidPtr(req, delcsr);
}

static PyObject *
parseRequest(PyObject *self, PyObject *args)
{
    void *tmp = NULL;
    PyObject *dict = NULL;
    X509_REQ *req = NULL;
    X509_REQ_INFO *ri = NULL;

    if (! PyArg_ParseTuple(args, "O", &tmp))
        return NULL;

    req = (X509_REQ *)PyCObject_AsVoidPtr(tmp);
    if (!req)
        return NULL;
    ri = req->req_info;
    if (!ri)
        return NULL;

    dict = PyDict_New();

    /* What follows isn't ideal, so if you have ideas how to improve it... */

    /* CSR version number */
    {
//        const char *neg=(ri->version->type == V_ASN1_NEG_INTEGER)?"-":"";
        long l=0;
        int i;
        for(i=0; i<ri->version->length; i++) {
            l<<=8; l+=ri->version->data[i]; 
        }
        /* version should be ("%s%lx", neg, l) */
        PyDict_SetItem(dict, Py_BuildValue("s", "Version"),
                             Py_BuildValue("l", l));
    }
    /* Public Key information */
    {
        EVP_PKEY *pkey = X509_REQ_get_pubkey(req);
        PyObject *pkdict = PyDict_New();

        /* Public Key Algorithm */
        {
            BIO *bp = BIO_new(BIO_s_mem());
            if (bp && i2a_ASN1_OBJECT(bp, ri->pubkey->algor->algorithm) > 0) {
                char *strPtr = NULL;
                long sz = BIO_get_mem_data(bp, &strPtr);
                (void)BIO_set_close(bp, BIO_NOCLOSE);
                BIO_free(bp);
                PyDict_SetItem(pkdict, Py_BuildValue("s", "algorithm"),
                               PyString_FromStringAndSize(strPtr, sz));
            }
        }
        /* Public key Strength */
        if (pkey->type == EVP_PKEY_RSA) 
            PyDict_SetItem(pkdict, Py_BuildValue("s", "strength"),
                           Py_BuildValue("l", BN_num_bits(pkey->pkey.rsa->n)));

        PyDict_SetItem(dict, Py_BuildValue("s", "public_key"), pkdict);
        EVP_PKEY_free(pkey);
    }
    /* CSR Subject */
    {
        PyObject *sdict = PyDict_New();
        int i;
        struct sParts {
            int nid;
            const char *dict;
        } parts[] = {
            { NID_commonName, "CN" },
            { NID_countryName, "C" },
            { NID_organizationName, "O" },
            { NID_organizationalUnitName, "OU" },
            { NID_localityName, "L" },
            { NID_stateOrProvinceName, "ST" },
            { NID_pkcs9_emailAddress, "Email" },
            { 0, NULL }
        };

        for (i = 0; parts[i].nid != 0; i++) {
            char buf[512];
            int rv = X509_NAME_get_text_by_NID(ri->subject, parts[i].nid, buf, 512);
            if (rv > 0)
                PyDict_SetItem(sdict, Py_BuildValue("s", parts[i].dict),
                               PyString_FromStringAndSize(buf, rv));
        }
        PyDict_SetItem(dict, Py_BuildValue("s", "subject"), sdict);
    }

    return dict;
}

static PyObject *
createRequest(PyObject *self, PyObject *args)
{
    PyObject *tmp = NULL;
    PyObject *dict = NULL;
    PyObject *sdict = NULL;
    X509_REQ *req = NULL;
    X509_NAME *subj = NULL;
    RSA *rsa = NULL;
    EVP_PKEY *pkey = NULL;
    unsigned long chtype = MBSTRING_ASC;
    PyObject *rv = NULL;
    const EVP_MD *digest=EVP_sha1();

    if (! PyArg_ParseTuple(args, "OO", &tmp, &dict))
        return NULL;
    rsa = (RSA *)PyCObject_AsVoidPtr(tmp);
    if (!rsa || !dict || RSA_check_key(rsa) != 1) {
        PyErr_SetString(PyExc_ValueError, "Invalid RSA key supplied");
        return NULL;
    }

/*    if (tmp = PyDict_GetItem(dict, Py_BuildValue("s", "chtype"))
*/
        /* set chtype correctly! */
    /* allow digest type to be altered... */

    sdict = PyDict_GetItem(dict, Py_BuildValue("s", "subject"));
    if (!sdict)
        return NULL;

    /* create a new PKEY and then assign the rsa key to it */
    pkey = EVP_PKEY_new();
    EVP_PKEY_assign_RSA(pkey, rsa);

    /* make the subject */
    subj = makeSubjectFromDict(sdict, chtype);

    req=X509_REQ_new();
    if (!req)
        goto err;

    if (!X509_REQ_set_version(req, 0L))
        goto err; /* version 1 */
    if (!X509_REQ_set_subject_name(req, subj))
        goto err;
    if (!X509_REQ_set_pubkey(req, pkey))
        goto err;
    if (!X509_REQ_sign(req, pkey, digest))
        goto err;

    /* Verify we have a correctly created CSR before returning */
    if (!X509_REQ_verify(req,pkey))
        goto err;

    rv = PyCObject_FromVoidPtr(req, delcsr);
    /* We add a reference to the rsa object passed in, but how do we
     * decrement it later? We probably need to when the req object is
     * destroyed, but how do we get access to it?
     */
    Py_INCREF(tmp);

err:
    if (req && rv == NULL)
        X509_REQ_free(req);
    EVP_PKEY_free(pkey);
    X509_NAME_free(subj);

    return rv;
}

static PyObject *
getRequestAsString(PyObject *self, PyObject *args)
{
    X509_REQ *req = NULL;
    PyObject *temp = NULL, *rv = NULL;
    BIO *out = NULL;

    if (! PyArg_ParseTuple(args, "O", &temp))
        return NULL;

    req = (X509_REQ *)PyCObject_AsVoidPtr(temp);
    if (!req)
        return NULL;

    out = BIO_new(BIO_s_mem());
    if (out) {
        if (PEM_write_bio_X509_REQ(out, req) > 0) {
            char *ptr = NULL;
            long sz = 0;
            (void)BIO_flush(out);
            sz = BIO_get_mem_data(out, &ptr);
            (void)BIO_set_close(out, BIO_NOCLOSE);
            rv = PyString_FromStringAndSize(ptr, sz);
        }
        BIO_free(out);        
    }

    return rv;
}

static PyMethodDef CSRMethods[] = {
    { "fromFile", readCSRFromFile, METH_VARARGS, "Read a CSR from a file" },
    { "fromMemory", readCSRFromMemory, METH_VARARGS, "Read a CSR from a block of memory" },
    { "parse", parseRequest, METH_VARARGS, "Parse a request into a python dict" },
    { "create", createRequest, METH_VARARGS, "Create a request from information supplied" },
    { "asString", getRequestAsString, METH_VARARGS, "Get request as string" },
    { NULL, NULL, 0, NULL },
};

PyMODINIT_FUNC
initcsr(void)
{
    PyObject *mod = NULL;

    CRYPTO_malloc_init();
    ERR_load_crypto_strings(); 
    OpenSSL_add_all_algorithms();

    mod = Py_InitModule("csr", CSRMethods);
}