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 <libloc/libloc.h>
  20 #include <libloc/as.h>
  21 #include <libloc/as-list.h>
  22 #include <libloc/database.h>
  23
  24 #include "locationmodule.h"
  25 #include "as.h"
  26 #include "country.h"
  27 #include "database.h"
  28 #include "network.h"
  29
  30 static PyObject* Database_new(PyTypeObject* type, PyObject* args, PyObject* kwds) {
  31         DatabaseObject* self = (DatabaseObject*)type->tp_alloc(type, 0);
  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
  40         if (self->path)
  41                 free(self->path);
  42
  43         Py_TYPE(self)->tp_free((PyObject* )self);
  44 }
  45
  46 static int Database_init(DatabaseObject* self, PyObject* args, PyObject* kwargs) {
  47         const char* path = NULL;
  48         FILE* f = NULL;
  49
  50         // Parse arguments
  51         if (!PyArg_ParseTuple(args, "s", &path))
  52                 return -1;
  53
  54         // Copy path
  55         self->path = strdup(path);
  56         if (!self->path)
  57                 goto ERROR;
  58
  59         // Open the file for reading
  60         f = fopen(self->path, "r");
  61         if (!f)
  62                 goto ERROR;
  63
  64         // Load the database
  65         int r = loc_database_new(loc_ctx, &self->db, f);
  66         if (r)
  67                 goto ERROR;
  68
  69         fclose(f);
  70         return 0;
  71
  72 ERROR:
  73         if (f)
  74                 fclose(f);
  75
  76         PyErr_SetFromErrno(PyExc_OSError);
  77         return -1;
  78 }
  79
  80 static PyObject* Database_repr(DatabaseObject* self) {
  81         return PyUnicode_FromFormat("<Database %s>", self->path);
  82 }
  83
  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);
 105
 106         if (r == 0)
 107                 Py_RETURN_TRUE;
 108
 109         Py_RETURN_FALSE;
 110 }
 111
 112 static PyObject* Database_get_description(DatabaseObject* self) {
 113         const char* description = loc_database_get_description(self->db);
 114         if (!description)
 115                 Py_RETURN_NONE;
 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);
 122         if (!vendor)
 123                 Py_RETURN_NONE;
 124
 125         return PyUnicode_FromString(vendor);
 126 }
 127
 128 static PyObject* Database_get_license(DatabaseObject* self) {
 129         const char* license = loc_database_get_license(self->db);
 130         if (!license)
 131                 Py_RETURN_NONE;
 132
 133         return PyUnicode_FromString(license);
 134 }
 135
 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
 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);
 151
 152         // We got an AS
 153         if (r == 0) {
 154                 PyObject* obj = new_as(&ASType, as);
 155                 loc_as_unref(as);
 156
 157                 return obj;
 158
 159         // Nothing found
 160         } else if (r == 1) {
 161                 Py_RETURN_NONE;
 162         }
 163
 164         // Unexpected error
 165         return NULL;
 166 }
 167
 168 static PyObject* Database_get_country(DatabaseObject* self, PyObject* args) {
 169         struct loc_country* country = NULL;
 170         const char* country_code = NULL;
 171
 172         if (!PyArg_ParseTuple(args, "s", &country_code))
 173                 return NULL;
 174
 175         // Fetch the country
 176         int r = loc_database_get_country(self->db, &country, country_code);
 177         if (r) {
 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;
 189         }
 190
 191         // No result
 192         if (!country)
 193                 Py_RETURN_NONE;
 194
 195         PyObject* obj = new_country(&CountryType, country);
 196         loc_country_unref(country);
 197
 198         return obj;
 199 }
 200
 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;
 217         }
 218
 219         // Nothing found
 220         if (!errno)
 221                 Py_RETURN_NONE;
 222
 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);
 230         }
 231
 232         return NULL;
 233 }
 234
 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
 244 static PyObject* Database_iterate_all(DatabaseObject* self,
 245                 enum loc_database_enumerator_mode what, int family, int flags) {
 246         struct loc_database_enumerator* enumerator;
 247
 248         int r = loc_database_enumerator_new(&enumerator, self->db, what, flags);
 249         if (r) {
 250                 PyErr_SetFromErrno(PyExc_SystemError);
 251                 return NULL;
 252         }
 253
 254         // Set family
 255         if (family)
 256                 loc_database_enumerator_set_family(enumerator, family);
 257
 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) {
 265         return Database_iterate_all(self, LOC_DB_ENUMERATE_ASES, AF_UNSPEC, 0);
 266 }
 267
 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
 276         int r = loc_database_enumerator_new(&enumerator, self->db, LOC_DB_ENUMERATE_ASES, 0);
 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
 291 static PyObject* Database_networks(DatabaseObject* self) {
 292         return Database_iterate_all(self, LOC_DB_ENUMERATE_NETWORKS, AF_UNSPEC, 0);
 293 }
 294
 295 static PyObject* Database_networks_flattened(DatabaseObject *self) {
 296         return Database_iterate_all(self, LOC_DB_ENUMERATE_NETWORKS, AF_UNSPEC,
 297                 LOC_DB_ENUMERATOR_FLAGS_FLATTEN);
 298 }
 299
 300 static PyObject* Database_search_networks(DatabaseObject* self, PyObject* args, PyObject* kwargs) {
 301         char* kwlist[] = { "country_codes", "asns", "flags", "family", "flatten", NULL };
 302         PyObject* country_codes = NULL;
 303         PyObject* asn_list = NULL;
 304         int flags = 0;
 305         int family = 0;
 306         int flatten = 0;
 307
 308         if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|O!O!iip", kwlist,
 309                         &PyList_Type, &country_codes, &PyList_Type, &asn_list, &flags, &family, &flatten))
 310                 return NULL;
 311
 312         struct loc_database_enumerator* enumerator;
 313         int r = loc_database_enumerator_new(&enumerator, self->db, LOC_DB_ENUMERATE_NETWORKS,
 314                 (flatten) ? LOC_DB_ENUMERATOR_FLAGS_FLATTEN : 0);
 315         if (r) {
 316                 PyErr_SetFromErrno(PyExc_SystemError);
 317                 return NULL;
 318         }
 319
 320         // Set country code we are searching for
 321         if (country_codes) {
 322                 struct loc_country_list* countries;
 323                 r = loc_country_list_new(loc_ctx, &countries);
 324                 if (r) {
 325                         PyErr_SetString(PyExc_SystemError, "Could not create country list");
 326                         return NULL;
 327                 }
 328
 329                 for (int i = 0; i < PyList_Size(country_codes); i++) {
 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
 366                 r = loc_database_enumerator_set_countries(enumerator, countries);
 367                 if (r) {
 368                         PyErr_SetFromErrno(PyExc_SystemError);
 369
 370                         loc_country_list_unref(countries);
 371                         return NULL;
 372                 }
 373
 374                 loc_country_list_unref(countries);
 375         }
 376
 377         // Set the ASN we are searching for
 378         if (asn_list) {
 379                 struct loc_as_list* asns;
 380                 r = loc_as_list_new(loc_ctx, &asns);
 381                 if (r) {
 382                         PyErr_SetFromErrno(PyExc_OSError);
 383                         return NULL;
 384                 }
 385
 386                 for (int i = 0; i < PyList_Size(asn_list); i++) {
 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);
 397
 398                         struct loc_as* as;
 399                         r = loc_as_new(loc_ctx, &as, number);
 400                         if (r) {
 401                                 PyErr_SetFromErrno(PyExc_OSError);
 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) {
 410                                 PyErr_SetFromErrno(PyExc_OSError);
 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);
 421                 if (r) {
 422                         PyErr_SetFromErrno(PyExc_OSError);
 423
 424                         loc_as_list_unref(asns);
 425                         return NULL;
 426                 }
 427
 428                 loc_as_list_unref(asns);
 429         }
 430
 431         // Set the flags we are searching for
 432         if (flags) {
 433                 r = loc_database_enumerator_set_flag(enumerator, flags);
 434
 435                 if (r) {
 436                         PyErr_SetFromErrno(PyExc_OSError);
 437                         return NULL;
 438                 }
 439         }
 440
 441         // Set the family we are searching for
 442         if (family) {
 443                 r = loc_database_enumerator_set_family(enumerator, family);
 444
 445                 if (r) {
 446                         PyErr_SetFromErrno(PyExc_OSError);
 447                         return NULL;
 448                 }
 449         }
 450
 451         PyObject* obj = new_database_enumerator(&DatabaseEnumeratorType, enumerator);
 452         loc_database_enumerator_unref(enumerator);
 453
 454         return obj;
 455 }
 456
 457 static PyObject* Database_countries(DatabaseObject* self) {
 458         return Database_iterate_all(self, LOC_DB_ENUMERATE_COUNTRIES, AF_UNSPEC, 0);
 459 }
 460
 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
 472 static struct PyMethodDef Database_methods[] = {
 473         {
 474                 "get_as",
 475                 (PyCFunction)Database_get_as,
 476                 METH_VARARGS,
 477                 NULL,
 478         },
 479         {
 480                 "get_country",
 481                 (PyCFunction)Database_get_country,
 482                 METH_VARARGS,
 483                 NULL,
 484         },
 485         {
 486                 "list_bogons",
 487                 (PyCFunction)Database_list_bogons,
 488                 METH_VARARGS|METH_KEYWORDS,
 489                 NULL,
 490         },
 491         {
 492                 "lookup",
 493                 (PyCFunction)Database_lookup,
 494                 METH_VARARGS,
 495                 NULL,
 496         },
 497         {
 498                 "search_as",
 499                 (PyCFunction)Database_search_as,
 500                 METH_VARARGS,
 501                 NULL,
 502         },
 503         {
 504                 "search_networks",
 505                 (PyCFunction)Database_search_networks,
 506                 METH_VARARGS|METH_KEYWORDS,
 507                 NULL,
 508         },
 509         {
 510                 "verify",
 511                 (PyCFunction)Database_verify,
 512                 METH_VARARGS,
 513                 NULL,
 514         },
 515         { NULL },
 516 };
 517
 518 static struct PyGetSetDef Database_getsetters[] = {
 519         {
 520                 "ases",
 521                 (getter)Database_ases,
 522                 NULL,
 523                 NULL,
 524                 NULL,
 525         },
 526         {
 527                 "countries",
 528                 (getter)Database_countries,
 529                 NULL,
 530                 NULL,
 531                 NULL,
 532         },
 533         {
 534                 "created_at",
 535                 (getter)Database_get_created_at,
 536                 NULL,
 537                 NULL,
 538                 NULL,
 539         },
 540         {
 541                 "description",
 542                 (getter)Database_get_description,
 543                 NULL,
 544                 NULL,
 545                 NULL,
 546         },
 547         {
 548                 "license",
 549                 (getter)Database_get_license,
 550                 NULL,
 551                 NULL,
 552                 NULL,
 553         },
 554         {
 555                 "networks",
 556                 (getter)Database_networks,
 557                 NULL,
 558                 NULL,
 559                 NULL,
 560         },
 561         {
 562                 "networks_flattened",
 563                 (getter)Database_networks_flattened,
 564                 NULL,
 565                 NULL,
 566                 NULL,
 567         },
 568         {
 569                 "vendor",
 570                 (getter)Database_get_vendor,
 571                 NULL,
 572                 NULL,
 573                 NULL,
 574         },
 575         { NULL },
 576 };
 577
 578 PyTypeObject DatabaseType = {
 579         PyVarObject_HEAD_INIT(NULL, 0)
 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,
 590 };
 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) {
 605         struct loc_network* network = NULL;
 606
 607         // Enumerate all networks
 608         int r = loc_database_enumerator_next_network(self->enumerator, &network);
 609         if (r) {
 610                 PyErr_SetFromErrno(PyExc_ValueError);
 611                 return NULL;
 612         }
 613
 614         // A network was found
 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
 623         struct loc_as* as = NULL;
 624
 625         r = loc_database_enumerator_next_as(self->enumerator, &as);
 626         if (r) {
 627                 PyErr_SetFromErrno(PyExc_ValueError);
 628                 return NULL;
 629         }
 630
 631         if (as) {
 632                 PyObject* obj = new_as(&ASType, as);
 633                 loc_as_unref(as);
 634
 635                 return obj;
 636         }
 637
 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
 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)
 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,
 669 };