src/_collectymodule.c

   1 /*
   2  * collecty
   3  * Copyright (C) 2015 IPFire Team (www.ipfire.org)
   4  *
   5  * This program is free software; you can redistribute it and/or modify
   6  * it under the terms of the GNU General Public License as published by
   7  * the Free Software Foundation; either version 3 of the License, or
   8  * (at your option) any later version.
   9  *
  10  * This program is distributed in the hope that it will be useful,
  11  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  12  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  13  * GNU General Public License for more details.
  14  *
  15  * You should have received a copy of the GNU General Public License
  16  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  17  */
  18
  19 #include <Python.h>
  20
  21 #include <atasmart.h>
  22 #include <errno.h>
  23 #include <fcntl.h>
  24 #include <linux/hdreg.h>
  25 #include <stdbool.h>
  26 #include <string.h>
  27 #include <sys/ioctl.h>
  28
  29 #define MODEL_SIZE  40
  30 #define SERIAL_SIZE 20
  31
  32 typedef struct {
  33         PyObject_HEAD
  34         char* path;
  35         struct hd_driveid identity;
  36         SkDisk* disk;
  37 } BlockDevice;
  38
  39 static void BlockDevice_dealloc(BlockDevice* self) {
  40         if (self->disk)
  41                 sk_disk_free(self->disk);
  42
  43         if (self->path)
  44                 free(self->path);
  45
  46         self->ob_type->tp_free((PyObject*)self);
  47 }
  48
  49 static int BlockDevice_get_identity(BlockDevice* device) {
  50         int fd;
  51
  52         if ((fd = open(device->path, O_RDONLY | O_NONBLOCK)) < 0) {
  53                 return 1;
  54         }
  55
  56         int r = ioctl(fd, HDIO_GET_IDENTITY, &device->identity);
  57         close(fd);
  58
  59         if (r)
  60                 return 1;
  61
  62         return 0;
  63 }
  64
  65 static int BlockDevice_smart_is_available(BlockDevice* device) {
  66         SkBool available = FALSE;
  67
  68         int r = sk_disk_smart_is_available(device->disk, &available);
  69         if (r)
  70                 return -1;
  71
  72         if (available)
  73                 return 0;
  74
  75         return 1;
  76 }
  77
  78 static int BlockDevice_check_sleep_mode(BlockDevice* device) {
  79         SkBool awake = FALSE;
  80
  81         int r = sk_disk_check_sleep_mode(device->disk, &awake);
  82         if (r)
  83                 return -1;
  84
  85         if (awake)
  86                 return 0;
  87
  88         return 1;
  89 }
  90
  91 static PyObject * BlockDevice_new(PyTypeObject* type, PyObject* args, PyObject* kwds) {
  92         BlockDevice* self = (BlockDevice*)type->tp_alloc(type, 0);
  93
  94         if (self) {
  95                 self->path = NULL;
  96
  97                 // libatasmart
  98                 self->disk = NULL;
  99         }
 100
 101         return (PyObject *)self;
 102 }
 103
 104 static int BlockDevice_init(BlockDevice* self, PyObject* args, PyObject* kwds) {
 105         const char* path = NULL;
 106
 107         if (!PyArg_ParseTuple(args, "s", &path))
 108                 return -1;
 109
 110         self->path = strdup(path);
 111
 112         int r = BlockDevice_get_identity(self);
 113         if (r) {
 114                 PyErr_Format(PyExc_OSError, "Could not open block device: %s", path);
 115                 return -1;
 116         }
 117
 118         r = sk_disk_open(path, &self->disk);
 119         if (r == 0) {
 120                 if (BlockDevice_smart_is_available(self) == 0) {
 121                         if (BlockDevice_check_sleep_mode(self) == 0) {
 122                                 r = sk_disk_smart_read_data(self->disk);
 123                                 if (r) {
 124                                         PyErr_Format(PyExc_OSError, "Could not open block device %s: %s", path,
 125                                                 strerror(errno));
 126                                         return -1;
 127                                 }
 128                         }
 129                 }
 130         } else {
 131                 PyErr_Format(PyExc_OSError, "Could not open block device %s: %s", path,
 132                         strerror(errno));
 133                 return -1;
 134         }
 135
 136         //sk_disk_identify_is_available
 137
 138         return 0;
 139 }
 140
 141 static PyObject* BlockDevice_get_path(PyObject* self) {
 142         BlockDevice* device = (BlockDevice*)self;
 143
 144         return PyString_FromString(device->path);
 145 }
 146
 147 static void clean_string(char *s) {
 148         for (char* e = s; *e; e++) {
 149                 if (*e < ' ' || *e >= 127)
 150                         *e = ' ';
 151         }
 152 }
 153
 154 static void drop_spaces(char *s) {
 155         char *d = s;
 156         bool prev_space = false;
 157
 158         s += strspn(s, " ");
 159
 160         for (; *s; s++) {
 161                 if (prev_space) {
 162                         if (*s != ' ') {
 163                                 prev_space = false;
 164                                 *(d++) = ' ';
 165                                 *(d++) = *s;
 166                         }
 167                 } else {
 168                         if (*s == ' ')
 169                                 prev_space = true;
 170                         else
 171                                 *(d++) = *s;
 172                 }
 173         }
 174
 175         *d = 0;
 176 }
 177
 178 static void copy_string(char* d, const char* s, size_t n) {
 179         // Copy the source buffer to the destination buffer up to n
 180         memcpy(d, s, n);
 181
 182         // Terminate the destination buffer with NULL
 183         d[n] = '\0';
 184
 185         // Clean up the string from non-printable characters
 186         clean_string(d);
 187         drop_spaces(d);
 188 }
 189
 190 static PyObject* BlockDevice_get_model(PyObject* self) {
 191         BlockDevice* device = (BlockDevice*)self;
 192
 193         char model[MODEL_SIZE + 1];
 194         copy_string(model, device->identity.model, sizeof(model));
 195
 196         return PyString_FromString(model);
 197 }
 198
 199 static PyObject* BlockDevice_get_serial(PyObject* self) {
 200         BlockDevice* device = (BlockDevice*)self;
 201
 202         char serial[SERIAL_SIZE + 1];
 203         copy_string(serial, device->identity.serial_no, sizeof(serial));
 204
 205         return PyString_FromString(serial);
 206 }
 207
 208 static PyObject* BlockDevice_is_smart_supported(PyObject* self) {
 209         BlockDevice* device = (BlockDevice*)self;
 210
 211         if (BlockDevice_smart_is_available(device) == 0)
 212                 Py_RETURN_TRUE;
 213
 214         Py_RETURN_FALSE;
 215 }
 216
 217 static PyObject* BlockDevice_is_awake(PyObject* self) {
 218         BlockDevice* device = (BlockDevice*)self;
 219
 220         if (BlockDevice_check_sleep_mode(device) == 0)
 221                 Py_RETURN_TRUE;
 222
 223         Py_RETURN_FALSE;
 224 }
 225
 226 static PyObject* BlockDevice_get_bad_sectors(PyObject* self) {
 227         BlockDevice* device = (BlockDevice*)self;
 228
 229         if (BlockDevice_smart_is_available(device)) {
 230                 PyErr_Format(PyExc_OSError, "Device does not support SMART");
 231                 return NULL;
 232         }
 233
 234         uint64_t bad_sectors;
 235         int r = sk_disk_smart_get_bad(device->disk, &bad_sectors);
 236         if (r)
 237                 return NULL;
 238
 239         return PyLong_FromUnsignedLongLong((unsigned long long)bad_sectors);
 240 }
 241
 242 static PyObject* BlockDevice_get_temperature(PyObject* self) {
 243         BlockDevice* device = (BlockDevice*)self;
 244
 245         if (BlockDevice_smart_is_available(device)) {
 246                 PyErr_Format(PyExc_OSError, "Device does not support SMART");
 247                 return NULL;
 248         }
 249
 250         uint64_t mkelvin;
 251         int r = sk_disk_smart_get_temperature(device->disk, &mkelvin);
 252         if (r)
 253                 return NULL;
 254
 255         return PyLong_FromUnsignedLongLong((unsigned long long)mkelvin);
 256 }
 257
 258 static PyGetSetDef BlockDevice_getsetters[] = {
 259         {"path", (getter)BlockDevice_get_path, NULL, NULL, NULL},
 260         {"model", (getter)BlockDevice_get_model, NULL, NULL, NULL},
 261         {"serial", (getter)BlockDevice_get_serial, NULL, NULL, NULL},
 262 };
 263
 264 static PyMethodDef BlockDevice_methods[] = {
 265         {"get_bad_sectors", (PyCFunction)BlockDevice_get_bad_sectors, METH_NOARGS, NULL},
 266         {"get_temperature", (PyCFunction)BlockDevice_get_temperature, METH_NOARGS, NULL},
 267         {"is_smart_supported", (PyCFunction)BlockDevice_is_smart_supported, METH_NOARGS, NULL},
 268         {"is_awake", (PyCFunction)BlockDevice_is_awake, METH_NOARGS, NULL},
 269         {NULL}
 270 };
 271
 272 static PyTypeObject BlockDeviceType = {
 273         PyObject_HEAD_INIT(NULL)
 274         0,                                  /*ob_size*/
 275         "_collecty.BlockDevice",            /*tp_name*/
 276         sizeof(BlockDevice),                /*tp_basicsize*/
 277         0,                                  /*tp_itemsize*/
 278         (destructor)BlockDevice_dealloc,    /*tp_dealloc*/
 279         0,                                  /*tp_print*/
 280         0,                                  /*tp_getattr*/
 281         0,                                  /*tp_setattr*/
 282         0,                                  /*tp_compare*/
 283         0,                                  /*tp_repr*/
 284         0,                                  /*tp_as_number*/
 285         0,                                  /*tp_as_sequence*/
 286         0,                                  /*tp_as_mapping*/
 287         0,                                  /*tp_hash */
 288         0,                                  /*tp_call*/
 289         0,                                  /*tp_str*/
 290         0,                                  /*tp_getattro*/
 291         0,                                  /*tp_setattro*/
 292         0,                                  /*tp_as_buffer*/
 293         Py_TPFLAGS_DEFAULT|Py_TPFLAGS_BASETYPE, /*tp_flags*/
 294         "BlockDevice objects",              /* tp_doc */
 295         0,                                          /* tp_traverse */
 296         0,                                          /* tp_clear */
 297         0,                                          /* tp_richcompare */
 298         0,                                          /* tp_weaklistoffset */
 299         0,                                          /* tp_iter */
 300         0,                                          /* tp_iternext */
 301         BlockDevice_methods,                /* tp_methods */
 302         0,                                  /* tp_members */
 303         BlockDevice_getsetters,             /* tp_getset */
 304         0,                                  /* tp_base */
 305         0,                                  /* tp_dict */
 306         0,                                  /* tp_descr_get */
 307         0,                                  /* tp_descr_set */
 308         0,                                  /* tp_dictoffset */
 309         (initproc)BlockDevice_init,         /* tp_init */
 310         0,                                  /* tp_alloc */
 311         BlockDevice_new,                    /* tp_new */
 312 };
 313
 314 static PyMethodDef collecty_module_methods[] = {
 315         {NULL},
 316 };
 317
 318 void init_collecty(void) {
 319         if (PyType_Ready(&BlockDeviceType) < 0)
 320                 return;
 321
 322         PyObject* m = Py_InitModule("_collecty", collecty_module_methods);
 323
 324         PyModule_AddObject(m, "BlockDevice", (PyObject*)&BlockDeviceType);
 325 }