[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 <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;
	FILE* f = NULL;

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

	// Copy path
	self->path = strdup(path);
	if (!self->path)
		goto ERROR;

	// Open the file for reading
	f = fopen(self->path, "r");
	if (!f)
		goto ERROR;

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

	fclose(f);
	return 0;

ERROR:
	if (f)
		fclose(f);

	PyErr_SetFromErrno(PyExc_OSError);
	return -1;
}

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
	if (!errno)
		Py_RETURN_NONE;

	// Handle any errors
	switch (errno) {
		case EINVAL:
			PyErr_Format(PyExc_ValueError, "Invalid IP address: %s", address);

		default:
			PyErr_SetFromErrno(PyExc_OSError);
	}

	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_SetFromErrno(PyExc_OSError);
			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_SetFromErrno(PyExc_OSError);

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

			r = loc_as_list_append(asns, as);
			if (r) {
				PyErr_SetFromErrno(PyExc_OSError);

				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_OSError);

			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_OSError);
			return NULL;
		}
	}

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

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