[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;

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

	// Try to retrieve a matching network
	int r = loc_database_lookup_from_string(self->db, address, &network);

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

		return obj;
	}

	// Nothing found
	if (!errno)
		Py_RETURN_NONE;

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

		default:
			PyErr_SetFromErrno(PyExc_OSError);
	}

	return NULL;
}

static PyObject* new_database_enumerator(PyTypeObject* type, struct loc_database_enumerator* enumerator) {
	DatabaseEnumeratorObject* self = (DatabaseEnumeratorObject*)type->tp_alloc(type, 0);
	if (self) {
		self->enumerator = loc_database_enumerator_ref(enumerator);
	}

	return (PyObject*)self;
}

static PyObject* Database_iterate_all(DatabaseObject* self,
		enum loc_database_enumerator_mode what, int family, int flags) {
	struct loc_database_enumerator* enumerator;

	int r = loc_database_enumerator_new(&enumerator, self->db, what, flags);
	if (r) {
		PyErr_SetFromErrno(PyExc_SystemError);
		return NULL;
	}

	// Set family
	if (family)
		loc_database_enumerator_set_family(enumerator, family);

	PyObject* obj = new_database_enumerator(&DatabaseEnumeratorType, enumerator);
	loc_database_enumerator_unref(enumerator);

	return obj;
}

static PyObject* Database_ases(DatabaseObject* self) {
	return Database_iterate_all(self, LOC_DB_ENUMERATE_ASES, AF_UNSPEC, 0);
}

static PyObject* Database_search_as(DatabaseObject* self, PyObject* args) {
	const char* string = NULL;

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

	struct loc_database_enumerator* enumerator;

	int r = loc_database_enumerator_new(&enumerator, self->db, LOC_DB_ENUMERATE_ASES, 0);
	if (r) {
		PyErr_SetFromErrno(PyExc_SystemError);
		return NULL;
	}

	// Search string we are searching for
	loc_database_enumerator_set_string(enumerator, string);

	PyObject* obj = new_database_enumerator(&DatabaseEnumeratorType, enumerator);
	loc_database_enumerator_unref(enumerator);

	return obj;
}

static PyObject* Database_networks(DatabaseObject* self) {
	return Database_iterate_all(self, LOC_DB_ENUMERATE_NETWORKS, AF_UNSPEC, 0);
}

static PyObject* Database_networks_flattened(DatabaseObject *self) {
	return Database_iterate_all(self, LOC_DB_ENUMERATE_NETWORKS, AF_UNSPEC,
		LOC_DB_ENUMERATOR_FLAGS_FLATTEN);
}

static PyObject* Database_search_networks(DatabaseObject* self, PyObject* args, PyObject* kwargs) {
	char* kwlist[] = { "country_codes", "asns", "flags", "family", "flatten", NULL };
	PyObject* country_codes = NULL;
	PyObject* asn_list = NULL;
	int flags = 0;
	int family = 0;
	int flatten = 0;

	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|O!O!iip", kwlist,
			&PyList_Type, &country_codes, &PyList_Type, &asn_list, &flags, &family, &flatten))
		return NULL;

	struct loc_database_enumerator* enumerator;
	int r = loc_database_enumerator_new(&enumerator, self->db, LOC_DB_ENUMERATE_NETWORKS,
		(flatten) ? LOC_DB_ENUMERATOR_FLAGS_FLATTEN : 0);
	if (r) {
		PyErr_SetFromErrno(PyExc_SystemError);
		return NULL;
	}

	// Set country code we are searching for
	if (country_codes) {
		struct loc_country_list* countries;
		r = loc_country_list_new(loc_ctx, &countries);
		if (r) {
			PyErr_SetString(PyExc_SystemError, "Could not create country list");
			return NULL;
		}

		for (int i = 0; i < PyList_Size(country_codes); i++) {
			PyObject* item = PyList_GetItem(country_codes, i);

			if (!PyUnicode_Check(item)) {
				PyErr_SetString(PyExc_TypeError, "Country codes must be strings");
				loc_country_list_unref(countries);
				return NULL;
			}

			const char* country_code = PyUnicode_AsUTF8(item);

			struct loc_country* country;
			r = loc_country_new(loc_ctx, &country, country_code);
			if (r) {
				if (r == -EINVAL) {
					PyErr_Format(PyExc_ValueError, "Invalid country code: %s", country_code);
				} else {
					PyErr_SetString(PyExc_SystemError, "Could not create country");
				}

				loc_country_list_unref(countries);
				return NULL;
			}

			// Append it to the list
			r = loc_country_list_append(countries, country);
			if (r) {
				PyErr_SetString(PyExc_SystemError, "Could not append country to the list");

				loc_country_list_unref(countries);
				loc_country_unref(country);
				return NULL;
			}

			loc_country_unref(country);
		}

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

			loc_country_list_unref(countries);
			return NULL;
		}

		loc_country_list_unref(countries);
	}

	// Set the ASN we are searching for
	if (asn_list) {
		struct loc_as_list* asns;
		r = loc_as_list_new(loc_ctx, &asns);
		if (r) {
			PyErr_SetFromErrno(PyExc_OSError);
			return NULL;
		}

		for (int i = 0; i < PyList_Size(asn_list); i++) {
			PyObject* item = PyList_GetItem(asn_list, i);

			if (!PyLong_Check(item)) {
				PyErr_SetString(PyExc_TypeError, "ASNs must be numbers");

				loc_as_list_unref(asns);
				return NULL;
			}

			unsigned long number = PyLong_AsLong(item);

			struct loc_as* as;
			r = loc_as_new(loc_ctx, &as, number);
			if (r) {
				PyErr_SetFromErrno(PyExc_OSError);

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

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

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

			loc_as_unref(as);
		}

		r = loc_database_enumerator_set_asns(enumerator, asns);
		if (r) {
			PyErr_SetFromErrno(PyExc_OSError);

			loc_as_list_unref(asns);
			return NULL;
		}

		loc_as_list_unref(asns);
	}

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

		if (r) {
			PyErr_SetFromErrno(PyExc_OSError);
			return NULL;
		}
	}

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

		if (r) {
			PyErr_SetFromErrno(PyExc_OSError);
			return NULL;
		}
	}

	PyObject* obj = new_database_enumerator(&DatabaseEnumeratorType, enumerator);
	loc_database_enumerator_unref(enumerator);

	return obj;
}

static PyObject* Database_countries(DatabaseObject* self) {
	return Database_iterate_all(self, LOC_DB_ENUMERATE_COUNTRIES, AF_UNSPEC, 0);
}

static PyObject* Database_list_bogons(DatabaseObject* self, PyObject* args, PyObject* kwargs) {
	char* kwlist[] = { "family", NULL };
	int family = AF_UNSPEC;

	// Parse arguments
	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|i", kwlist, &family))
		return NULL;

	return Database_iterate_all(self, LOC_DB_ENUMERATE_BOGONS, family, 0);
}

static struct PyMethodDef Database_methods[] = {
	{
		"get_as",
		(PyCFunction)Database_get_as,
		METH_VARARGS,
		NULL,
	},
	{
		"get_country",
		(PyCFunction)Database_get_country,
		METH_VARARGS,
		NULL,
	},
	{
		"list_bogons",
		(PyCFunction)Database_list_bogons,
		METH_VARARGS|METH_KEYWORDS,
		NULL,
	},
	{
		"lookup",
		(PyCFunction)Database_lookup,
		METH_VARARGS,
		NULL,
	},
	{
		"search_as",
		(PyCFunction)Database_search_as,
		METH_VARARGS,
		NULL,
	},
	{
		"search_networks",
		(PyCFunction)Database_search_networks,
		METH_VARARGS|METH_KEYWORDS,
		NULL,
	},
	{
		"verify",
		(PyCFunction)Database_verify,
		METH_VARARGS,
		NULL,
	},
	{ NULL },
};

static struct PyGetSetDef Database_getsetters[] = {
	{
		"ases",
		(getter)Database_ases,
		NULL,
		NULL,
		NULL,
	},
	{
		"countries",
		(getter)Database_countries,
		NULL,
		NULL,
		NULL,
	},
	{
		"created_at",
		(getter)Database_get_created_at,
		NULL,
		NULL,
		NULL,
	},
	{
		"description",
		(getter)Database_get_description,
		NULL,
		NULL,
		NULL,
	},
	{
		"license",
		(getter)Database_get_license,
		NULL,
		NULL,
		NULL,
	},
	{
		"networks",
		(getter)Database_networks,
		NULL,
		NULL,
		NULL,
	},
	{
		"networks_flattened",
		(getter)Database_networks_flattened,
		NULL,
		NULL,
		NULL,
	},
	{
		"vendor",
		(getter)Database_get_vendor,
		NULL,
		NULL,
		NULL,
	},
	{ NULL },
};

PyTypeObject DatabaseType = {
	PyVarObject_HEAD_INIT(NULL, 0)
	.tp_name =               "location.Database",
	.tp_basicsize =          sizeof(DatabaseObject),
	.tp_flags =              Py_TPFLAGS_DEFAULT|Py_TPFLAGS_BASETYPE,
	.tp_new =                Database_new,
	.tp_dealloc =            (destructor)Database_dealloc,
	.tp_init =               (initproc)Database_init,
	.tp_doc =                "Database object",
	.tp_methods =            Database_methods,
	.tp_getset =             Database_getsetters,
	.tp_repr =               (reprfunc)Database_repr,
};

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

	return (PyObject*)self;
}

static void DatabaseEnumerator_dealloc(DatabaseEnumeratorObject* self) {
	loc_database_enumerator_unref(self->enumerator);

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

static PyObject* DatabaseEnumerator_next(DatabaseEnumeratorObject* self) {
	struct loc_network* network = NULL;

	// Enumerate all networks
	int r = loc_database_enumerator_next_network(self->enumerator, &network);
	if (r) {
		PyErr_SetFromErrno(PyExc_ValueError);
		return NULL;
	}

	// A network was found
	if (network) {
		PyObject* obj = new_network(&NetworkType, network);
		loc_network_unref(network);

		return obj;
	}

	// Enumerate all ASes
	struct loc_as* as = NULL;

	r = loc_database_enumerator_next_as(self->enumerator, &as);
	if (r) {
		PyErr_SetFromErrno(PyExc_ValueError);
		return NULL;
	}

	if (as) {
		PyObject* obj = new_as(&ASType, as);
		loc_as_unref(as);

		return obj;
	}

	// Enumerate all countries
	struct loc_country* country = NULL;

	r = loc_database_enumerator_next_country(self->enumerator, &country);
	if (r) {
		PyErr_SetFromErrno(PyExc_ValueError);
		return NULL;
	}

	if (country) {
		PyObject* obj = new_country(&CountryType, country);
		loc_country_unref(country);

		return obj;
	}

	// Nothing found, that means the end
	PyErr_SetNone(PyExc_StopIteration);
	return NULL;
}

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