[location/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 <libloc/libloc.h>
#include <libloc/as.h>
#include <libloc/as-list.h>
#include <libloc/database.h>

#include "locationmodule.h"
#include "as.h"
#include "country.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_verify(DatabaseObject* self, PyObject* args) {
	PyObject* public_key = NULL;
	FILE* f = NULL;

	// Parse arguments
	if (!PyArg_ParseTuple(args, "O", &public_key))
		return NULL;

	// Convert into FILE*
	int fd = PyObject_AsFileDescriptor(public_key);
	if (fd < 0)
		return NULL;

	// Re-open file descriptor
	f = fdopen(fd, "r");
	if (!f) {
		PyErr_SetFromErrno(PyExc_IOError);
		return NULL;
	}

	int r = loc_database_verify(self->db, f);

	if (r == 0)
		Py_RETURN_TRUE;

	Py_RETURN_FALSE;
}

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

	return PyUnicode_FromString(description);
}

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

	return PyUnicode_FromString(vendor);
}

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

	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_get_country(DatabaseObject* self, PyObject* args) {
	const char* country_code = NULL;

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

	struct loc_country* country;
	int r = loc_database_get_country(self->db, &country, country_code);
	if (r) {
		Py_RETURN_NONE;
	}

	PyObject* obj = new_country(&CountryType, country);
	loc_country_unref(country);

	return obj;
}

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_iterate_all(DatabaseObject* self,
		enum loc_database_enumerator_mode what, int family, int flags) {
	struct loc_database_enumerator* enumerator;

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

	// Set family
	if (family)
		loc_database_enumerator_set_family(enumerator, family);

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

	return obj;
}

static PyObject* Database_ases(DatabaseObject* self) {
	return Database_iterate_all(self, LOC_DB_ENUMERATE_ASES, AF_UNSPEC, 0);
}

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, 0);
	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_networks(DatabaseObject* self) {
	return Database_iterate_all(self, LOC_DB_ENUMERATE_NETWORKS, AF_UNSPEC, 0);
}

static PyObject* Database_networks_flattened(DatabaseObject *self) {
	return Database_iterate_all(self, LOC_DB_ENUMERATE_NETWORKS, AF_UNSPEC,
		LOC_DB_ENUMERATOR_FLAGS_FLATTEN);
}

static PyObject* Database_search_networks(DatabaseObject* self, PyObject* args, PyObject* kwargs) {
	char* kwlist[] = { "country_codes", "asns", "flags", "family", "flatten", NULL };
	PyObject* country_codes = NULL;
	PyObject* asn_list = NULL;
	int flags = 0;
	int family = 0;
	int flatten = 0;

	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|O!O!iip", kwlist,
			&PyList_Type, &country_codes, &PyList_Type, &asn_list, &flags, &family, &flatten))
		return NULL;

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

	// Set country code we are searching for
	if (country_codes) {
		struct loc_country_list* countries;
		r = loc_country_list_new(loc_ctx, &countries);
		if (r) {
			PyErr_SetString(PyExc_SystemError, "Could not create country list");
			return NULL;
		}

		for (int i = 0; i < PyList_Size(country_codes); i++) {
			PyObject* item = PyList_GetItem(country_codes, i);

			if (!PyUnicode_Check(item)) {
				PyErr_SetString(PyExc_TypeError, "Country codes must be strings");
				loc_country_list_unref(countries);
				return NULL;
			}

			const char* country_code = PyUnicode_AsUTF8(item);

			struct loc_country* country;
			r = loc_country_new(loc_ctx, &country, country_code);
			if (r) {
				if (r == -EINVAL) {
					PyErr_Format(PyExc_ValueError, "Invalid country code: %s", country_code);
				} else {
					PyErr_SetString(PyExc_SystemError, "Could not create country");
				}

				loc_country_list_unref(countries);
				return NULL;
			}

			// Append it to the list
			r = loc_country_list_append(countries, country);
			if (r) {
				PyErr_SetString(PyExc_SystemError, "Could not append country to the list");

				loc_country_list_unref(countries);
				loc_country_unref(country);
				return NULL;
			}

			loc_country_unref(country);
		}

		r = loc_database_enumerator_set_countries(enumerator, countries);
		if (r) {
			PyErr_SetFromErrno(PyExc_SystemError);

			loc_country_list_unref(countries);
			return NULL;
		}

		loc_country_list_unref(countries);
	}

	// Set the ASN we are searching for
	if (asn_list) {
		struct loc_as_list* asns;
		r = loc_as_list_new(loc_ctx, &asns);
		if (r) {
			PyErr_SetString(PyExc_SystemError, "Could not create AS list");
			return NULL;
		}

		for (int i = 0; i < PyList_Size(asn_list); i++) {
			PyObject* item = PyList_GetItem(asn_list, i);

			if (!PyLong_Check(item)) {
				PyErr_SetString(PyExc_TypeError, "ASNs must be numbers");

				loc_as_list_unref(asns);
				return NULL;
			}

			unsigned long number = PyLong_AsLong(item);

			struct loc_as* as;
			r = loc_as_new(loc_ctx, &as, number);
			if (r) {
				PyErr_SetString(PyExc_SystemError, "Could not create AS");

				loc_as_list_unref(asns);
				loc_as_unref(as);
				return NULL;
			}

			r = loc_as_list_append(asns, as);
			if (r) {
				PyErr_SetString(PyExc_SystemError, "Could not append AS to the list");

				loc_as_list_unref(asns);
				loc_as_unref(as);
				return NULL;
			}

			loc_as_unref(as);
		}

		r = loc_database_enumerator_set_asns(enumerator, asns);
		if (r) {
			PyErr_SetFromErrno(PyExc_SystemError);

			loc_as_list_unref(asns);
			return NULL;
		}

		loc_as_list_unref(asns);
	}

	// Set the flags we are searching for
	if (flags) {
		r = loc_database_enumerator_set_flag(enumerator, flags);

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

	// Set the family we are searching for
	if (family) {
		r = loc_database_enumerator_set_family(enumerator, family);

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

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

	return obj;
}

static PyObject* Database_countries(DatabaseObject* self) {
	return Database_iterate_all(self, LOC_DB_ENUMERATE_COUNTRIES, AF_UNSPEC, 0);
}

static PyObject* Database_list_bogons(DatabaseObject* self, PyObject* args, PyObject* kwargs) {
	char* kwlist[] = { "family", NULL };
	int family = AF_UNSPEC;

	// Parse arguments
	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|i", kwlist, &family))
		return NULL;

	return Database_iterate_all(self, LOC_DB_ENUMERATE_BOGONS, family, 0);
}

static struct PyMethodDef Database_methods[] = {
	{
		"get_as",
		(PyCFunction)Database_get_as,
		METH_VARARGS,
		NULL,
	},
	{
		"get_country",
		(PyCFunction)Database_get_country,
		METH_VARARGS,
		NULL,
	},
	{
		"list_bogons",
		(PyCFunction)Database_list_bogons,
		METH_VARARGS|METH_KEYWORDS,
		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,
	},
	{
		"verify",
		(PyCFunction)Database_verify,
		METH_VARARGS,
		NULL,
	},
	{ NULL },
};

static struct PyGetSetDef Database_getsetters[] = {
	{
		"ases",
		(getter)Database_ases,
		NULL,
		NULL,
		NULL,
	},
	{
		"countries",
		(getter)Database_countries,
		NULL,
		NULL,
		NULL,
	},
	{
		"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,
	},
	{
		"networks",
		(getter)Database_networks,
		NULL,
		NULL,
		NULL,
	},
	{
		"networks_flattened",
		(getter)Database_networks_flattened,
		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) {
	struct loc_network* network = NULL;

	// Enumerate all networks
	int r = loc_database_enumerator_next_network(self->enumerator, &network);
	if (r) {
		PyErr_SetFromErrno(PyExc_ValueError);
		return NULL;
	}

	// A network was found
	if (network) {
		PyObject* obj = new_network(&NetworkType, network);
		loc_network_unref(network);

		return obj;
	}

	// Enumerate all ASes
	struct loc_as* as = NULL;

	r = loc_database_enumerator_next_as(self->enumerator, &as);
	if (r) {
		PyErr_SetFromErrno(PyExc_ValueError);
		return NULL;
	}

	if (as) {
		PyObject* obj = new_as(&ASType, as);
		loc_as_unref(as);

		return obj;
	}

	// Enumerate all countries
	struct loc_country* country = NULL;

	r = loc_database_enumerator_next_country(self->enumerator, &country);
	if (r) {
		PyErr_SetFromErrno(PyExc_ValueError);
		return NULL;
	}

	if (country) {
		PyObject* obj = new_country(&CountryType, country);
		loc_country_unref(country);

		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
c12a1385 MT	19	#include <libloc/libloc.h>
	20	#include <libloc/as.h>
	21	#include <libloc/as-list.h>
	22	#include <libloc/database.h>
9fc7f001	23
1da9cd39	24	#include "locationmodule.h"
86ca7ef7	25	#include "as.h"
7c922e9c	26	#include "country.h"
9cdf6c53	27	#include "database.h"
31edab76	28	#include "network.h"
9cdf6c53 MT	29
9cdf6c53 MT	30	static PyObject* Database_new(PyTypeObject* type, PyObject* args, PyObject* kwds) {
9cdf6c53	31	DatabaseObject* self = (DatabaseObject*)type->tp_alloc(type, 0);
9cdf6c53 MT	32
	33	return (PyObject*)self;
	34	}
	35
	36	static void Database_dealloc(DatabaseObject* self) {
	37	if (self->db)
	38	loc_database_unref(self->db);
	39
6fd96715 MT	40	if (self->path)
	41	free(self->path);
	42
9cdf6c53 MT	43	Py_TYPE(self)->tp_free((PyObject* )self);
	44	}
	45
	46	static int Database_init(DatabaseObject* self, PyObject* args, PyObject* kwargs) {
	47	const char* path = NULL;
	48
	49	if (!PyArg_ParseTuple(args, "s", &path))
	50	return -1;
	51
6fd96715 MT	52	self->path = strdup(path);
6fd96715 MT	53
9cdf6c53	54	// Open the file for reading
6fd96715	55	FILE* f = fopen(self->path, "r");
774eea26 MT	56	if (!f) {
774eea26 MT	57	PyErr_SetFromErrno(PyExc_IOError);
9cdf6c53	58	return -1;
774eea26	59	}
9cdf6c53 MT	60
9cdf6c53 MT	61	// Load the database
38e07ee0	62	int r = loc_database_new(loc_ctx, &self->db, f);
9cdf6c53 MT	63	fclose(f);
	64
	65	// Return on any errors
	66	if (r)
	67	return -1;
	68
	69	return 0;
	70	}
	71
6fd96715 MT	72	static PyObject* Database_repr(DatabaseObject* self) {
	73	return PyUnicode_FromFormat("<Database %s>", self->path);
	74	}
	75
726f9984 MT	76	static PyObject* Database_verify(DatabaseObject* self, PyObject* args) {
	77	PyObject* public_key = NULL;
	78	FILE* f = NULL;
	79
	80	// Parse arguments
	81	if (!PyArg_ParseTuple(args, "O", &public_key))
	82	return NULL;
	83
	84	// Convert into FILE*
	85	int fd = PyObject_AsFileDescriptor(public_key);
	86	if (fd < 0)
	87	return NULL;
	88
	89	// Re-open file descriptor
	90	f = fdopen(fd, "r");
	91	if (!f) {
	92	PyErr_SetFromErrno(PyExc_IOError);
	93	return NULL;
	94	}
	95
	96	int r = loc_database_verify(self->db, f);
b1720435 MT	97
	98	if (r == 0)
	99	Py_RETURN_TRUE;
	100
	101	Py_RETURN_FALSE;
	102	}
	103
d99b0256 MT	104	static PyObject* Database_get_description(DatabaseObject* self) {
d99b0256 MT	105	const char* description = loc_database_get_description(self->db);
f35adb4c MT	106	if (!description)
f35adb4c MT	107	Py_RETURN_NONE;
d99b0256 MT	108
	109	return PyUnicode_FromString(description);
	110	}
	111
	112	static PyObject* Database_get_vendor(DatabaseObject* self) {
	113	const char* vendor = loc_database_get_vendor(self->db);
f35adb4c MT	114	if (!vendor)
f35adb4c MT	115	Py_RETURN_NONE;
d99b0256 MT	116
	117	return PyUnicode_FromString(vendor);
	118	}
	119
4bf49d00 MT	120	static PyObject* Database_get_license(DatabaseObject* self) {
4bf49d00 MT	121	const char* license = loc_database_get_license(self->db);
f35adb4c MT	122	if (!license)
f35adb4c MT	123	Py_RETURN_NONE;
4bf49d00 MT	124
	125	return PyUnicode_FromString(license);
	126	}
	127
53524b2d MT	128	static PyObject* Database_get_created_at(DatabaseObject* self) {
	129	time_t created_at = loc_database_created_at(self->db);
	130
	131	return PyLong_FromLong(created_at);
	132	}
	133
86ca7ef7 MT	134	static PyObject* Database_get_as(DatabaseObject* self, PyObject* args) {
	135	struct loc_as* as = NULL;
	136	uint32_t number = 0;
	137
	138	if (!PyArg_ParseTuple(args, "i", &number))
	139	return NULL;
	140
	141	// Try to retrieve the AS
	142	int r = loc_database_get_as(self->db, &as, number);
86ca7ef7	143
4a0a0f7e MT	144	// We got an AS
4a0a0f7e MT	145	if (r == 0) {
86ca7ef7 MT	146	PyObject* obj = new_as(&ASType, as);
	147	loc_as_unref(as);
	148
	149	return obj;
86ca7ef7 MT	150
86ca7ef7 MT	151	// Nothing found
4a0a0f7e MT	152	} else if (r == 1) {
	153	Py_RETURN_NONE;
	154	}
	155
	156	// Unexpected error
	157	return NULL;
86ca7ef7 MT	158	}
86ca7ef7 MT	159
7c922e9c MT	160	static PyObject* Database_get_country(DatabaseObject* self, PyObject* args) {
	161	const char* country_code = NULL;
	162
	163	if (!PyArg_ParseTuple(args, "s", &country_code))
	164	return NULL;
	165
	166	struct loc_country* country;
	167	int r = loc_database_get_country(self->db, &country, country_code);
	168	if (r) {
	169	Py_RETURN_NONE;
	170	}
	171
	172	PyObject* obj = new_country(&CountryType, country);
	173	loc_country_unref(country);
	174
	175	return obj;
	176	}
	177
31edab76 MT	178	static PyObject* Database_lookup(DatabaseObject* self, PyObject* args) {
	179	struct loc_network* network = NULL;
	180	const char* address = NULL;
	181
	182	if (!PyArg_ParseTuple(args, "s", &address))
	183	return NULL;
	184
	185	// Try to retrieve a matching network
	186	int r = loc_database_lookup_from_string(self->db, address, &network);
	187
	188	// We got a network
	189	if (r == 0) {
	190	PyObject* obj = new_network(&NetworkType, network);
	191	loc_network_unref(network);
	192
	193	return obj;
	194
	195	// Nothing found
	196	} else if (r == 1) {
	197	Py_RETURN_NONE;
927e82f2 MT	198
	199	// Invalid input
	200	} else if (r == -EINVAL) {
	201	PyErr_Format(PyExc_ValueError, "Invalid IP address: %s", address);
	202	return NULL;
31edab76 MT	203	}
	204
	205	// Unexpected error
	206	return NULL;
	207	}
	208
afb426df MT	209	static PyObject* new_database_enumerator(PyTypeObject* type, struct loc_database_enumerator* enumerator) {
	210	DatabaseEnumeratorObject* self = (DatabaseEnumeratorObject*)type->tp_alloc(type, 0);
	211	if (self) {
	212	self->enumerator = loc_database_enumerator_ref(enumerator);
	213	}
	214
	215	return (PyObject*)self;
	216	}
	217
d34b669c MT	218	static PyObject* Database_iterate_all(DatabaseObject* self,
d34b669c MT	219	enum loc_database_enumerator_mode what, int family, int flags) {
a68a46f5 MT	220	struct loc_database_enumerator* enumerator;
a68a46f5 MT	221
681ff05c	222	int r = loc_database_enumerator_new(&enumerator, self->db, what, flags);
a68a46f5 MT	223	if (r) {
	224	PyErr_SetFromErrno(PyExc_SystemError);
	225	return NULL;
	226	}
	227
d34b669c MT	228	// Set family
	229	if (family)
	230	loc_database_enumerator_set_family(enumerator, family);
	231
a68a46f5 MT	232	PyObject* obj = new_database_enumerator(&DatabaseEnumeratorType, enumerator);
	233	loc_database_enumerator_unref(enumerator);
	234
	235	return obj;
	236	}
	237
	238	static PyObject* Database_ases(DatabaseObject* self) {
d34b669c	239	return Database_iterate_all(self, LOC_DB_ENUMERATE_ASES, AF_UNSPEC, 0);
a68a46f5 MT	240	}
a68a46f5 MT	241
afb426df MT	242	static PyObject* Database_search_as(DatabaseObject* self, PyObject* args) {
	243	const char* string = NULL;
	244
	245	if (!PyArg_ParseTuple(args, "s", &string))
	246	return NULL;
	247
	248	struct loc_database_enumerator* enumerator;
	249
681ff05c	250	int r = loc_database_enumerator_new(&enumerator, self->db, LOC_DB_ENUMERATE_ASES, 0);
afb426df MT	251	if (r) {
	252	PyErr_SetFromErrno(PyExc_SystemError);
	253	return NULL;
	254	}
	255
	256	// Search string we are searching for
	257	loc_database_enumerator_set_string(enumerator, string);
	258
	259	PyObject* obj = new_database_enumerator(&DatabaseEnumeratorType, enumerator);
	260	loc_database_enumerator_unref(enumerator);
	261
	262	return obj;
	263	}
	264
a68a46f5	265	static PyObject* Database_networks(DatabaseObject* self) {
d34b669c	266	return Database_iterate_all(self, LOC_DB_ENUMERATE_NETWORKS, AF_UNSPEC, 0);
681ff05c MT	267	}
	268
	269	static PyObject* Database_networks_flattened(DatabaseObject *self) {
d34b669c MT	270	return Database_iterate_all(self, LOC_DB_ENUMERATE_NETWORKS, AF_UNSPEC,
d34b669c MT	271	LOC_DB_ENUMERATOR_FLAGS_FLATTEN);
a68a46f5 MT	272	}
a68a46f5 MT	273
ccc7ab4e	274	static PyObject* Database_search_networks(DatabaseObject* self, PyObject* args, PyObject* kwargs) {
84a2f0c2	275	char* kwlist[] = { "country_codes", "asns", "flags", "family", "flatten", NULL };
e646a8f3	276	PyObject* country_codes = NULL;
84a2f0c2	277	PyObject* asn_list = NULL;
bbdb2e0a	278	int flags = 0;
44e5ef71	279	int family = 0;
c242f732	280	int flatten = 0;
ccc7ab4e	281
84a2f0c2 MT	282	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "\|O!O!iip", kwlist,
84a2f0c2 MT	283	&PyList_Type, &country_codes, &PyList_Type, &asn_list, &flags, &family, &flatten))
ccc7ab4e MT	284	return NULL;
	285
	286	struct loc_database_enumerator* enumerator;
c242f732 MT	287	int r = loc_database_enumerator_new(&enumerator, self->db, LOC_DB_ENUMERATE_NETWORKS,
c242f732 MT	288	(flatten) ? LOC_DB_ENUMERATOR_FLAGS_FLATTEN : 0);
ccc7ab4e MT	289	if (r) {
	290	PyErr_SetFromErrno(PyExc_SystemError);
	291	return NULL;
	292	}
	293
	294	// Set country code we are searching for
e646a8f3 MT	295	if (country_codes) {
	296	struct loc_country_list* countries;
	297	r = loc_country_list_new(loc_ctx, &countries);
ccc7ab4e	298	if (r) {
e646a8f3	299	PyErr_SetString(PyExc_SystemError, "Could not create country list");
ccc7ab4e MT	300	return NULL;
ccc7ab4e MT	301	}
e646a8f3	302
fa9eb94e	303	for (int i = 0; i < PyList_Size(country_codes); i++) {
e646a8f3 MT	304	PyObject* item = PyList_GetItem(country_codes, i);
	305
	306	if (!PyUnicode_Check(item)) {
	307	PyErr_SetString(PyExc_TypeError, "Country codes must be strings");
	308	loc_country_list_unref(countries);
	309	return NULL;
	310	}
	311
	312	const char* country_code = PyUnicode_AsUTF8(item);
	313
	314	struct loc_country* country;
	315	r = loc_country_new(loc_ctx, &country, country_code);
	316	if (r) {
	317	if (r == -EINVAL) {
	318	PyErr_Format(PyExc_ValueError, "Invalid country code: %s", country_code);
	319	} else {
	320	PyErr_SetString(PyExc_SystemError, "Could not create country");
	321	}
	322
	323	loc_country_list_unref(countries);
	324	return NULL;
	325	}
	326
	327	// Append it to the list
	328	r = loc_country_list_append(countries, country);
	329	if (r) {
	330	PyErr_SetString(PyExc_SystemError, "Could not append country to the list");
	331
	332	loc_country_list_unref(countries);
	333	loc_country_unref(country);
	334	return NULL;
	335	}
	336
	337	loc_country_unref(country);
	338	}
	339
c98ebf8a MT	340	r = loc_database_enumerator_set_countries(enumerator, countries);
	341	if (r) {
	342	PyErr_SetFromErrno(PyExc_SystemError);
	343
a1a00053	344	loc_country_list_unref(countries);
c98ebf8a MT	345	return NULL;
c98ebf8a MT	346	}
e646a8f3	347
e646a8f3	348	loc_country_list_unref(countries);
ccc7ab4e MT	349	}
	350
	351	// Set the ASN we are searching for
84a2f0c2 MT	352	if (asn_list) {
	353	struct loc_as_list* asns;
	354	r = loc_as_list_new(loc_ctx, &asns);
	355	if (r) {
	356	PyErr_SetString(PyExc_SystemError, "Could not create AS list");
	357	return NULL;
	358	}
	359
fa9eb94e	360	for (int i = 0; i < PyList_Size(asn_list); i++) {
84a2f0c2 MT	361	PyObject* item = PyList_GetItem(asn_list, i);
	362
	363	if (!PyLong_Check(item)) {
	364	PyErr_SetString(PyExc_TypeError, "ASNs must be numbers");
	365
	366	loc_as_list_unref(asns);
	367	return NULL;
	368	}
	369
	370	unsigned long number = PyLong_AsLong(item);
ccc7ab4e	371
84a2f0c2 MT	372	struct loc_as* as;
	373	r = loc_as_new(loc_ctx, &as, number);
	374	if (r) {
	375	PyErr_SetString(PyExc_SystemError, "Could not create AS");
	376
	377	loc_as_list_unref(asns);
	378	loc_as_unref(as);
	379	return NULL;
	380	}
	381
	382	r = loc_as_list_append(asns, as);
	383	if (r) {
	384	PyErr_SetString(PyExc_SystemError, "Could not append AS to the list");
	385
	386	loc_as_list_unref(asns);
	387	loc_as_unref(as);
	388	return NULL;
	389	}
	390
	391	loc_as_unref(as);
	392	}
	393
	394	r = loc_database_enumerator_set_asns(enumerator, asns);
ccc7ab4e MT	395	if (r) {
ccc7ab4e MT	396	PyErr_SetFromErrno(PyExc_SystemError);
84a2f0c2 MT	397
84a2f0c2 MT	398	loc_as_list_unref(asns);
ccc7ab4e MT	399	return NULL;
ccc7ab4e MT	400	}
84a2f0c2 MT	401
84a2f0c2 MT	402	loc_as_list_unref(asns);
ccc7ab4e MT	403	}
ccc7ab4e MT	404
bbdb2e0a MT	405	// Set the flags we are searching for
	406	if (flags) {
	407	r = loc_database_enumerator_set_flag(enumerator, flags);
	408
	409	if (r) {
	410	PyErr_SetFromErrno(PyExc_SystemError);
	411	return NULL;
	412	}
	413	}
	414
44e5ef71 MT	415	// Set the family we are searching for
	416	if (family) {
	417	r = loc_database_enumerator_set_family(enumerator, family);
	418
	419	if (r) {
	420	PyErr_SetFromErrno(PyExc_SystemError);
	421	return NULL;
	422	}
	423	}
	424
ccc7ab4e MT	425	PyObject* obj = new_database_enumerator(&DatabaseEnumeratorType, enumerator);
	426	loc_database_enumerator_unref(enumerator);
	427
	428	return obj;
	429	}
	430
fa9a3663	431	static PyObject* Database_countries(DatabaseObject* self) {
d34b669c	432	return Database_iterate_all(self, LOC_DB_ENUMERATE_COUNTRIES, AF_UNSPEC, 0);
fa9a3663 MT	433	}
fa9a3663 MT	434
8c37d8a7 MT	435	static PyObject* Database_list_bogons(DatabaseObject* self, PyObject* args, PyObject* kwargs) {
	436	char* kwlist[] = { "family", NULL };
	437	int family = AF_UNSPEC;
	438
	439	// Parse arguments
	440	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "\|i", kwlist, &family))
	441	return NULL;
	442
	443	return Database_iterate_all(self, LOC_DB_ENUMERATE_BOGONS, family, 0);
	444	}
	445
9cdf6c53	446	static struct PyMethodDef Database_methods[] = {
86ca7ef7 MT	447	{
	448	"get_as",
	449	(PyCFunction)Database_get_as,
	450	METH_VARARGS,
	451	NULL,
	452	},
7c922e9c MT	453	{
	454	"get_country",
	455	(PyCFunction)Database_get_country,
	456	METH_VARARGS,
	457	NULL,
	458	},
8c37d8a7 MT	459	{
	460	"list_bogons",
	461	(PyCFunction)Database_list_bogons,
	462	METH_VARARGS\|METH_KEYWORDS,
	463	NULL,
	464	},
31edab76 MT	465	{
	466	"lookup",
	467	(PyCFunction)Database_lookup,
	468	METH_VARARGS,
	469	NULL,
	470	},
afb426df MT	471	{
	472	"search_as",
	473	(PyCFunction)Database_search_as,
	474	METH_VARARGS,
	475	NULL,
	476	},
ccc7ab4e MT	477	{
	478	"search_networks",
	479	(PyCFunction)Database_search_networks,
	480	METH_VARARGS\|METH_KEYWORDS,
	481	NULL,
	482	},
b1720435 MT	483	{
	484	"verify",
	485	(PyCFunction)Database_verify,
726f9984	486	METH_VARARGS,
b1720435 MT	487	NULL,
b1720435 MT	488	},
9cdf6c53 MT	489	{ NULL },
	490	};
	491
d99b0256	492	static struct PyGetSetDef Database_getsetters[] = {
a68a46f5 MT	493	{
	494	"ases",
	495	(getter)Database_ases,
	496	NULL,
	497	NULL,
	498	NULL,
	499	},
fa9a3663 MT	500	{
	501	"countries",
	502	(getter)Database_countries,
	503	NULL,
	504	NULL,
	505	NULL,
	506	},
53524b2d MT	507	{
	508	"created_at",
	509	(getter)Database_get_created_at,
	510	NULL,
	511	NULL,
	512	NULL,
	513	},
d99b0256 MT	514	{
	515	"description",
	516	(getter)Database_get_description,
	517	NULL,
	518	NULL,
	519	NULL,
	520	},
4bf49d00 MT	521	{
	522	"license",
	523	(getter)Database_get_license,
	524	NULL,
	525	NULL,
	526	NULL,
	527	},
a68a46f5 MT	528	{
	529	"networks",
	530	(getter)Database_networks,
	531	NULL,
	532	NULL,
	533	NULL,
	534	},
681ff05c MT	535	{
	536	"networks_flattened",
	537	(getter)Database_networks_flattened,
	538	NULL,
	539	NULL,
	540	NULL,
	541	},
d99b0256 MT	542	{
	543	"vendor",
	544	(getter)Database_get_vendor,
	545	NULL,
	546	NULL,
53524b2d	547	NULL,
d99b0256 MT	548	},
	549	{ NULL },
	550	};
	551
9cdf6c53 MT	552	PyTypeObject DatabaseType = {
9cdf6c53 MT	553	PyVarObject_HEAD_INIT(NULL, 0)
d42e1dcd MT	554	.tp_name = "location.Database",
	555	.tp_basicsize = sizeof(DatabaseObject),
	556	.tp_flags = Py_TPFLAGS_DEFAULT\|Py_TPFLAGS_BASETYPE,
	557	.tp_new = Database_new,
	558	.tp_dealloc = (destructor)Database_dealloc,
	559	.tp_init = (initproc)Database_init,
	560	.tp_doc = "Database object",
	561	.tp_methods = Database_methods,
	562	.tp_getset = Database_getsetters,
	563	.tp_repr = (reprfunc)Database_repr,
9cdf6c53	564	};
afb426df MT	565
	566	static PyObject* DatabaseEnumerator_new(PyTypeObject* type, PyObject* args, PyObject* kwds) {
	567	DatabaseEnumeratorObject* self = (DatabaseEnumeratorObject*)type->tp_alloc(type, 0);
	568
	569	return (PyObject*)self;
	570	}
	571
	572	static void DatabaseEnumerator_dealloc(DatabaseEnumeratorObject* self) {
	573	loc_database_enumerator_unref(self->enumerator);
	574
	575	Py_TYPE(self)->tp_free((PyObject* )self);
	576	}
	577
	578	static PyObject* DatabaseEnumerator_next(DatabaseEnumeratorObject* self) {
15f79e2d MT	579	struct loc_network* network = NULL;
15f79e2d MT	580
ccc7ab4e	581	// Enumerate all networks
15f79e2d MT	582	int r = loc_database_enumerator_next_network(self->enumerator, &network);
15f79e2d MT	583	if (r) {
20bb6d0c	584	PyErr_SetFromErrno(PyExc_ValueError);
15f79e2d MT	585	return NULL;
	586	}
	587
	588	// A network was found
ccc7ab4e MT	589	if (network) {
	590	PyObject* obj = new_network(&NetworkType, network);
	591	loc_network_unref(network);
	592
	593	return obj;
	594	}
	595
	596	// Enumerate all ASes
15f79e2d MT	597	struct loc_as* as = NULL;
	598
	599	r = loc_database_enumerator_next_as(self->enumerator, &as);
	600	if (r) {
20bb6d0c	601	PyErr_SetFromErrno(PyExc_ValueError);
15f79e2d MT	602	return NULL;
	603	}
	604
afb426df MT	605	if (as) {
	606	PyObject* obj = new_as(&ASType, as);
	607	loc_as_unref(as);
	608
	609	return obj;
	610	}
	611
fa9a3663 MT	612	// Enumerate all countries
	613	struct loc_country* country = NULL;
	614
	615	r = loc_database_enumerator_next_country(self->enumerator, &country);
	616	if (r) {
	617	PyErr_SetFromErrno(PyExc_ValueError);
	618	return NULL;
	619	}
	620
	621	if (country) {
	622	PyObject* obj = new_country(&CountryType, country);
	623	loc_country_unref(country);
	624
	625	return obj;
	626	}
	627
afb426df MT	628	// Nothing found, that means the end
	629	PyErr_SetNone(PyExc_StopIteration);
	630	return NULL;
	631	}
	632
	633	PyTypeObject DatabaseEnumeratorType = {
	634	PyVarObject_HEAD_INIT(NULL, 0)
d42e1dcd MT	635	.tp_name = "location.DatabaseEnumerator",
	636	.tp_basicsize = sizeof(DatabaseEnumeratorObject),
	637	.tp_flags = Py_TPFLAGS_DEFAULT,
	638	.tp_alloc = PyType_GenericAlloc,
	639	.tp_new = DatabaseEnumerator_new,
	640	.tp_dealloc = (destructor)DatabaseEnumerator_dealloc,
	641	.tp_iter = PyObject_SelfIter,
	642	.tp_iternext = (iternextfunc)DatabaseEnumerator_next,
afb426df	643	};