#include <Python.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <signal.h>

PyObject *g_socketerror; // socket.error

static PyObject *ucred_sendcreds(PyObject *self, PyObject *args, PyObject *keywds);
static PyObject *ucred_recvcreds(PyObject *self, PyObject *args, PyObject *keywds);

static PyMethodDef ucredMethods[] = {
    {"sendcreds", (PyCFunction)ucred_sendcreds, METH_VARARGS|METH_KEYWORDS, NULL},
    {"recvcreds", (PyCFunction)ucred_recvcreds, METH_VARARGS|METH_KEYWORDS, NULL},
    {NULL, NULL, 0, NULL}
};

void initucred(void) {
    PyObject *module;
    module = Py_InitModule("ucred", ucredMethods);

	if(-1 == PyModule_AddIntConstant(module, "SO_PASSCRED", SO_PASSCRED)) {
		return;
	}

    module = PyImport_ImportModule("socket");
    if(!module) {
        return;
    }
    g_socketerror = PyObject_GetAttrString(module, "error");
    if(!g_socketerror) {
        return;
    }
}

static PyObject *ucred_sendcreds(PyObject *self, PyObject *args, PyObject *keywds) {
    int fd;
    int ret;
    int pid=-1, uid=-1, gid=-1;
    struct msghdr msg;
    struct iovec iov[1];
	struct ucred uc;
	unsigned char control[CMSG_SPACE(sizeof(struct ucred))];
	struct cmsghdr *cur;

    static char *kwlist[] = {"fd",
							 "pid",
							 "uid",
							 "gid",
							 NULL};

//    kill(0, SIGTRAP);
    if (!PyArg_ParseTupleAndKeywords(args, keywds, "i|iii", kwlist, 
				&fd,
				&pid,
				&uid,
				&gid)) {
        return NULL;
    }
//    kill(0, SIGTRAP);

	if(pid!=-1)
		uc.pid=pid;
	else
		uc.pid=getpid();
	if(uid!=-1)
		uc.uid=uid;
	else
		uc.uid=getuid();
	if(gid!=-1)
		uc.gid=gid;
	else
		uc.gid=getgid();

	iov[0].iov_base="SCM_CREDENTIALS";
	iov[0].iov_len=15;

	memset(&msg, 0, sizeof(msg));
    msg.msg_iov = iov;
    msg.msg_iovlen = 1;
	msg.msg_control = &control;
	msg.msg_controllen = sizeof(control);

	cur = CMSG_FIRSTHDR(&msg);
	cur->cmsg_level = SOL_SOCKET;
	cur->cmsg_type = SCM_CREDENTIALS;
	cur->cmsg_len = CMSG_SPACE(sizeof(struct ucred));
	memcpy(CMSG_DATA(cur), &uc, sizeof(struct ucred));

    ret = sendmsg(fd, &msg, 0);
    if(ret < 0) {
        PyErr_SetFromErrno(g_socketerror);
        return NULL;
    }

    return Py_BuildValue("i", ret);
}

static PyObject *ucred_recvcreds(PyObject *self, PyObject *args, PyObject *keywds) {
    int fd;
    int flags=0;
    size_t maxsize=1024;
	char control[CMSG_SPACE(sizeof(struct ucred))];
    int ret;
	int pid=-1, uid=-1, gid=-1;
    struct msghdr msg;
    struct iovec iov[1];
	struct cmsghdr *cur;
	struct ucred *uc;


    static char *kwlist[] = {"fd", "flags", "maxsize", NULL};

    if (!PyArg_ParseTupleAndKeywords(args, keywds, "i|ii", kwlist,
            &fd, &flags, &maxsize)) {
        return NULL;
    }

	memset(&msg, 0, sizeof(msg));

    iov[0].iov_len = maxsize;
    iov[0].iov_base = malloc(maxsize);
    if (!iov[0].iov_base) {
        PyErr_NoMemory();
        return NULL;
    }
    msg.msg_iov = iov;
    msg.msg_iovlen = 1;

    msg.msg_control = &control;
    msg.msg_controllen = sizeof(control);

    ret = recvmsg(fd, &msg, flags);
    if (ret < 0) {
        PyErr_SetFromErrno(g_socketerror);
        free(iov[0].iov_base);
        return NULL;
    }

	cur=CMSG_FIRSTHDR(&msg);

	if(cur && cur->cmsg_type==SCM_CREDENTIALS && cur->cmsg_level==SOL_SOCKET)
	{
		assert(cur->cmsg_len >= sizeof(struct cmsghdr));
						// no null ancillary data messages?

		uc=(struct ucred*)CMSG_DATA(cur);
		pid=uc->pid;
		uid=uc->uid;
		gid=uc->gid;
	}

    {
        PyObject *r;
        r = Py_BuildValue("s#iii",
                iov[0].iov_base, ret,
                pid, uid, gid);
        free(iov[0].iov_base);
        return r;
    }
}