Beispiel #1
0
    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))
Beispiel #2
0
    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))
Beispiel #3
0
 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)))
Beispiel #4
0
 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)))
Beispiel #5
0
    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))
Beispiel #6
0
    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))
Beispiel #7
0
    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])
Beispiel #8
0
    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])
Beispiel #9
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())
Beispiel #10
0
    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),))
Beispiel #11
0
    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()
Beispiel #12
0
    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),))
Beispiel #13
0
    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()
Beispiel #14
0
 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])
Beispiel #15
0
 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])
Beispiel #16
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)
Beispiel #17
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)
Beispiel #18
0
    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)
Beispiel #19
0
    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)
Beispiel #20
0
    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)
Beispiel #21
0
 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)