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

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

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

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

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

#include <Python.h>

#include <loc/libloc.h>
#include <loc/as.h>
#include <loc/as-list.h>
#include <loc/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);

	return PyUnicode_FromString(description);
}

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

	return PyUnicode_FromString(vendor);
}

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

	return PyUnicode_FromString(license);
}

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

	return PyLong_FromLong(created_at);
}

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

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

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

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

		return obj;

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

	// Unexpected error
	return NULL;
}

static PyObject* Database_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 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;
	}

	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, 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, 0);
}

static PyObject* Database_networks_flattened(DatabaseObject *self) {
	return Database_iterate_all(self, LOC_DB_ENUMERATE_NETWORKS, 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 (unsigned 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);
		}

		loc_database_enumerator_set_countries(enumerator, countries);

		Py_DECREF(country_codes);
		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 (unsigned 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, 0);
}

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