def __internal_get_eq_attrs_values(self, obj): """ :param obj: The object from which I need a unique string. :return: A string with all the values from the get_eq_attrs() method concatenated. This should represent the object in an unique way. """ if type(obj).__name__ in dir(__builtin__): return cpickle_dumps(obj) elif obj is None: return cpickle_dumps(obj) elif isinstance(obj, DiskItem): result = '' for attr in obj.get_eq_attrs(): value = getattr(obj, attr) result += self.__internal_get_eq_attrs_values(value) return result else: msg = ('Complex classes like %s need to inherit from DiskItem to' ' be stored.') raise Exception(msg % type(obj))
def __setitem__(self, key, value): # Test if it is already in the DB: if key in self: query = 'UPDATE %s SET value = ? WHERE key=?' % self.table_name self.db.execute(query, (cpickle_dumps(value), cpickle_dumps(key))) else: query = "INSERT INTO %s VALUES (NULL, ?, ?)" % self.table_name self.db.execute(query, (cpickle_dumps(key), cpickle_dumps(value)))
def update(self, old_info, update_info): """ :param old_info: The info/vuln instance to be updated in the kb. :param update_info: The info/vuln instance with new information :return: Nothing """ old_not_info = not isinstance(old_info, (Info, InfoSet, Shell)) update_not_info = not isinstance(update_info, (Info, InfoSet, Shell)) if old_not_info or update_not_info: msg = ('You MUST use raw_write/raw_read to store non-info objects' ' to the KnowledgeBase.') raise TypeError(msg) old_uniq_id = old_info.get_uniq_id() new_uniq_id = update_info.get_uniq_id() pickle = cpickle_dumps(update_info) # Update the pickle and unique_id after finding by original uniq_id query = "UPDATE %s SET pickle = ?, uniq_id = ? WHERE uniq_id = ?" params = (pickle, new_uniq_id, old_uniq_id) result = self.db.execute(query % self.table_name, params).result() if result.rowcount: self._notify_observers(self.UPDATE, old_info, update_info) else: ex = ('Failed to update() %s instance because' ' the original unique_id (%s) does not exist in the DB,' ' or the new unique_id (%s) is invalid.') raise DBException(ex % (old_info.__class__.__name__, old_uniq_id, new_uniq_id))
def update(self, old_info, update_info): """ :param old_info: The info/vuln instance to be updated in the kb. :param update_info: The info/vuln instance with new information :return: Nothing """ old_not_info = not isinstance(old_info, (Info, InfoSet, Shell)) update_not_info = not isinstance(update_info, (Info, InfoSet, Shell)) if old_not_info or update_not_info: msg = 'You MUST use raw_write/raw_read to store non-info objects'\ ' to the KnowledgeBase.' raise TypeError(msg) old_uniq_id = old_info.get_uniq_id() new_uniq_id = update_info.get_uniq_id() pickle = cpickle_dumps(update_info) # Update the pickle and unique_id after finding by original uniq_id query = "UPDATE %s SET pickle = ?, uniq_id = ? WHERE uniq_id = ?" params = (pickle, new_uniq_id, old_uniq_id) result = self.db.execute(query % self.table_name, params).result() if result.rowcount: self._notify_observers(self.UPDATE, old_info, update_info) else: ex = 'Failed to update() %s instance because' \ ' the original unique_id (%s) does not exist in the DB,' \ ' or the new unique_id (%s) is invalid.' raise DBException(ex % (old_info.__class__.__name__, old_uniq_id, new_uniq_id))
def __getitem__(self, key): query = 'SELECT value FROM %s WHERE key=? limit 1' % self.table_name r = self.db.select(query, (cpickle_dumps(key),)) if not r: raise KeyError('%s not in DiskDict.' % key) return cPickle.loads(r[0][0])
def test_pickle(self): i = MockInfo() iset1 = InfoSet([i]) pickled_iset1 = cpickle_dumps(iset1) iset1_clone = loads(pickled_iset1) self.assertEqual(iset1.get_uniq_id(), iset1_clone.get_uniq_id())
def __delitem__(self, key): """ Delete the key from the dict :param key: The key to delete :return: None """ query = 'DELETE FROM %s WHERE key = ?' % self.table_name self.db.execute(query, (cpickle_dumps(key),))
def _get_eq_attrs_values(self, obj): """ :param obj: The object from which I need a hash. :return: A hash representing the eq_attrs specified in the DiskItem. """ attr_values = self._get_attr_values_as_builtin(obj) concatenated_eq_attrs = cpickle_dumps(attr_values) return hashlib.md5(concatenated_eq_attrs).hexdigest()
def __contains__(self, key): """ :return: True if the value is in keys """ # Adding the "limit 1" to the query makes it faster, as it won't # have to scan through all the table/index, it just stops on the # first match. query = 'SELECT count(*) FROM %s WHERE key=? limit 1' % self.table_name r = self.db.select_one(query, (cpickle_dumps(key),)) return bool(r[0])
def _dump(self, obj): """ Serialize an object, using the dump function provided in __init__ or cPickle if None was specified. :return: A string containing the object """ if self.dump is not None: return self.dump(obj) return cpickle_dumps(obj)
def append(self, value): """ Append a value to the DiskList. :param value: The value to append. """ assert self._state == OPEN pickled_obj = cpickle_dumps(value) eq_attrs = self._get_eq_attrs_values(value) t = (eq_attrs, pickled_obj) query = "INSERT INTO %s VALUES (NULL, ?, ?)" % self.table_name self.db.execute(query, t)
def append(self, location_a, location_b, value, ignore_type=False): """ This method appends the location_b value to a dict. """ if not ignore_type and not isinstance(value, (Info, Shell)): msg = 'You MUST use raw_write/raw_read to store non-info objects'\ ' to the KnowledgeBase.' raise TypeError(msg) location_a = self._get_real_name(location_a) uniq_id = self._get_uniq_id(value) pickled_obj = cpickle_dumps(value) t = (location_a, location_b, uniq_id, pickled_obj) query = "INSERT INTO %s VALUES (?, ?, ?, ?)" % self.table_name self.db.execute(query, t) self._notify(location_a, location_b, value)