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", NULL };
 235         const char* country_code = NULL;
 236         unsigned int asn = 0;
 237         int flags = 0;
 238
 239     if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|sii", kwlist, &country_code, &asn, &flags))
 240                 return NULL;
 241
 242         struct loc_database_enumerator* enumerator;
 243         int r = loc_database_enumerator_new(&enumerator, self->db, LOC_DB_ENUMERATE_NETWORKS);
 244         if (r) {
 245                 PyErr_SetFromErrno(PyExc_SystemError);
 246                 return NULL;
 247         }
 248
 249         // Set country code we are searching for
 250         if (country_code) {
 251                 r = loc_database_enumerator_set_country_code(enumerator, country_code);
 252
 253                 if (r) {
 254                         PyErr_SetFromErrno(PyExc_SystemError);
 255                         return NULL;
 256                 }
 257         }
 258
 259         // Set the ASN we are searching for
 260         if (asn) {
 261                 r = loc_database_enumerator_set_asn(enumerator, asn);
 262
 263                 if (r) {
 264                         PyErr_SetFromErrno(PyExc_SystemError);
 265                         return NULL;
 266                 }
 267         }
 268
 269         // Set the flags we are searching for
 270         if (flags) {
 271                 r = loc_database_enumerator_set_flag(enumerator, flags);
 272
 273                 if (r) {
 274                         PyErr_SetFromErrno(PyExc_SystemError);
 275                         return NULL;
 276                 }
 277         }
 278
 279         PyObject* obj = new_database_enumerator(&DatabaseEnumeratorType, enumerator);
 280         loc_database_enumerator_unref(enumerator);
 281
 282         return obj;
 283 }
 284
 285 static struct PyMethodDef Database_methods[] = {
 286         {
 287                 "get_as",
 288                 (PyCFunction)Database_get_as,
 289                 METH_VARARGS,
 290                 NULL,
 291         },
 292         {
 293                 "get_country",
 294                 (PyCFunction)Database_get_country,
 295                 METH_VARARGS,
 296                 NULL,
 297         },
 298         {
 299                 "lookup",
 300                 (PyCFunction)Database_lookup,
 301                 METH_VARARGS,
 302                 NULL,
 303         },
 304         {
 305                 "search_as",
 306                 (PyCFunction)Database_search_as,
 307                 METH_VARARGS,
 308                 NULL,
 309         },
 310         {
 311                 "search_networks",
 312                 (PyCFunction)Database_search_networks,
 313                 METH_VARARGS|METH_KEYWORDS,
 314                 NULL,
 315         },
 316         {
 317                 "verify",
 318                 (PyCFunction)Database_verify,
 319                 METH_VARARGS,
 320                 NULL,
 321         },
 322         { NULL },
 323 };
 324
 325 static struct PyGetSetDef Database_getsetters[] = {
 326         {
 327                 "created_at",
 328                 (getter)Database_get_created_at,
 329                 NULL,
 330                 NULL,
 331                 NULL,
 332         },
 333         {
 334                 "description",
 335                 (getter)Database_get_description,
 336                 NULL,
 337                 NULL,
 338                 NULL,
 339         },
 340         {
 341                 "license",
 342                 (getter)Database_get_license,
 343                 NULL,
 344                 NULL,
 345                 NULL,
 346         },
 347         {
 348                 "vendor",
 349                 (getter)Database_get_vendor,
 350                 NULL,
 351                 NULL,
 352                 NULL,
 353         },
 354         { NULL },
 355 };
 356
 357 PyTypeObject DatabaseType = {
 358         PyVarObject_HEAD_INIT(NULL, 0)
 359         .tp_name =               "location.Database",
 360         .tp_basicsize =          sizeof(DatabaseObject),
 361         .tp_flags =              Py_TPFLAGS_DEFAULT|Py_TPFLAGS_BASETYPE,
 362         .tp_new =                Database_new,
 363         .tp_dealloc =            (destructor)Database_dealloc,
 364         .tp_init =               (initproc)Database_init,
 365         .tp_doc =                "Database object",
 366         .tp_methods =            Database_methods,
 367         .tp_getset =             Database_getsetters,
 368         .tp_repr =               (reprfunc)Database_repr,
 369 };
 370
 371 static PyObject* DatabaseEnumerator_new(PyTypeObject* type, PyObject* args, PyObject* kwds) {
 372         DatabaseEnumeratorObject* self = (DatabaseEnumeratorObject*)type->tp_alloc(type, 0);
 373
 374         return (PyObject*)self;
 375 }
 376
 377 static void DatabaseEnumerator_dealloc(DatabaseEnumeratorObject* self) {
 378         loc_database_enumerator_unref(self->enumerator);
 379
 380         Py_TYPE(self)->tp_free((PyObject* )self);
 381 }
 382
 383 static PyObject* DatabaseEnumerator_next(DatabaseEnumeratorObject* self) {
 384         struct loc_network* network = NULL;
 385
 386         // Enumerate all networks
 387         int r = loc_database_enumerator_next_network(self->enumerator, &network);
 388         if (r) {
 389                 PyErr_SetFromErrno(PyExc_ValueError);
 390                 return NULL;
 391         }
 392
 393         // A network was found
 394         if (network) {
 395                 PyObject* obj = new_network(&NetworkType, network);
 396                 loc_network_unref(network);
 397
 398                 return obj;
 399         }
 400
 401         // Enumerate all ASes
 402         struct loc_as* as = NULL;
 403
 404         r = loc_database_enumerator_next_as(self->enumerator, &as);
 405         if (r) {
 406                 PyErr_SetFromErrno(PyExc_ValueError);
 407                 return NULL;
 408         }
 409
 410         if (as) {
 411                 PyObject* obj = new_as(&ASType, as);
 412                 loc_as_unref(as);
 413
 414                 return obj;
 415         }
 416
 417         // Nothing found, that means the end
 418         PyErr_SetNone(PyExc_StopIteration);
 419         return NULL;
 420 }
 421
 422 PyTypeObject DatabaseEnumeratorType = {
 423         PyVarObject_HEAD_INIT(NULL, 0)
 424         .tp_name =               "location.DatabaseEnumerator",
 425         .tp_basicsize =          sizeof(DatabaseEnumeratorObject),
 426         .tp_flags =              Py_TPFLAGS_DEFAULT,
 427         .tp_alloc =              PyType_GenericAlloc,
 428         .tp_new =                DatabaseEnumerator_new,
 429         .tp_dealloc =            (destructor)DatabaseEnumerator_dealloc,
 430         .tp_iter =               PyObject_SelfIter,
 431         .tp_iternext =           (iternextfunc)DatabaseEnumerator_next,
 432 };