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

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

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

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

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

#include <Python.h>

#include <libloc/libloc.h>
#include <libloc/as.h>
#include <libloc/as-list.h>
#include <libloc/database.h>

#include "locationmodule.h"
#include "as.h"
#include "country.h"
#include "database.h"
#include "network.h"

static PyObject* Database_new(PyTypeObject* type, PyObject* args, PyObject* kwds) {
	DatabaseObject* self = (DatabaseObject*)type->tp_alloc(type, 0);

	return (PyObject*)self;
}

static void Database_dealloc(DatabaseObject* self) {
	if (self->db)
		loc_database_unref(self->db);

	if (self->path)
		free(self->path);

	Py_TYPE(self)->tp_free((PyObject* )self);
}

static int Database_init(DatabaseObject* self, PyObject* args, PyObject* kwargs) {
	const char* path = NULL;
	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) {
	struct loc_country* country = NULL;
	const char* country_code = NULL;

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

	// Fetch the country
	int r = loc_database_get_country(self->db, &country, country_code);
	if (r) {
		switch (errno) {
			case EINVAL:
				PyErr_SetString(PyExc_ValueError, "Invalid country code");
				break;

			default:
				PyErr_SetFromErrno(PyExc_OSError);
				break;
		}

		return NULL;
	}

	// No result
	if (!country)
		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;
	int r;

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

	// Try to retrieve a matching network
	r = loc_database_lookup_from_string(self->db, address, &network);
	if (r) {
		// Handle any errors
		switch (errno) {
			case EINVAL:
				PyErr_Format(PyExc_ValueError, "Invalid IP address: %s", address);
				break;

			default:
				PyErr_SetFromErrno(PyExc_OSError);
				break;
		}

		return NULL;
	}

	// Nothing found
	if (!network)
		Py_RETURN_NONE;

	// We got a network
	PyObject* obj = new_network(&NetworkType, network);
	loc_network_unref(network);

	return obj;
}

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) {
	const 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", (char**)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) {
	const char* kwlist[] = { "family", NULL };
	int family = AF_UNSPEC;

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