src/python/database.c

   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
  19 #include <loc/libloc.h>
  20 #include <loc/database.h>
  21
  22 #include "locationmodule.h"
  23 #include "as.h"
  24 #include "country.h"
  25 #include "database.h"
  26 #include "network.h"
  27
  28 static PyObject* Database_new(PyTypeObject* type, PyObject* args, PyObject* kwds) {
  29         DatabaseObject* self = (DatabaseObject*)type->tp_alloc(type, 0);
  30
  31         return (PyObject*)self;
  32 }
  33
  34 static void Database_dealloc(DatabaseObject* self) {
  35         if (self->db)
  36                 loc_database_unref(self->db);
  37
  38         if (self->path)
  39                 free(self->path);
  40
  41         Py_TYPE(self)->tp_free((PyObject* )self);
  42 }
  43
  44 static int Database_init(DatabaseObject* self, PyObject* args, PyObject* kwargs) {
  45         const char* path = NULL;
  46
  47         if (!PyArg_ParseTuple(args, "s", &path))
  48                 return -1;
  49
  50         self->path = strdup(path);
  51
  52         // Open the file for reading
  53         FILE* f = fopen(self->path, "r");
  54         if (!f) {
  55                 PyErr_SetFromErrno(PyExc_IOError);
  56                 return -1;
  57         }
  58
  59         // Load the database
  60         int r = loc_database_new(loc_ctx, &self->db, f);
  61         fclose(f);
  62
  63         // Return on any errors
  64         if (r)
  65                 return -1;
  66
  67         return 0;
  68 }
  69
  70 static PyObject* Database_repr(DatabaseObject* self) {
  71         return PyUnicode_FromFormat("<Database %s>", self->path);
  72 }
  73
  74 static PyObject* Database_verify(DatabaseObject* self) {
  75         int r = loc_database_verify(self->db);
  76
  77         if (r == 0)
  78                 Py_RETURN_TRUE;
  79
  80         Py_RETURN_FALSE;
  81 }
  82
  83 static PyObject* Database_get_description(DatabaseObject* self) {
  84         const char* description = loc_database_get_description(self->db);
  85
  86         return PyUnicode_FromString(description);
  87 }
  88
  89 static PyObject* Database_get_vendor(DatabaseObject* self) {
  90         const char* vendor = loc_database_get_vendor(self->db);
  91
  92         return PyUnicode_FromString(vendor);
  93 }
  94
  95 static PyObject* Database_get_license(DatabaseObject* self) {
  96         const char* license = loc_database_get_license(self->db);
  97
  98         return PyUnicode_FromString(license);
  99 }
 100
 101 static PyObject* Database_get_created_at(DatabaseObject* self) {
 102         time_t created_at = loc_database_created_at(self->db);
 103
 104         return PyLong_FromLong(created_at);
 105 }
 106
 107 static PyObject* Database_get_as(DatabaseObject* self, PyObject* args) {
 108         struct loc_as* as = NULL;
 109         uint32_t number = 0;
 110
 111         if (!PyArg_ParseTuple(args, "i", &number))
 112                 return NULL;
 113
 114         // Try to retrieve the AS
 115         int r = loc_database_get_as(self->db, &as, number);
 116
 117         // We got an AS
 118         if (r == 0) {
 119                 PyObject* obj = new_as(&ASType, as);
 120                 loc_as_unref(as);
 121
 122                 return obj;
 123
 124         // Nothing found
 125         } else if (r == 1) {
 126                 Py_RETURN_NONE;
 127         }
 128
 129         // Unexpected error
 130         return NULL;
 131 }
 132
 133 static PyObject* Database_get_country(DatabaseObject* self, PyObject* args) {
 134         const char* country_code = NULL;
 135
 136         if (!PyArg_ParseTuple(args, "s", &country_code))
 137                 return NULL;
 138
 139         struct loc_country* country;
 140         int r = loc_database_get_country(self->db, &country, country_code);
 141         if (r) {
 142                 Py_RETURN_NONE;
 143         }
 144
 145         PyObject* obj = new_country(&CountryType, country);
 146         loc_country_unref(country);
 147
 148         return obj;
 149 }
 150
 151 static PyObject* Database_lookup(DatabaseObject* self, PyObject* args) {
 152         struct loc_network* network = NULL;
 153         const char* address = NULL;
 154
 155         if (!PyArg_ParseTuple(args, "s", &address))
 156                 return NULL;
 157
 158         // Try to retrieve a matching network
 159         int r = loc_database_lookup_from_string(self->db, address, &network);
 160
 161         // We got a network
 162         if (r == 0) {
 163                 PyObject* obj = new_network(&NetworkType, network);
 164                 loc_network_unref(network);
 165
 166                 return obj;
 167
 168         // Nothing found
 169         } else if (r == 1) {
 170                 Py_RETURN_NONE;
 171
 172         // Invalid input
 173         } else if (r == -EINVAL) {
 174                 PyErr_Format(PyExc_ValueError, "Invalid IP address: %s", address);
 175                 return NULL;
 176         }
 177
 178         // Unexpected error
 179         return NULL;
 180 }
 181
 182 static PyObject* new_database_enumerator(PyTypeObject* type, struct loc_database_enumerator* enumerator) {
 183         DatabaseEnumeratorObject* self = (DatabaseEnumeratorObject*)type->tp_alloc(type, 0);
 184         if (self) {
 185                 self->enumerator = loc_database_enumerator_ref(enumerator);
 186         }
 187
 188         return (PyObject*)self;
 189 }
 190
 191 static PyObject* Database_search_as(DatabaseObject* self, PyObject* args) {
 192         const char* string = NULL;
 193
 194         if (!PyArg_ParseTuple(args, "s", &string))
 195                 return NULL;
 196
 197         struct loc_database_enumerator* enumerator;
 198
 199         int r = loc_database_enumerator_new(&enumerator, self->db, LOC_DB_ENUMERATE_ASES);
 200         if (r) {
 201                 PyErr_SetFromErrno(PyExc_SystemError);
 202                 return NULL;
 203         }
 204
 205         // Search string we are searching for
 206         loc_database_enumerator_set_string(enumerator, string);
 207
 208         PyObject* obj = new_database_enumerator(&DatabaseEnumeratorType, enumerator);
 209         loc_database_enumerator_unref(enumerator);
 210
 211         return obj;
 212 }
 213
 214 static PyObject* Database_search_networks(DatabaseObject* self, PyObject* args, PyObject* kwargs) {
 215         char* kwlist[] = { "country_code", "asn", "flags", NULL };
 216         const char* country_code = NULL;
 217         unsigned int asn = 0;
 218         int flags = 0;
 219
 220     if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|sii", kwlist, &country_code, &asn, &flags))
 221                 return NULL;
 222
 223         struct loc_database_enumerator* enumerator;
 224         int r = loc_database_enumerator_new(&enumerator, self->db, LOC_DB_ENUMERATE_NETWORKS);
 225         if (r) {
 226                 PyErr_SetFromErrno(PyExc_SystemError);
 227                 return NULL;
 228         }
 229
 230         // Set country code we are searching for
 231         if (country_code) {
 232                 r = loc_database_enumerator_set_country_code(enumerator, country_code);
 233
 234                 if (r) {
 235                         PyErr_SetFromErrno(PyExc_SystemError);
 236                         return NULL;
 237                 }
 238         }
 239
 240         // Set the ASN we are searching for
 241         if (asn) {
 242                 r = loc_database_enumerator_set_asn(enumerator, asn);
 243
 244                 if (r) {
 245                         PyErr_SetFromErrno(PyExc_SystemError);
 246                         return NULL;
 247                 }
 248         }
 249
 250         // Set the flags we are searching for
 251         if (flags) {
 252                 r = loc_database_enumerator_set_flag(enumerator, flags);
 253
 254                 if (r) {
 255                         PyErr_SetFromErrno(PyExc_SystemError);
 256                         return NULL;
 257                 }
 258         }
 259
 260         PyObject* obj = new_database_enumerator(&DatabaseEnumeratorType, enumerator);
 261         loc_database_enumerator_unref(enumerator);
 262
 263         return obj;
 264 }
 265
 266 static struct PyMethodDef Database_methods[] = {
 267         {
 268                 "get_as",
 269                 (PyCFunction)Database_get_as,
 270                 METH_VARARGS,
 271                 NULL,
 272         },
 273         {
 274                 "get_country",
 275                 (PyCFunction)Database_get_country,
 276                 METH_VARARGS,
 277                 NULL,
 278         },
 279         {
 280                 "lookup",
 281                 (PyCFunction)Database_lookup,
 282                 METH_VARARGS,
 283                 NULL,
 284         },
 285         {
 286                 "search_as",
 287                 (PyCFunction)Database_search_as,
 288                 METH_VARARGS,
 289                 NULL,
 290         },
 291         {
 292                 "search_networks",
 293                 (PyCFunction)Database_search_networks,
 294                 METH_VARARGS|METH_KEYWORDS,
 295                 NULL,
 296         },
 297         {
 298                 "verify",
 299                 (PyCFunction)Database_verify,
 300                 METH_NOARGS,
 301                 NULL,
 302         },
 303         { NULL },
 304 };
 305
 306 static struct PyGetSetDef Database_getsetters[] = {
 307         {
 308                 "created_at",
 309                 (getter)Database_get_created_at,
 310                 NULL,
 311                 NULL,
 312                 NULL,
 313         },
 314         {
 315                 "description",
 316                 (getter)Database_get_description,
 317                 NULL,
 318                 NULL,
 319                 NULL,
 320         },
 321         {
 322                 "license",
 323                 (getter)Database_get_license,
 324                 NULL,
 325                 NULL,
 326                 NULL,
 327         },
 328         {
 329                 "vendor",
 330                 (getter)Database_get_vendor,
 331                 NULL,
 332                 NULL,
 333                 NULL,
 334         },
 335         { NULL },
 336 };
 337
 338 PyTypeObject DatabaseType = {
 339         PyVarObject_HEAD_INIT(NULL, 0)
 340         .tp_name =               "location.Database",
 341         .tp_basicsize =          sizeof(DatabaseObject),
 342         .tp_flags =              Py_TPFLAGS_DEFAULT|Py_TPFLAGS_BASETYPE,
 343         .tp_new =                Database_new,
 344         .tp_dealloc =            (destructor)Database_dealloc,
 345         .tp_init =               (initproc)Database_init,
 346         .tp_doc =                "Database object",
 347         .tp_methods =            Database_methods,
 348         .tp_getset =             Database_getsetters,
 349         .tp_repr =               (reprfunc)Database_repr,
 350 };
 351
 352 static PyObject* DatabaseEnumerator_new(PyTypeObject* type, PyObject* args, PyObject* kwds) {
 353         DatabaseEnumeratorObject* self = (DatabaseEnumeratorObject*)type->tp_alloc(type, 0);
 354
 355         return (PyObject*)self;
 356 }
 357
 358 static void DatabaseEnumerator_dealloc(DatabaseEnumeratorObject* self) {
 359         loc_database_enumerator_unref(self->enumerator);
 360
 361         Py_TYPE(self)->tp_free((PyObject* )self);
 362 }
 363
 364 static PyObject* DatabaseEnumerator_next(DatabaseEnumeratorObject* self) {
 365         struct loc_network* network = NULL;
 366
 367         // Enumerate all networks
 368         int r = loc_database_enumerator_next_network(self->enumerator, &network);
 369         if (r) {
 370                 PyErr_SetFromErrno(PyExc_ValueError);
 371                 return NULL;
 372         }
 373
 374         // A network was found
 375         if (network) {
 376                 PyObject* obj = new_network(&NetworkType, network);
 377                 loc_network_unref(network);
 378
 379                 return obj;
 380         }
 381
 382         // Enumerate all ASes
 383         struct loc_as* as = NULL;
 384
 385         r = loc_database_enumerator_next_as(self->enumerator, &as);
 386         if (r) {
 387                 PyErr_SetFromErrno(PyExc_ValueError);
 388                 return NULL;
 389         }
 390
 391         if (as) {
 392                 PyObject* obj = new_as(&ASType, as);
 393                 loc_as_unref(as);
 394
 395                 return obj;
 396         }
 397
 398         // Nothing found, that means the end
 399         PyErr_SetNone(PyExc_StopIteration);
 400         return NULL;
 401 }
 402
 403 PyTypeObject DatabaseEnumeratorType = {
 404         PyVarObject_HEAD_INIT(NULL, 0)
 405         .tp_name =               "location.DatabaseEnumerator",
 406         .tp_basicsize =          sizeof(DatabaseEnumeratorObject),
 407         .tp_flags =              Py_TPFLAGS_DEFAULT,
 408         .tp_alloc =              PyType_GenericAlloc,
 409         .tp_new =                DatabaseEnumerator_new,
 410         .tp_dealloc =            (destructor)DatabaseEnumerator_dealloc,
 411         .tp_iter =               PyObject_SelfIter,
 412         .tp_iternext =           (iternextfunc)DatabaseEnumerator_next,
 413 };