tag.py (4332B)
1 # standard imports 2 import hashlib 3 import logging 4 5 logg = logging.getLogger().getChild(__name__) 6 7 8 9 class TagPool: 10 11 def __init__(self): 12 self.pool = {} 13 14 15 def register(self, account): 16 z = account.sum() 17 if self.pool.get(z) == None: 18 self.pool[z] = [] 19 20 self.pool[z].append(account) 21 22 logg.debug('pool register account {} to tag sum {}'.format(account, z.hex())) 23 24 return True 25 26 27 class Tag: 28 """Represents a collection of tags for a cached object. 29 30 When a new tag is added, the tag collection is deterministically ordered and summed. 31 """ 32 def __init__(self): 33 self.tags = [] 34 self.tag_values = {} 35 self.sum = b'\x00' * 32 36 self.dirty = False 37 38 39 def get(self): 40 """The current deterministic sum of the tags. 41 42 :rtype: bytes 43 :return: Tag digest sum 44 """ 45 if self.dirty: 46 self.tags.sort() 47 h = hashlib.new('sha256') 48 for tag in self.tags: 49 h.update(tag) 50 self.sum = h.digest() 51 return self.sum 52 53 54 def add(self, tag, value=None): 55 """Add a tag to the collection. 56 57 Client code should call Tag.create() instead. 58 59 :param tag: Tag value digest 60 :type tag: bytes 61 :param value: Tag value 62 :type value: bytes 63 :rtype: boolean 64 :returns: False if tag digest already exists in object 65 """ 66 if tag in self.tags: 67 return False 68 self.tags.append(tag) 69 self.tag_values[tag] = value 70 self.dirty = True 71 return True 72 73 74 def create(self, value): 75 """Create a new tag record to add to the collection. 76 77 :param value: Tag value 78 :type value: bytes 79 :rtype: bytes 80 :return: Digest of newly added tag 81 """ 82 h = hashlib.new('sha256') 83 h.update(value) 84 tag = h.digest() 85 self.add(tag, value) 86 return tag 87 88 89 def merge(self, tags): 90 """Merge contents of two tag objects. After this operation the sum of each of the tag objects will be identical. 91 92 :param tags: Tag collection to merge with 93 :type tags: taint.tag.Tag 94 :raises TypeError: If argument is not a taint.tag.Tag instance 95 """ 96 if not isinstance(tags, Tag): 97 raise TypeError('tags must be type taint.tag.Tag') 98 for tag in tags.tags: 99 self.add(tag) 100 self.tag_values[tag] = tags.tag_values[tag] 101 102 for tag in self.tags: 103 tags.add(tag) 104 tags.tag_values[tag] = self.tag_values[tag] 105 106 107 def serialize(self): 108 """Serialize tags for storage. 109 110 Serialized tags are deterministically ordered. 111 112 :rtype: bytes 113 :returns: Serialized tags 114 """ 115 b = self.get() 116 for tag in self.tags: 117 b += tag 118 return b 119 120 121 def deserialize(self, b, skip_check=False): 122 """Deserialize tags into currently instantiated object. 123 124 Deserialization will ADD tags to the current object. If different tags already exist in the object, the resulting collection will not be identical to the serialized data. 125 126 :param b: Serialized tag data 127 :type b: bytes 128 :raises ValueError: If skip_check is not set, and serialized data does not match tag object sum 129 """ 130 if len(b) % 32 > 0: 131 raise ValueError('invalid data length; remainder {} from 32'.format(len(b) % 32)) 132 cursor = 32 133 z = b[:cursor] 134 135 for i in range(cursor, len(b), 32): 136 tag = b[i:i+32] 137 logg.debug('deserialize add {}'.format(tag)) 138 self.add(tag) 139 140 if not skip_check: 141 zz = self.get() 142 if z != zz: 143 raise ValueError('data sum does not match content; expected {}, found {}'.format(zz.hex(), z.hex())) 144 145 146 def __str__(self): 147 tag_list = [] 148 for tag in self.tags: 149 v = self.tag_values[tag] 150 if v == None: 151 v = tag.hex() 152 else: 153 try: 154 v = v.decode('utf-8') 155 except UnicodeDecodeError: 156 v = v.hex() 157 tag_list.append(v) 158 tag_list.sort() 159 return ','.join(tag_list)