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