[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
	} 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_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;
	202
	203	// Nothing found
	204	} else if (r == 1) {
	205	Py_RETURN_NONE;
927e82f2 MT	206
	207	// Invalid input
	208	} else if (r == -EINVAL) {
	209	PyErr_Format(PyExc_ValueError, "Invalid IP address: %s", address);
	210	return NULL;
31edab76 MT	211	}
	212
	213	// Unexpected error
	214	return NULL;
	215	}
	216
afb426df MT	217	static PyObject* new_database_enumerator(PyTypeObject* type, struct loc_database_enumerator* enumerator) {
	218	DatabaseEnumeratorObject* self = (DatabaseEnumeratorObject*)type->tp_alloc(type, 0);
	219	if (self) {
	220	self->enumerator = loc_database_enumerator_ref(enumerator);
	221	}
	222
	223	return (PyObject*)self;
	224	}
	225
d34b669c MT	226	static PyObject* Database_iterate_all(DatabaseObject* self,
d34b669c MT	227	enum loc_database_enumerator_mode what, int family, int flags) {
a68a46f5 MT	228	struct loc_database_enumerator* enumerator;
a68a46f5 MT	229
681ff05c	230	int r = loc_database_enumerator_new(&enumerator, self->db, what, flags);
a68a46f5 MT	231	if (r) {
	232	PyErr_SetFromErrno(PyExc_SystemError);
	233	return NULL;
	234	}
	235
d34b669c MT	236	// Set family
	237	if (family)
	238	loc_database_enumerator_set_family(enumerator, family);
	239
a68a46f5 MT	240	PyObject* obj = new_database_enumerator(&DatabaseEnumeratorType, enumerator);
	241	loc_database_enumerator_unref(enumerator);
	242
	243	return obj;
	244	}
	245
	246	static PyObject* Database_ases(DatabaseObject* self) {
d34b669c	247	return Database_iterate_all(self, LOC_DB_ENUMERATE_ASES, AF_UNSPEC, 0);
a68a46f5 MT	248	}
a68a46f5 MT	249
afb426df MT	250	static PyObject* Database_search_as(DatabaseObject* self, PyObject* args) {
	251	const char* string = NULL;
	252
	253	if (!PyArg_ParseTuple(args, "s", &string))
	254	return NULL;
	255
	256	struct loc_database_enumerator* enumerator;
	257
681ff05c	258	int r = loc_database_enumerator_new(&enumerator, self->db, LOC_DB_ENUMERATE_ASES, 0);
afb426df MT	259	if (r) {
	260	PyErr_SetFromErrno(PyExc_SystemError);
	261	return NULL;
	262	}
	263
	264	// Search string we are searching for
	265	loc_database_enumerator_set_string(enumerator, string);
	266
	267	PyObject* obj = new_database_enumerator(&DatabaseEnumeratorType, enumerator);
	268	loc_database_enumerator_unref(enumerator);
	269
	270	return obj;
	271	}
	272
a68a46f5	273	static PyObject* Database_networks(DatabaseObject* self) {
d34b669c	274	return Database_iterate_all(self, LOC_DB_ENUMERATE_NETWORKS, AF_UNSPEC, 0);
681ff05c MT	275	}
	276
	277	static PyObject* Database_networks_flattened(DatabaseObject *self) {
d34b669c MT	278	return Database_iterate_all(self, LOC_DB_ENUMERATE_NETWORKS, AF_UNSPEC,
d34b669c MT	279	LOC_DB_ENUMERATOR_FLAGS_FLATTEN);
a68a46f5 MT	280	}
a68a46f5 MT	281
ccc7ab4e	282	static PyObject* Database_search_networks(DatabaseObject* self, PyObject* args, PyObject* kwargs) {
84a2f0c2	283	char* kwlist[] = { "country_codes", "asns", "flags", "family", "flatten", NULL };
e646a8f3	284	PyObject* country_codes = NULL;
84a2f0c2	285	PyObject* asn_list = NULL;
bbdb2e0a	286	int flags = 0;
44e5ef71	287	int family = 0;
c242f732	288	int flatten = 0;
ccc7ab4e	289
84a2f0c2 MT	290	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "\|O!O!iip", kwlist,
84a2f0c2 MT	291	&PyList_Type, &country_codes, &PyList_Type, &asn_list, &flags, &family, &flatten))
ccc7ab4e MT	292	return NULL;
	293
	294	struct loc_database_enumerator* enumerator;
c242f732 MT	295	int r = loc_database_enumerator_new(&enumerator, self->db, LOC_DB_ENUMERATE_NETWORKS,
c242f732 MT	296	(flatten) ? LOC_DB_ENUMERATOR_FLAGS_FLATTEN : 0);
ccc7ab4e MT	297	if (r) {
	298	PyErr_SetFromErrno(PyExc_SystemError);
	299	return NULL;
	300	}
	301
	302	// Set country code we are searching for
e646a8f3 MT	303	if (country_codes) {
	304	struct loc_country_list* countries;
	305	r = loc_country_list_new(loc_ctx, &countries);
ccc7ab4e	306	if (r) {
e646a8f3	307	PyErr_SetString(PyExc_SystemError, "Could not create country list");
ccc7ab4e MT	308	return NULL;
ccc7ab4e MT	309	}
e646a8f3	310
fa9eb94e	311	for (int i = 0; i < PyList_Size(country_codes); i++) {
e646a8f3 MT	312	PyObject* item = PyList_GetItem(country_codes, i);
	313
	314	if (!PyUnicode_Check(item)) {
	315	PyErr_SetString(PyExc_TypeError, "Country codes must be strings");
	316	loc_country_list_unref(countries);
	317	return NULL;
	318	}
	319
	320	const char* country_code = PyUnicode_AsUTF8(item);
	321
	322	struct loc_country* country;
	323	r = loc_country_new(loc_ctx, &country, country_code);
	324	if (r) {
	325	if (r == -EINVAL) {
	326	PyErr_Format(PyExc_ValueError, "Invalid country code: %s", country_code);
	327	} else {
	328	PyErr_SetString(PyExc_SystemError, "Could not create country");
	329	}
	330
	331	loc_country_list_unref(countries);
	332	return NULL;
	333	}
	334
	335	// Append it to the list
	336	r = loc_country_list_append(countries, country);
	337	if (r) {
	338	PyErr_SetString(PyExc_SystemError, "Could not append country to the list");
	339
	340	loc_country_list_unref(countries);
	341	loc_country_unref(country);
	342	return NULL;
	343	}
	344
	345	loc_country_unref(country);
	346	}
	347
c98ebf8a MT	348	r = loc_database_enumerator_set_countries(enumerator, countries);
	349	if (r) {
	350	PyErr_SetFromErrno(PyExc_SystemError);
	351
a1a00053	352	loc_country_list_unref(countries);
c98ebf8a MT	353	return NULL;
c98ebf8a MT	354	}
e646a8f3	355
e646a8f3	356	loc_country_list_unref(countries);
ccc7ab4e MT	357	}
	358
	359	// Set the ASN we are searching for
84a2f0c2 MT	360	if (asn_list) {
	361	struct loc_as_list* asns;
	362	r = loc_as_list_new(loc_ctx, &asns);
	363	if (r) {
198e382c	364	PyErr_SetFromErrno(PyExc_OSError);
84a2f0c2 MT	365	return NULL;
	366	}
	367
fa9eb94e	368	for (int i = 0; i < PyList_Size(asn_list); i++) {
84a2f0c2 MT	369	PyObject* item = PyList_GetItem(asn_list, i);
	370
	371	if (!PyLong_Check(item)) {
	372	PyErr_SetString(PyExc_TypeError, "ASNs must be numbers");
	373
	374	loc_as_list_unref(asns);
	375	return NULL;
	376	}
	377
	378	unsigned long number = PyLong_AsLong(item);
ccc7ab4e	379
84a2f0c2 MT	380	struct loc_as* as;
	381	r = loc_as_new(loc_ctx, &as, number);
	382	if (r) {
198e382c	383	PyErr_SetFromErrno(PyExc_OSError);
84a2f0c2 MT	384
	385	loc_as_list_unref(asns);
	386	loc_as_unref(as);
	387	return NULL;
	388	}
	389
	390	r = loc_as_list_append(asns, as);
	391	if (r) {
198e382c	392	PyErr_SetFromErrno(PyExc_OSError);
84a2f0c2 MT	393
	394	loc_as_list_unref(asns);
	395	loc_as_unref(as);
	396	return NULL;
	397	}
	398
	399	loc_as_unref(as);
	400	}
	401
	402	r = loc_database_enumerator_set_asns(enumerator, asns);
ccc7ab4e	403	if (r) {
198e382c	404	PyErr_SetFromErrno(PyExc_OSError);
84a2f0c2 MT	405
84a2f0c2 MT	406	loc_as_list_unref(asns);
ccc7ab4e MT	407	return NULL;
ccc7ab4e MT	408	}
84a2f0c2 MT	409
84a2f0c2 MT	410	loc_as_list_unref(asns);
ccc7ab4e MT	411	}
ccc7ab4e MT	412
bbdb2e0a MT	413	// Set the flags we are searching for
	414	if (flags) {
	415	r = loc_database_enumerator_set_flag(enumerator, flags);
	416
	417	if (r) {
198e382c	418	PyErr_SetFromErrno(PyExc_OSError);
bbdb2e0a MT	419	return NULL;
	420	}
	421	}
	422
44e5ef71 MT	423	// Set the family we are searching for
	424	if (family) {
	425	r = loc_database_enumerator_set_family(enumerator, family);
	426
	427	if (r) {
198e382c	428	PyErr_SetFromErrno(PyExc_OSError);
44e5ef71 MT	429	return NULL;
	430	}
	431	}
	432
ccc7ab4e MT	433	PyObject* obj = new_database_enumerator(&DatabaseEnumeratorType, enumerator);
	434	loc_database_enumerator_unref(enumerator);
	435
	436	return obj;
	437	}
	438
fa9a3663	439	static PyObject* Database_countries(DatabaseObject* self) {
d34b669c	440	return Database_iterate_all(self, LOC_DB_ENUMERATE_COUNTRIES, AF_UNSPEC, 0);
fa9a3663 MT	441	}
fa9a3663 MT	442
8c37d8a7 MT	443	static PyObject* Database_list_bogons(DatabaseObject* self, PyObject* args, PyObject* kwargs) {
	444	char* kwlist[] = { "family", NULL };
	445	int family = AF_UNSPEC;
	446
	447	// Parse arguments
	448	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "\|i", kwlist, &family))
	449	return NULL;
	450
	451	return Database_iterate_all(self, LOC_DB_ENUMERATE_BOGONS, family, 0);
	452	}
	453
9cdf6c53	454	static struct PyMethodDef Database_methods[] = {
86ca7ef7 MT	455	{
	456	"get_as",
	457	(PyCFunction)Database_get_as,
	458	METH_VARARGS,
	459	NULL,
	460	},
7c922e9c MT	461	{
	462	"get_country",
	463	(PyCFunction)Database_get_country,
	464	METH_VARARGS,
	465	NULL,
	466	},
8c37d8a7 MT	467	{
	468	"list_bogons",
	469	(PyCFunction)Database_list_bogons,
	470	METH_VARARGS\|METH_KEYWORDS,
	471	NULL,
	472	},
31edab76 MT	473	{
	474	"lookup",
	475	(PyCFunction)Database_lookup,
	476	METH_VARARGS,
	477	NULL,
	478	},
afb426df MT	479	{
	480	"search_as",
	481	(PyCFunction)Database_search_as,
	482	METH_VARARGS,
	483	NULL,
	484	},
ccc7ab4e MT	485	{
	486	"search_networks",
	487	(PyCFunction)Database_search_networks,
	488	METH_VARARGS\|METH_KEYWORDS,
	489	NULL,
	490	},
b1720435 MT	491	{
	492	"verify",
	493	(PyCFunction)Database_verify,
726f9984	494	METH_VARARGS,
b1720435 MT	495	NULL,
b1720435 MT	496	},
9cdf6c53 MT	497	{ NULL },
	498	};
	499
d99b0256	500	static struct PyGetSetDef Database_getsetters[] = {
a68a46f5 MT	501	{
	502	"ases",
	503	(getter)Database_ases,
	504	NULL,
	505	NULL,
	506	NULL,
	507	},
fa9a3663 MT	508	{
	509	"countries",
	510	(getter)Database_countries,
	511	NULL,
	512	NULL,
	513	NULL,
	514	},
53524b2d MT	515	{
	516	"created_at",
	517	(getter)Database_get_created_at,
	518	NULL,
	519	NULL,
	520	NULL,
	521	},
d99b0256 MT	522	{
	523	"description",
	524	(getter)Database_get_description,
	525	NULL,
	526	NULL,
	527	NULL,
	528	},
4bf49d00 MT	529	{
	530	"license",
	531	(getter)Database_get_license,
	532	NULL,
	533	NULL,
	534	NULL,
	535	},
a68a46f5 MT	536	{
	537	"networks",
	538	(getter)Database_networks,
	539	NULL,
	540	NULL,
	541	NULL,
	542	},
681ff05c MT	543	{
	544	"networks_flattened",
	545	(getter)Database_networks_flattened,
	546	NULL,
	547	NULL,
	548	NULL,
	549	},
d99b0256 MT	550	{
	551	"vendor",
	552	(getter)Database_get_vendor,
	553	NULL,
	554	NULL,
53524b2d	555	NULL,
d99b0256 MT	556	},
	557	{ NULL },
	558	};
	559
9cdf6c53 MT	560	PyTypeObject DatabaseType = {
9cdf6c53 MT	561	PyVarObject_HEAD_INIT(NULL, 0)
d42e1dcd MT	562	.tp_name = "location.Database",
	563	.tp_basicsize = sizeof(DatabaseObject),
	564	.tp_flags = Py_TPFLAGS_DEFAULT\|Py_TPFLAGS_BASETYPE,
	565	.tp_new = Database_new,
	566	.tp_dealloc = (destructor)Database_dealloc,
	567	.tp_init = (initproc)Database_init,
	568	.tp_doc = "Database object",
	569	.tp_methods = Database_methods,
	570	.tp_getset = Database_getsetters,
	571	.tp_repr = (reprfunc)Database_repr,
9cdf6c53	572	};
afb426df MT	573
	574	static PyObject* DatabaseEnumerator_new(PyTypeObject* type, PyObject* args, PyObject* kwds) {
	575	DatabaseEnumeratorObject* self = (DatabaseEnumeratorObject*)type->tp_alloc(type, 0);
	576
	577	return (PyObject*)self;
	578	}
	579
	580	static void DatabaseEnumerator_dealloc(DatabaseEnumeratorObject* self) {
	581	loc_database_enumerator_unref(self->enumerator);
	582
	583	Py_TYPE(self)->tp_free((PyObject* )self);
	584	}
	585
	586	static PyObject* DatabaseEnumerator_next(DatabaseEnumeratorObject* self) {
15f79e2d MT	587	struct loc_network* network = NULL;
15f79e2d MT	588
ccc7ab4e	589	// Enumerate all networks
15f79e2d MT	590	int r = loc_database_enumerator_next_network(self->enumerator, &network);
15f79e2d MT	591	if (r) {
20bb6d0c	592	PyErr_SetFromErrno(PyExc_ValueError);
15f79e2d MT	593	return NULL;
	594	}
	595
	596	// A network was found
ccc7ab4e MT	597	if (network) {
	598	PyObject* obj = new_network(&NetworkType, network);
	599	loc_network_unref(network);
	600
	601	return obj;
	602	}
	603
	604	// Enumerate all ASes
15f79e2d MT	605	struct loc_as* as = NULL;
	606
	607	r = loc_database_enumerator_next_as(self->enumerator, &as);
	608	if (r) {
20bb6d0c	609	PyErr_SetFromErrno(PyExc_ValueError);
15f79e2d MT	610	return NULL;
	611	}
	612
afb426df MT	613	if (as) {
	614	PyObject* obj = new_as(&ASType, as);
	615	loc_as_unref(as);
	616
	617	return obj;
	618	}
	619
fa9a3663 MT	620	// Enumerate all countries
	621	struct loc_country* country = NULL;
	622
	623	r = loc_database_enumerator_next_country(self->enumerator, &country);
	624	if (r) {
	625	PyErr_SetFromErrno(PyExc_ValueError);
	626	return NULL;
	627	}
	628
	629	if (country) {
	630	PyObject* obj = new_country(&CountryType, country);
	631	loc_country_unref(country);
	632
	633	return obj;
	634	}
	635
afb426df MT	636	// Nothing found, that means the end
	637	PyErr_SetNone(PyExc_StopIteration);
	638	return NULL;
	639	}
	640
	641	PyTypeObject DatabaseEnumeratorType = {
	642	PyVarObject_HEAD_INIT(NULL, 0)
d42e1dcd MT	643	.tp_name = "location.DatabaseEnumerator",
	644	.tp_basicsize = sizeof(DatabaseEnumeratorObject),
	645	.tp_flags = Py_TPFLAGS_DEFAULT,
	646	.tp_alloc = PyType_GenericAlloc,
	647	.tp_new = DatabaseEnumerator_new,
	648	.tp_dealloc = (destructor)DatabaseEnumerator_dealloc,
	649	.tp_iter = PyObject_SelfIter,
	650	.tp_iternext = (iternextfunc)DatabaseEnumerator_next,
afb426df	651	};