Statistics
| Branch: | Tag: | Revision:

root / snf-cyclades-app / synnefo / db / pools / __init__.py @ f8714db8

History | View | Annotate | Download (10.3 kB)

1
from bitarray import bitarray
2
from base64 import b64encode, b64decode
3
import ipaddr
4

    
5
AVAILABLE = True
6
UNAVAILABLE = False
7

    
8

    
9
class PoolManager(object):
10
    """PoolManager for DB PoolTable models.
11

12
    This class implements a persistent Pool mechanism based on rows of
13
    PoolTable objects. Values that are pooled by this class, are mapped to an
14
    index on a bitarray, which is the one that is stored on the DB.
15

16
    The object that will be used in order to initialize this pool, must have
17
    two string attributes (available_map and reserved_map) and the size of the
18
    pool.
19

20
    Subclasses of PoolManager must implement value_to_index and index_to_value
21
    method's in order to denote how the value will be mapped to the index in
22
    the bitarray.
23

24
    Important!!: Updates on a PoolManager object are not reflected to the DB,
25
    until save() method is called.
26

27
    """
28
    def __init__(self, pool_table):
29
        self.pool_table = pool_table
30
        self.pool_size = pool_table.size
31
        if pool_table.available_map:
32
            self.available = _bitarray_from_string(pool_table.available_map)
33
            self.reserved = _bitarray_from_string(pool_table.reserved_map)
34
        else:
35
            self.available = self._create_empty_pool(self.pool_size)
36
            self.reserved = self._create_empty_pool(self.pool_size)
37
            self.add_padding(self.pool_size)
38

    
39
    def _create_empty_pool(self, size):
40
        ba = bitarray(size)
41
        ba.setall(AVAILABLE)
42
        return ba
43

    
44
    def add_padding(self, pool_size):
45
        bits = find_padding(pool_size)
46
        self.available.extend([UNAVAILABLE] * bits)
47
        self.reserved.extend([UNAVAILABLE] * bits)
48

    
49
    def cut_padding(self, pool_size):
50
        bits = find_padding(pool_size)
51
        self.available = self.available[:-bits]
52
        self.reserved = self.reserved[:-bits]
53

    
54
    @property
55
    def pool(self):
56
        return (self.available & self.reserved)
57

    
58
    def get(self, value=None):
59
        """Get a value from the pool."""
60
        if value is None:
61
            if self.empty():
62
                raise EmptyPool
63
            # Get the first available index
64
            index = int(self.pool.index(AVAILABLE))
65
            assert(index < self.pool_size)
66
            self._reserve(index)
67
            return self.index_to_value(index)
68
        else:
69
            if not self.contains(value):
70
                raise InvalidValue("Value %s does not belong to pool." % value)
71
            if self.is_available(value):
72
                self.reserve(value)
73
                return value
74
            else:
75
                raise ValueNotAvailable("Value %s is not available" % value)
76

    
77
    def put(self, value, external=False):
78
        """Return a value to the pool."""
79
        if value is None:
80
            raise ValueError
81
        if not self.contains(value):
82
            raise InvalidValue("%s does not belong to pool." % value)
83
        index = self.value_to_index(value)
84
        self._release(index, external)
85

    
86
    def reserve(self, value, external=False):
87
        """Reserve a value."""
88
        if not self.contains(value):
89
            raise InvalidValue("%s does not belong to pool." % value)
90
        index = self.value_to_index(value)
91
        self._reserve(index, external)
92
        return True
93

    
94
    def save(self, db=True):
95
        """Save changes to the DB."""
96
        self.pool_table.available_map = _bitarray_to_string(self.available)
97
        self.pool_table.reserved_map = _bitarray_to_string(self.reserved)
98
        if db:
99
            self.pool_table.save()
100

    
101
    def empty(self):
102
        """Return True when pool is empty."""
103
        return not self.pool.any()
104

    
105
    def size(self):
106
        """Return the size of the bitarray(original size + padding)."""
107
        return self.pool.length()
108

    
109
    def _reserve(self, index, external=False):
110
        if external:
111
            self.reserved[index] = UNAVAILABLE
112
        else:
113
            self.available[index] = UNAVAILABLE
114

    
115
    def _release(self, index, external=False):
116
        if external:
117
            self.reserved[index] = AVAILABLE
118
        else:
119
            self.available[index] = AVAILABLE
120

    
121
    def contains(self, value, index=False):
122
        if index is False:
123
            index = self.value_to_index(value)
124
        return index >= 0 and index < self.pool_size
125

    
126
    def count_available(self):
127
        return self.pool.count(AVAILABLE)
128

    
129
    def count_unavailable(self):
130
        return self.pool_size - self.count_available()
131

    
132
    def count_reserved(self):
133
        return self.reserved[:self.pool_size].count(UNAVAILABLE)
134

    
135
    def count_unreserved(self):
136
        return self.pool_size - self.count_reserved()
137

    
138
    def is_available(self, value, index=False):
139
        if not self.contains(value, index=index):
140
            raise InvalidValue("%s does not belong to pool." % value)
141
        if not index:
142
            idx = self.value_to_index(value)
143
        else:
144
            idx = value
145
        return self.pool[idx] == AVAILABLE
146

    
147
    def is_reserved(self, value, index=False):
148
        if not self.contains(value, index=index):
149
            raise InvalidValue("%s does not belong to pool." % value)
150
        if not index:
151
            idx = self.value_to_index(value)
152
        else:
153
            idx = value
154
        return self.reserved[idx] == UNAVAILABLE
155

    
156
    def to_01(self):
157
        return self.pool[:self.pool_size].to01()
158

    
159
    def to_map(self):
160
        return self.to_01().replace("0", "X").replace("1", ".")
161

    
162
    def extend(self, bits_num):
163
        assert(bits_num >= 0)
164
        self.resize(bits_num)
165

    
166
    def shrink(self, bits_num):
167
        assert(bits_num >= 0)
168
        size = self.pool_size
169
        tmp = self.available[(size - bits_num): size]
170
        if tmp.count(UNAVAILABLE):
171
            raise Exception("Can not shrink. In use")
172
        self.resize(-bits_num)
173

    
174
    def resize(self, bits_num):
175
        if bits_num == 0:
176
            return
177
        # Cut old padding
178
        self.cut_padding(self.pool_size)
179
        # Do the resize
180
        if bits_num > 0:
181
            self.available.extend([AVAILABLE] * bits_num)
182
            self.reserved.extend([AVAILABLE] * bits_num)
183
        else:
184
            self.available = self.available[:bits_num]
185
            self.reserved = self.reserved[:bits_num]
186
        # Add new padding
187
        self.pool_size = self.pool_size + bits_num
188
        self.add_padding(self.pool_size)
189
        self.pool_table.size = self.pool_size
190

    
191
    def index_to_value(self, index):
192
        raise NotImplementedError
193

    
194
    def value_to_index(self, value):
195
        raise NotImplementedError
196

    
197
    def __repr__(self):
198
        return repr(self.pool_table)
199

    
200

    
201
class EmptyPool(Exception):
202
    pass
203

    
204

    
205
class ValueNotAvailable(Exception):
206
    pass
207

    
208

    
209
class InvalidValue(Exception):
210
    pass
211

    
212

    
213
def find_padding(size):
214
    extra = size % 8
215
    return extra and (8 - extra) or 0
216

    
217

    
218
def bitarray_to_01(bitarray_):
219
    return bitarray_.to01()
220

    
221

    
222
def bitarray_to_map(bitarray_):
223
    return bitarray_to_01(bitarray_).replace("0", "X").replace("1", ".")
224

    
225

    
226
def _bitarray_from_string(bitarray_):
227
    ba = bitarray()
228
    ba.frombytes(b64decode(bitarray_))
229
    return ba
230

    
231

    
232
def _bitarray_to_string(bitarray_):
233
    return b64encode(bitarray_.tobytes())
234

    
235
##
236
## Custom pools
237
##
238

    
239

    
240
class BridgePool(PoolManager):
241
    def index_to_value(self, index):
242
        # Bridge indexes should start from 1
243
        return self.pool_table.base + str(index + 1)
244

    
245
    def value_to_index(self, value):
246
        return int(value.replace(self.pool_table.base, "")) - 1
247

    
248

    
249
class MacPrefixPool(PoolManager):
250
    def __init__(self, pool_table):
251
        do_init = False if pool_table.available_map else True
252
        super(MacPrefixPool, self).__init__(pool_table)
253
        if do_init:
254
            for i in xrange(1, self.pool_size):
255
                if not self.validate_mac(self.index_to_value(i)):
256
                    self._reserve(i, external=True)
257
            # Reserve the first mac-prefix for public-networks
258
            self._reserve(0, external=True)
259

    
260
    def index_to_value(self, index):
261
        """Convert index to mac_prefix"""
262
        base = self.pool_table.base
263
        a = hex(int(base.replace(":", ""), 16) + index).replace("0x", '')
264
        mac_prefix = ":".join([a[x:x + 2] for x in xrange(0, len(a), 2)])
265
        return mac_prefix
266

    
267
    def value_to_index(self, value):
268
        base = self.pool_table.base
269
        return int(value.replace(":", ""), 16) - int(base.replace(":", ""), 16)
270

    
271
    @staticmethod
272
    def validate_mac(value):
273
        hex_ = value.replace(":", "")
274
        bin_ = bin(int(hex_, 16))[2:].zfill(8)
275
        return bin_[6] == '1' and bin_[7] == '0'
276

    
277

    
278
class IPPool(PoolManager):
279
    def __init__(self, pool_table):
280
        subnet = pool_table.subnet
281
        self.net = ipaddr.IPNetwork(subnet.cidr)
282
        self.offset = pool_table.offset
283
        self.base = pool_table.base
284
        if pool_table.available_map:
285
            initialized_pool = True
286
        else:
287
            initialized_pool = False
288
        super(IPPool, self).__init__(pool_table)
289
        if not initialized_pool:
290
            self.check_pool_integrity()
291

    
292
    def check_pool_integrity(self):
293
        """Check the integrity of the IP pool
294

295
        This check is required only for old IP pools(one IP pool per network
296
        that contained the whole subnet cidr) that had not been initialized
297
        before the migration. This checks that the network, gateway and
298
        broadcast IP addresses are externally reserved, and if not reserves
299
        them.
300

301
        """
302
        subnet = self.pool_table.subnet
303
        check_ips = [str(self.net.network), str(self.net.broadcast)]
304
        if subnet.gateway is not None:
305
            check_ips.append(subnet.gateway)
306
        for ip in check_ips:
307
            if self.contains(ip):
308
                self.reserve(ip, external=True)
309

    
310
    def value_to_index(self, value):
311
        addr = ipaddr.IPAddress(value)
312
        return int(addr) - int(self.net.network) - int(self.offset)
313

    
314
    def index_to_value(self, index):
315
        return str(self.net[index + int(self.offset)])
316

    
317
    def contains(self, address, index=False):
318
        if index is False:
319
            addr = ipaddr.IPAddress(address)
320
            if addr not in self.net:
321
                return False
322
        return super(IPPool, self).contains(address, index=False)
323

    
324
    def return_start(self):
325
        return str(ipaddr.IPAddress(ipaddr.IPNetwork(self.base).network) +
326
                   self.offset)
327

    
328
    def return_end(self):
329
        return str(ipaddr.IPAddress(ipaddr.IPNetwork(self.base).network) +
330
                   self.offset + self.pool_size - 1)