[location/libloc.git] / src / python / database.c

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

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

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

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

#include <Python.h>

#include <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);
		}

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

			loc_country_list_unref(countries);
			return NULL;
		}

		loc_country_list_unref(countries);
	}

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

		for (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
c98ebf8a MT	328	r = loc_database_enumerator_set_countries(enumerator, countries);
	329	if (r) {
	330	PyErr_SetFromErrno(PyExc_SystemError);
	331
a1a00053	332	loc_country_list_unref(countries);
c98ebf8a MT	333	return NULL;
c98ebf8a MT	334	}
e646a8f3	335
e646a8f3	336	loc_country_list_unref(countries);
ccc7ab4e MT	337	}
	338
	339	// Set the ASN we are searching for
84a2f0c2 MT	340	if (asn_list) {
	341	struct loc_as_list* asns;
	342	r = loc_as_list_new(loc_ctx, &asns);
	343	if (r) {
	344	PyErr_SetString(PyExc_SystemError, "Could not create AS list");
	345	return NULL;
	346	}
	347
	348	for (unsigned int i = 0; i < PyList_Size(asn_list); i++) {
	349	PyObject* item = PyList_GetItem(asn_list, i);
	350
	351	if (!PyLong_Check(item)) {
	352	PyErr_SetString(PyExc_TypeError, "ASNs must be numbers");
	353
	354	loc_as_list_unref(asns);
	355	return NULL;
	356	}
	357
	358	unsigned long number = PyLong_AsLong(item);
ccc7ab4e	359
84a2f0c2 MT	360	struct loc_as* as;
	361	r = loc_as_new(loc_ctx, &as, number);
	362	if (r) {
	363	PyErr_SetString(PyExc_SystemError, "Could not create AS");
	364
	365	loc_as_list_unref(asns);
	366	loc_as_unref(as);
	367	return NULL;
	368	}
	369
	370	r = loc_as_list_append(asns, as);
	371	if (r) {
	372	PyErr_SetString(PyExc_SystemError, "Could not append AS to the list");
	373
	374	loc_as_list_unref(asns);
	375	loc_as_unref(as);
	376	return NULL;
	377	}
	378
	379	loc_as_unref(as);
	380	}
	381
	382	r = loc_database_enumerator_set_asns(enumerator, asns);
ccc7ab4e MT	383	if (r) {
ccc7ab4e MT	384	PyErr_SetFromErrno(PyExc_SystemError);
84a2f0c2 MT	385
84a2f0c2 MT	386	loc_as_list_unref(asns);
ccc7ab4e MT	387	return NULL;
ccc7ab4e MT	388	}
84a2f0c2 MT	389
84a2f0c2 MT	390	loc_as_list_unref(asns);
ccc7ab4e MT	391	}
ccc7ab4e MT	392
bbdb2e0a MT	393	// Set the flags we are searching for
	394	if (flags) {
	395	r = loc_database_enumerator_set_flag(enumerator, flags);
	396
	397	if (r) {
	398	PyErr_SetFromErrno(PyExc_SystemError);
	399	return NULL;
	400	}
	401	}
	402
44e5ef71 MT	403	// Set the family we are searching for
	404	if (family) {
	405	r = loc_database_enumerator_set_family(enumerator, family);
	406
	407	if (r) {
	408	PyErr_SetFromErrno(PyExc_SystemError);
	409	return NULL;
	410	}
	411	}
	412
ccc7ab4e MT	413	PyObject* obj = new_database_enumerator(&DatabaseEnumeratorType, enumerator);
	414	loc_database_enumerator_unref(enumerator);
	415
	416	return obj;
	417	}
	418
fa9a3663	419	static PyObject* Database_countries(DatabaseObject* self) {
681ff05c	420	return Database_iterate_all(self, LOC_DB_ENUMERATE_COUNTRIES, 0);
fa9a3663 MT	421	}
fa9a3663 MT	422
9cdf6c53	423	static struct PyMethodDef Database_methods[] = {
86ca7ef7 MT	424	{
	425	"get_as",
	426	(PyCFunction)Database_get_as,
	427	METH_VARARGS,
	428	NULL,
	429	},
7c922e9c MT	430	{
	431	"get_country",
	432	(PyCFunction)Database_get_country,
	433	METH_VARARGS,
	434	NULL,
	435	},
31edab76 MT	436	{
	437	"lookup",
	438	(PyCFunction)Database_lookup,
	439	METH_VARARGS,
	440	NULL,
	441	},
afb426df MT	442	{
	443	"search_as",
	444	(PyCFunction)Database_search_as,
	445	METH_VARARGS,
	446	NULL,
	447	},
ccc7ab4e MT	448	{
	449	"search_networks",
	450	(PyCFunction)Database_search_networks,
	451	METH_VARARGS\|METH_KEYWORDS,
	452	NULL,
	453	},
b1720435 MT	454	{
	455	"verify",
	456	(PyCFunction)Database_verify,
726f9984	457	METH_VARARGS,
b1720435 MT	458	NULL,
b1720435 MT	459	},
9cdf6c53 MT	460	{ NULL },
	461	};
	462
d99b0256	463	static struct PyGetSetDef Database_getsetters[] = {
a68a46f5 MT	464	{
	465	"ases",
	466	(getter)Database_ases,
	467	NULL,
	468	NULL,
	469	NULL,
	470	},
fa9a3663 MT	471	{
	472	"countries",
	473	(getter)Database_countries,
	474	NULL,
	475	NULL,
	476	NULL,
	477	},
53524b2d MT	478	{
	479	"created_at",
	480	(getter)Database_get_created_at,
	481	NULL,
	482	NULL,
	483	NULL,
	484	},
d99b0256 MT	485	{
	486	"description",
	487	(getter)Database_get_description,
	488	NULL,
	489	NULL,
	490	NULL,
	491	},
4bf49d00 MT	492	{
	493	"license",
	494	(getter)Database_get_license,
	495	NULL,
	496	NULL,
	497	NULL,
	498	},
a68a46f5 MT	499	{
	500	"networks",
	501	(getter)Database_networks,
	502	NULL,
	503	NULL,
	504	NULL,
	505	},
681ff05c MT	506	{
	507	"networks_flattened",
	508	(getter)Database_networks_flattened,
	509	NULL,
	510	NULL,
	511	NULL,
	512	},
d99b0256 MT	513	{
	514	"vendor",
	515	(getter)Database_get_vendor,
	516	NULL,
	517	NULL,
53524b2d	518	NULL,
d99b0256 MT	519	},
	520	{ NULL },
	521	};
	522
9cdf6c53 MT	523	PyTypeObject DatabaseType = {
9cdf6c53 MT	524	PyVarObject_HEAD_INIT(NULL, 0)
d42e1dcd MT	525	.tp_name = "location.Database",
	526	.tp_basicsize = sizeof(DatabaseObject),
	527	.tp_flags = Py_TPFLAGS_DEFAULT\|Py_TPFLAGS_BASETYPE,
	528	.tp_new = Database_new,
	529	.tp_dealloc = (destructor)Database_dealloc,
	530	.tp_init = (initproc)Database_init,
	531	.tp_doc = "Database object",
	532	.tp_methods = Database_methods,
	533	.tp_getset = Database_getsetters,
	534	.tp_repr = (reprfunc)Database_repr,
9cdf6c53	535	};
afb426df MT	536
	537	static PyObject* DatabaseEnumerator_new(PyTypeObject* type, PyObject* args, PyObject* kwds) {
	538	DatabaseEnumeratorObject* self = (DatabaseEnumeratorObject*)type->tp_alloc(type, 0);
	539
	540	return (PyObject*)self;
	541	}
	542
	543	static void DatabaseEnumerator_dealloc(DatabaseEnumeratorObject* self) {
	544	loc_database_enumerator_unref(self->enumerator);
	545
	546	Py_TYPE(self)->tp_free((PyObject* )self);
	547	}
	548
	549	static PyObject* DatabaseEnumerator_next(DatabaseEnumeratorObject* self) {
15f79e2d MT	550	struct loc_network* network = NULL;
15f79e2d MT	551
ccc7ab4e	552	// Enumerate all networks
15f79e2d MT	553	int r = loc_database_enumerator_next_network(self->enumerator, &network);
15f79e2d MT	554	if (r) {
20bb6d0c	555	PyErr_SetFromErrno(PyExc_ValueError);
15f79e2d MT	556	return NULL;
	557	}
	558
	559	// A network was found
ccc7ab4e MT	560	if (network) {
	561	PyObject* obj = new_network(&NetworkType, network);
	562	loc_network_unref(network);
	563
	564	return obj;
	565	}
	566
	567	// Enumerate all ASes
15f79e2d MT	568	struct loc_as* as = NULL;
	569
	570	r = loc_database_enumerator_next_as(self->enumerator, &as);
	571	if (r) {
20bb6d0c	572	PyErr_SetFromErrno(PyExc_ValueError);
15f79e2d MT	573	return NULL;
	574	}
	575
afb426df MT	576	if (as) {
	577	PyObject* obj = new_as(&ASType, as);
	578	loc_as_unref(as);
	579
	580	return obj;
	581	}
	582
fa9a3663 MT	583	// Enumerate all countries
	584	struct loc_country* country = NULL;
	585
	586	r = loc_database_enumerator_next_country(self->enumerator, &country);
	587	if (r) {
	588	PyErr_SetFromErrno(PyExc_ValueError);
	589	return NULL;
	590	}
	591
	592	if (country) {
	593	PyObject* obj = new_country(&CountryType, country);
	594	loc_country_unref(country);
	595
	596	return obj;
	597	}
	598
afb426df MT	599	// Nothing found, that means the end
	600	PyErr_SetNone(PyExc_StopIteration);
	601	return NULL;
	602	}
	603
	604	PyTypeObject DatabaseEnumeratorType = {
	605	PyVarObject_HEAD_INIT(NULL, 0)
d42e1dcd MT	606	.tp_name = "location.DatabaseEnumerator",
	607	.tp_basicsize = sizeof(DatabaseEnumeratorObject),
	608	.tp_flags = Py_TPFLAGS_DEFAULT,
	609	.tp_alloc = PyType_GenericAlloc,
	610	.tp_new = DatabaseEnumerator_new,
	611	.tp_dealloc = (destructor)DatabaseEnumerator_dealloc,
	612	.tp_iter = PyObject_SelfIter,
	613	.tp_iternext = (iternextfunc)DatabaseEnumerator_next,
afb426df	614	};