#include <Python.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include "scramble_crypt.h"

static PyObject *ScrambleError = 0;
static int pass_bits4=0;
static int pass_bits6=0;
static int _initialized=0;

static PyObject *
_scramble_init(PyObject *self, PyObject *args, PyObject *keywds)
{
    const char *keyfile="";
    static char *kwlist[] = {"keyfile", "pass4", "pass6", NULL};
    int r;

    if (!PyArg_ParseTupleAndKeywords(args, keywds, "s|ii", kwlist,
                                     &keyfile, &pass_bits4, &pass_bits6))
        return NULL;

    if (pass_bits4 < 0 || pass_bits4 > 32) {
        PyErr_SetString(ScrambleError, "pass4 doesn't make sense");
        return NULL;
    }
    if (pass_bits6 < 0 || pass_bits6 > 128) {
        PyErr_SetString(ScrambleError, "pass6 doesn't make sense");
        return NULL;
    }
    r = scramble_init_from_file(keyfile, 
				SCRAMBLE_BLOWFISH, 
				SCRAMBLE_BLOWFISH, NULL);
    if (r < 0) {
        PyErr_SetString(ScrambleError, "Initialization from file failed");
        return NULL;
    }

    _initialized = 1;
    Py_INCREF(Py_None);
    return Py_None;
}

static PyObject *
_scramble_ip4(PyObject *self, PyObject *args) {
    const char *ip_txt;
    struct in_addr ip4;
    char cbuf[256];

    if (!_initialized) {
	PyErr_SetString(ScrambleError, "Module not initialized");
	return NULL;
    }
    if (!PyArg_ParseTuple(args, "s", &ip_txt)) {
	PyErr_SetString(ScrambleError, "Function needs string argument");
	return NULL;
    }
    if (inet_pton(AF_INET, ip_txt, &ip4) <= 0) {
	PyErr_SetString(ScrambleError, "Can't parse IPv4 address");
	return NULL;
    }
    ip4.s_addr = scramble_ip4(ip4.s_addr, pass_bits4);
    if (cbuf != inet_ntop(AF_INET, &ip4, cbuf, sizeof(cbuf))) {
	PyErr_SetString(ScrambleError, "internal error(4)");
	return NULL;
    }
    return Py_BuildValue("s", cbuf);
}

static PyObject *
_scramble_ip6(PyObject *self, PyObject *args) {
    const char *ip_txt;
    struct in6_addr ip6;
    char cbuf[256];

    if (!_initialized) {
	PyErr_SetString(ScrambleError, "Module not initialized");
	return NULL;
    }
    if (!PyArg_ParseTuple(args, "s", &ip_txt)) {
	PyErr_SetString(ScrambleError, "Function needs string argument");
	return NULL;
    }
    if (inet_pton(AF_INET6, ip_txt, &ip6) <= 0) {
	PyErr_SetString(ScrambleError, "Can't parse IPv6 address");
	return NULL;
    }
    scramble_ip6(&ip6, pass_bits6);
    if (cbuf != inet_ntop(AF_INET6, &ip6, cbuf, sizeof(cbuf))) {
	PyErr_SetString(ScrambleError, "internal error(6)");
	return NULL;
    }
    return Py_BuildValue("s", cbuf);
}

static PyMethodDef ScrambleMethods[] = {
    {"init", (PyCFunction)_scramble_init, METH_VARARGS | METH_KEYWORDS,
     "init(keyfile=..., [pass4=..., [pass6=...]])\n"
     "Load secret scramble key from keyfile and optionally set pass values for\n"
     "IPv4 and IPv6 addresses.  keyfile will be created if does not exist."},
    {"ip4", _scramble_ip4, METH_VARARGS,
     "Scramble IPv4 address."},
    {"ip6", _scramble_ip6, METH_VARARGS,
     "Scramble IPv6 address."},

    {NULL, NULL, 0, NULL}        /* Sentinel */
};


PyMODINIT_FUNC
initscramble(void) {
    PyObject *m = Py_InitModule("scramble", ScrambleMethods);
    if (!m) return;
    ScrambleError = PyErr_NewException("scramble.error", NULL, NULL);
    _initialized = 0;
}

