[people/ms/libloc.git] / src / python / database.c

/*
	libloc - A library to determine the location of someone on the Internet

	Copyright (C) 2017 IPFire Development Team <info@ipfire.org>

	This library is free software; you can redistribute it and/or
	modify it under the terms of the GNU Lesser General Public
	License as published by the Free Software Foundation; either
	version 2.1 of the License, or (at your option) any later version.

	This library is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
	Lesser General Public License for more details.
*/

#include <Python.h>

#include <loc/libloc.h>
#include <loc/database.h>

#include "locationmodule.h"
#include "as.h"
#include "database.h"
#include "network.h"

static PyObject* Database_new(PyTypeObject* type, PyObject* args, PyObject* kwds) {
	DatabaseObject* self = (DatabaseObject*)type->tp_alloc(type, 0);

	return (PyObject*)self;
}

static void Database_dealloc(DatabaseObject* self) {
	if (self->db)
		loc_database_unref(self->db);

	if (self->path)
		free(self->path);

	Py_TYPE(self)->tp_free((PyObject* )self);
}

static int Database_init(DatabaseObject* self, PyObject* args, PyObject* kwargs) {
	const char* path = NULL;

	if (!PyArg_ParseTuple(args, "s", &path))
		return -1;

	self->path = strdup(path);

	// Open the file for reading
	FILE* f = fopen(self->path, "r");
	if (!f) {
		PyErr_SetFromErrno(PyExc_IOError);
		return -1;
	}

	// Load the database
	int r = loc_database_new(loc_ctx, &self->db, f);
	fclose(f);

	// Return on any errors
	if (r)
		return -1;

	return 0;
}

static PyObject* Database_repr(DatabaseObject* self) {
	return PyUnicode_FromFormat("<Database %s>", self->path);
}

static PyObject* Database_get_description(DatabaseObject* self) {
	const char* description = loc_database_get_description(self->db);

	return PyUnicode_FromString(description);
}

static PyObject* Database_get_vendor(DatabaseObject* self) {
	const char* vendor = loc_database_get_vendor(self->db);

	return PyUnicode_FromString(vendor);
}

static PyObject* Database_get_license(DatabaseObject* self) {
	const char* license = loc_database_get_license(self->db);

	return PyUnicode_FromString(license);
}

static PyObject* Database_get_created_at(DatabaseObject* self) {
	time_t created_at = loc_database_created_at(self->db);

	return PyLong_FromLong(created_at);
}

static PyObject* Database_get_as(DatabaseObject* self, PyObject* args) {
	struct loc_as* as = NULL;
	uint32_t number = 0;

	if (!PyArg_ParseTuple(args, "i", &number))
		return NULL;

	// Try to retrieve the AS
	int r = loc_database_get_as(self->db, &as, number);

	// We got an AS
	if (r == 0) {
		PyObject* obj = new_as(&ASType, as);
		loc_as_unref(as);

		return obj;

	// Nothing found
	} else if (r == 1) {
		Py_RETURN_NONE;
	}

	// Unexpected error
	return NULL;
}

static PyObject* Database_lookup(DatabaseObject* self, PyObject* args) {
	struct loc_network* network = NULL;
	const char* address = NULL;

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

	// Try to retrieve a matching network
	int r = loc_database_lookup_from_string(self->db, address, &network);

	// We got a network
	if (r == 0) {
		PyObject* obj = new_network(&NetworkType, network);
		loc_network_unref(network);

		return obj;

	// Nothing found
	} else if (r == 1) {
		Py_RETURN_NONE;

	// Invalid input
	} else if (r == -EINVAL) {
		PyErr_Format(PyExc_ValueError, "Invalid IP address: %s", address);
		return NULL;
	}

	// Unexpected error
	return NULL;
}

static PyObject* new_database_enumerator(PyTypeObject* type, struct loc_database_enumerator* enumerator) {
	DatabaseEnumeratorObject* self = (DatabaseEnumeratorObject*)type->tp_alloc(type, 0);
	if (self) {
		self->enumerator = loc_database_enumerator_ref(enumerator);
	}

	return (PyObject*)self;
}

static PyObject* Database_search_as(DatabaseObject* self, PyObject* args) {
	const char* string = NULL;

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

	struct loc_database_enumerator* enumerator;

	int r = loc_database_enumerator_new(&enumerator, self->db, LOC_DB_ENUMERATE_ASES);
	if (r) {
		PyErr_SetFromErrno(PyExc_SystemError);
		return NULL;
	}

	// Search string we are searching for
	loc_database_enumerator_set_string(enumerator, string);

	PyObject* obj = new_database_enumerator(&DatabaseEnumeratorType, enumerator);
	loc_database_enumerator_unref(enumerator);

	return obj;
}

static PyObject* Database_search_networks(DatabaseObject* self, PyObject* args, PyObject* kwargs) {
	char* kwlist[] = { "country_code", "asn", NULL };
	const char* country_code = NULL;
	unsigned int asn = 0;

    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|si", kwlist, &country_code, &asn))
		return NULL;

	struct loc_database_enumerator* enumerator;
	int r = loc_database_enumerator_new(&enumerator, self->db, LOC_DB_ENUMERATE_NETWORKS);
	if (r) {
		PyErr_SetFromErrno(PyExc_SystemError);
		return NULL;
	}

	// Set country code we are searching for
	if (country_code) {
		r = loc_database_enumerator_set_country_code(enumerator, country_code);

		if (r) {
			PyErr_SetFromErrno(PyExc_SystemError);
			return NULL;
		}
	}

	// Set the ASN we are searching for
	if (asn) {
		r = loc_database_enumerator_set_asn(enumerator, asn);

		if (r) {
			PyErr_SetFromErrno(PyExc_SystemError);
			return NULL;
		}
	}

	PyObject* obj = new_database_enumerator(&DatabaseEnumeratorType, enumerator);
	loc_database_enumerator_unref(enumerator);

	return obj;
}

static struct PyMethodDef Database_methods[] = {
	{
		"get_as",
		(PyCFunction)Database_get_as,
		METH_VARARGS,
		NULL,
	},
	{
		"lookup",
		(PyCFunction)Database_lookup,
		METH_VARARGS,
		NULL,
	},
	{
		"search_as",
		(PyCFunction)Database_search_as,
		METH_VARARGS,
		NULL,
	},
	{
		"search_networks",
		(PyCFunction)Database_search_networks,
		METH_VARARGS|METH_KEYWORDS,
		NULL,
	},
	{ NULL },
};

static struct PyGetSetDef Database_getsetters[] = {
	{
		"created_at",
		(getter)Database_get_created_at,
		NULL,
		NULL,
		NULL,
	},
	{
		"description",
		(getter)Database_get_description,
		NULL,
		NULL,
		NULL,
	},
	{
		"license",
		(getter)Database_get_license,
		NULL,
		NULL,
		NULL,
	},
	{
		"vendor",
		(getter)Database_get_vendor,
		NULL,
		NULL,
		NULL,
	},
	{ NULL },
};

PyTypeObject DatabaseType = {
	PyVarObject_HEAD_INIT(NULL, 0)
	.tp_name =               "location.Database",
	.tp_basicsize =          sizeof(DatabaseObject),
	.tp_flags =              Py_TPFLAGS_DEFAULT|Py_TPFLAGS_BASETYPE,
	.tp_new =                Database_new,
	.tp_dealloc =            (destructor)Database_dealloc,
	.tp_init =               (initproc)Database_init,
	.tp_doc =                "Database object",
	.tp_methods =            Database_methods,
	.tp_getset =             Database_getsetters,
	.tp_repr =               (reprfunc)Database_repr,
};

static PyObject* DatabaseEnumerator_new(PyTypeObject* type, PyObject* args, PyObject* kwds) {
	DatabaseEnumeratorObject* self = (DatabaseEnumeratorObject*)type->tp_alloc(type, 0);

	return (PyObject*)self;
}

static void DatabaseEnumerator_dealloc(DatabaseEnumeratorObject* self) {
	loc_database_enumerator_unref(self->enumerator);

	Py_TYPE(self)->tp_free((PyObject* )self);
}

static PyObject* DatabaseEnumerator_next(DatabaseEnumeratorObject* self) {
	// Enumerate all networks
	struct loc_network* network = loc_database_enumerator_next_network(self->enumerator);
	if (network) {
		PyObject* obj = new_network(&NetworkType, network);
		loc_network_unref(network);

		return obj;
	}

	// Enumerate all ASes
	struct loc_as* as = loc_database_enumerator_next_as(self->enumerator);
	if (as) {
		PyObject* obj = new_as(&ASType, as);
		loc_as_unref(as);

		return obj;
	}

	// Nothing found, that means the end
	PyErr_SetNone(PyExc_StopIteration);
	return NULL;
}

PyTypeObject DatabaseEnumeratorType = {
	PyVarObject_HEAD_INIT(NULL, 0)
	.tp_name =               "location.DatabaseEnumerator",
	.tp_basicsize =          sizeof(DatabaseEnumeratorObject),
	.tp_flags =              Py_TPFLAGS_DEFAULT,
	.tp_alloc =              PyType_GenericAlloc,
	.tp_new =                DatabaseEnumerator_new,
	.tp_dealloc =            (destructor)DatabaseEnumerator_dealloc,
	.tp_iter =               PyObject_SelfIter,
	.tp_iternext =           (iternextfunc)DatabaseEnumerator_next,
};
Commit	Line	Data
9cdf6c53 MT	1	/*
	2	libloc - A library to determine the location of someone on the Internet
	3
	4	Copyright (C) 2017 IPFire Development Team <info@ipfire.org>
	5
	6	This library is free software; you can redistribute it and/or
	7	modify it under the terms of the GNU Lesser General Public
	8	License as published by the Free Software Foundation; either
	9	version 2.1 of the License, or (at your option) any later version.
	10
	11	This library is distributed in the hope that it will be useful,
	12	but WITHOUT ANY WARRANTY; without even the implied warranty of
	13	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
	14	Lesser General Public License for more details.
	15	*/
	16
	17	#include <Python.h>
	18
9fc7f001 MT	19	#include <loc/libloc.h>
	20	#include <loc/database.h>
	21
1da9cd39	22	#include "locationmodule.h"
86ca7ef7	23	#include "as.h"
9cdf6c53	24	#include "database.h"
31edab76	25	#include "network.h"
9cdf6c53 MT	26
9cdf6c53 MT	27	static PyObject* Database_new(PyTypeObject* type, PyObject* args, PyObject* kwds) {
9cdf6c53	28	DatabaseObject* self = (DatabaseObject*)type->tp_alloc(type, 0);
9cdf6c53 MT	29
	30	return (PyObject*)self;
	31	}
	32
	33	static void Database_dealloc(DatabaseObject* self) {
	34	if (self->db)
	35	loc_database_unref(self->db);
	36
6fd96715 MT	37	if (self->path)
	38	free(self->path);
	39
9cdf6c53 MT	40	Py_TYPE(self)->tp_free((PyObject* )self);
	41	}
	42
	43	static int Database_init(DatabaseObject* self, PyObject* args, PyObject* kwargs) {
	44	const char* path = NULL;
	45
	46	if (!PyArg_ParseTuple(args, "s", &path))
	47	return -1;
	48
6fd96715 MT	49	self->path = strdup(path);
6fd96715 MT	50
9cdf6c53	51	// Open the file for reading
6fd96715	52	FILE* f = fopen(self->path, "r");
774eea26 MT	53	if (!f) {
774eea26 MT	54	PyErr_SetFromErrno(PyExc_IOError);
9cdf6c53	55	return -1;
774eea26	56	}
9cdf6c53 MT	57
9cdf6c53 MT	58	// Load the database
38e07ee0	59	int r = loc_database_new(loc_ctx, &self->db, f);
9cdf6c53 MT	60	fclose(f);
	61
	62	// Return on any errors
	63	if (r)
	64	return -1;
	65
	66	return 0;
	67	}
	68
6fd96715 MT	69	static PyObject* Database_repr(DatabaseObject* self) {
	70	return PyUnicode_FromFormat("<Database %s>", self->path);
	71	}
	72
d99b0256 MT	73	static PyObject* Database_get_description(DatabaseObject* self) {
	74	const char* description = loc_database_get_description(self->db);
	75
	76	return PyUnicode_FromString(description);
	77	}
	78
	79	static PyObject* Database_get_vendor(DatabaseObject* self) {
	80	const char* vendor = loc_database_get_vendor(self->db);
	81
	82	return PyUnicode_FromString(vendor);
	83	}
	84
4bf49d00 MT	85	static PyObject* Database_get_license(DatabaseObject* self) {
	86	const char* license = loc_database_get_license(self->db);
	87
	88	return PyUnicode_FromString(license);
	89	}
	90
53524b2d MT	91	static PyObject* Database_get_created_at(DatabaseObject* self) {
	92	time_t created_at = loc_database_created_at(self->db);
	93
	94	return PyLong_FromLong(created_at);
	95	}
	96
86ca7ef7 MT	97	static PyObject* Database_get_as(DatabaseObject* self, PyObject* args) {
	98	struct loc_as* as = NULL;
	99	uint32_t number = 0;
	100
	101	if (!PyArg_ParseTuple(args, "i", &number))
	102	return NULL;
	103
	104	// Try to retrieve the AS
	105	int r = loc_database_get_as(self->db, &as, number);
86ca7ef7	106
4a0a0f7e MT	107	// We got an AS
4a0a0f7e MT	108	if (r == 0) {
86ca7ef7 MT	109	PyObject* obj = new_as(&ASType, as);
	110	loc_as_unref(as);
	111
	112	return obj;
86ca7ef7 MT	113
86ca7ef7 MT	114	// Nothing found
4a0a0f7e MT	115	} else if (r == 1) {
	116	Py_RETURN_NONE;
	117	}
	118
	119	// Unexpected error
	120	return NULL;
86ca7ef7 MT	121	}
86ca7ef7 MT	122
31edab76 MT	123	static PyObject* Database_lookup(DatabaseObject* self, PyObject* args) {
	124	struct loc_network* network = NULL;
	125	const char* address = NULL;
	126
	127	if (!PyArg_ParseTuple(args, "s", &address))
	128	return NULL;
	129
	130	// Try to retrieve a matching network
	131	int r = loc_database_lookup_from_string(self->db, address, &network);
	132
	133	// We got a network
	134	if (r == 0) {
	135	PyObject* obj = new_network(&NetworkType, network);
	136	loc_network_unref(network);
	137
	138	return obj;
	139
	140	// Nothing found
	141	} else if (r == 1) {
	142	Py_RETURN_NONE;
927e82f2 MT	143
	144	// Invalid input
	145	} else if (r == -EINVAL) {
	146	PyErr_Format(PyExc_ValueError, "Invalid IP address: %s", address);
	147	return NULL;
31edab76 MT	148	}
	149
	150	// Unexpected error
	151	return NULL;
	152	}
	153
afb426df MT	154	static PyObject* new_database_enumerator(PyTypeObject* type, struct loc_database_enumerator* enumerator) {
	155	DatabaseEnumeratorObject* self = (DatabaseEnumeratorObject*)type->tp_alloc(type, 0);
	156	if (self) {
	157	self->enumerator = loc_database_enumerator_ref(enumerator);
	158	}
	159
	160	return (PyObject*)self;
	161	}
	162
	163	static PyObject* Database_search_as(DatabaseObject* self, PyObject* args) {
	164	const char* string = NULL;
	165
	166	if (!PyArg_ParseTuple(args, "s", &string))
	167	return NULL;
	168
	169	struct loc_database_enumerator* enumerator;
	170
ccc7ab4e	171	int r = loc_database_enumerator_new(&enumerator, self->db, LOC_DB_ENUMERATE_ASES);
afb426df MT	172	if (r) {
	173	PyErr_SetFromErrno(PyExc_SystemError);
	174	return NULL;
	175	}
	176
	177	// Search string we are searching for
	178	loc_database_enumerator_set_string(enumerator, string);
	179
	180	PyObject* obj = new_database_enumerator(&DatabaseEnumeratorType, enumerator);
	181	loc_database_enumerator_unref(enumerator);
	182
	183	return obj;
	184	}
	185
ccc7ab4e MT	186	static PyObject* Database_search_networks(DatabaseObject* self, PyObject* args, PyObject* kwargs) {
	187	char* kwlist[] = { "country_code", "asn", NULL };
	188	const char* country_code = NULL;
	189	unsigned int asn = 0;
	190
	191	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "\|si", kwlist, &country_code, &asn))
	192	return NULL;
	193
	194	struct loc_database_enumerator* enumerator;
	195	int r = loc_database_enumerator_new(&enumerator, self->db, LOC_DB_ENUMERATE_NETWORKS);
	196	if (r) {
	197	PyErr_SetFromErrno(PyExc_SystemError);
	198	return NULL;
	199	}
	200
	201	// Set country code we are searching for
	202	if (country_code) {
	203	r = loc_database_enumerator_set_country_code(enumerator, country_code);
	204
	205	if (r) {
	206	PyErr_SetFromErrno(PyExc_SystemError);
	207	return NULL;
	208	}
	209	}
	210
	211	// Set the ASN we are searching for
	212	if (asn) {
	213	r = loc_database_enumerator_set_asn(enumerator, asn);
	214
	215	if (r) {
	216	PyErr_SetFromErrno(PyExc_SystemError);
	217	return NULL;
	218	}
	219	}
	220
	221	PyObject* obj = new_database_enumerator(&DatabaseEnumeratorType, enumerator);
	222	loc_database_enumerator_unref(enumerator);
	223
	224	return obj;
	225	}
	226
9cdf6c53	227	static struct PyMethodDef Database_methods[] = {
86ca7ef7 MT	228	{
	229	"get_as",
	230	(PyCFunction)Database_get_as,
	231	METH_VARARGS,
	232	NULL,
	233	},
31edab76 MT	234	{
	235	"lookup",
	236	(PyCFunction)Database_lookup,
	237	METH_VARARGS,
	238	NULL,
	239	},
afb426df MT	240	{
	241	"search_as",
	242	(PyCFunction)Database_search_as,
	243	METH_VARARGS,
	244	NULL,
	245	},
ccc7ab4e MT	246	{
	247	"search_networks",
	248	(PyCFunction)Database_search_networks,
	249	METH_VARARGS\|METH_KEYWORDS,
	250	NULL,
	251	},
9cdf6c53 MT	252	{ NULL },
	253	};
	254
d99b0256	255	static struct PyGetSetDef Database_getsetters[] = {
53524b2d MT	256	{
	257	"created_at",
	258	(getter)Database_get_created_at,
	259	NULL,
	260	NULL,
	261	NULL,
	262	},
d99b0256 MT	263	{
	264	"description",
	265	(getter)Database_get_description,
	266	NULL,
	267	NULL,
	268	NULL,
	269	},
4bf49d00 MT	270	{
	271	"license",
	272	(getter)Database_get_license,
	273	NULL,
	274	NULL,
	275	NULL,
	276	},
d99b0256 MT	277	{
	278	"vendor",
	279	(getter)Database_get_vendor,
	280	NULL,
	281	NULL,
53524b2d	282	NULL,
d99b0256 MT	283	},
	284	{ NULL },
	285	};
	286
9cdf6c53 MT	287	PyTypeObject DatabaseType = {
9cdf6c53 MT	288	PyVarObject_HEAD_INIT(NULL, 0)
d42e1dcd MT	289	.tp_name = "location.Database",
	290	.tp_basicsize = sizeof(DatabaseObject),
	291	.tp_flags = Py_TPFLAGS_DEFAULT\|Py_TPFLAGS_BASETYPE,
	292	.tp_new = Database_new,
	293	.tp_dealloc = (destructor)Database_dealloc,
	294	.tp_init = (initproc)Database_init,
	295	.tp_doc = "Database object",
	296	.tp_methods = Database_methods,
	297	.tp_getset = Database_getsetters,
	298	.tp_repr = (reprfunc)Database_repr,
9cdf6c53	299	};
afb426df MT	300
	301	static PyObject* DatabaseEnumerator_new(PyTypeObject* type, PyObject* args, PyObject* kwds) {
	302	DatabaseEnumeratorObject* self = (DatabaseEnumeratorObject*)type->tp_alloc(type, 0);
	303
	304	return (PyObject*)self;
	305	}
	306
	307	static void DatabaseEnumerator_dealloc(DatabaseEnumeratorObject* self) {
	308	loc_database_enumerator_unref(self->enumerator);
	309
	310	Py_TYPE(self)->tp_free((PyObject* )self);
	311	}
	312
	313	static PyObject* DatabaseEnumerator_next(DatabaseEnumeratorObject* self) {
ccc7ab4e MT	314	// Enumerate all networks
	315	struct loc_network* network = loc_database_enumerator_next_network(self->enumerator);
	316	if (network) {
	317	PyObject* obj = new_network(&NetworkType, network);
	318	loc_network_unref(network);
	319
	320	return obj;
	321	}
	322
	323	// Enumerate all ASes
afb426df MT	324	struct loc_as* as = loc_database_enumerator_next_as(self->enumerator);
	325	if (as) {
	326	PyObject* obj = new_as(&ASType, as);
	327	loc_as_unref(as);
	328
	329	return obj;
	330	}
	331
	332	// Nothing found, that means the end
	333	PyErr_SetNone(PyExc_StopIteration);
	334	return NULL;
	335	}
	336
	337	PyTypeObject DatabaseEnumeratorType = {
	338	PyVarObject_HEAD_INIT(NULL, 0)
d42e1dcd MT	339	.tp_name = "location.DatabaseEnumerator",
	340	.tp_basicsize = sizeof(DatabaseEnumeratorObject),
	341	.tp_flags = Py_TPFLAGS_DEFAULT,
	342	.tp_alloc = PyType_GenericAlloc,
	343	.tp_new = DatabaseEnumerator_new,
	344	.tp_dealloc = (destructor)DatabaseEnumerator_dealloc,
	345	.tp_iter = PyObject_SelfIter,
	346	.tp_iternext = (iternextfunc)DatabaseEnumerator_next,
afb426df	347	};