util.py 6.28 KB
Newer Older
Aurel's avatar
Aurel committed
1
#
Grégory Wisniewski's avatar
Grégory Wisniewski committed
2
# Copyright (C) 2006-2010  Nexedi SA
3
#
Aurel's avatar
Aurel committed
4 5 6 7
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
8
#
Aurel's avatar
Aurel committed
9 10 11 12 13 14 15
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
16
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
Aurel's avatar
Aurel committed
17 18


19
import re
20
import socket
21
from binascii import a2b_hex, b2a_hex
22
from hashlib import sha1
23
from Queue import deque
24 25
from struct import pack, unpack

26
SOCKET_CONNECTORS_DICT = {
Olivier Cros's avatar
Olivier Cros committed
27 28 29 30
    socket.AF_INET : 'SocketConnectorIPv4',
    socket.AF_INET6: 'SocketConnectorIPv6',
}

31 32 33 34 35
def u64(s):
    return unpack('!Q', s)[0]

def p64(n):
    return pack('!Q', n)
36

37 38 39 40
def add64(packed, offset):
    """Add a python number to a 64-bits packed value"""
    return p64(u64(packed) + offset)

Yoshinori Okuji's avatar
Yoshinori Okuji committed
41
def dump(s):
42
    """Dump a binary string in hex."""
43 44 45
    if s is not None:
        if isinstance(s, str):
            return b2a_hex(s)
Aurel's avatar
Aurel committed
46
        return repr(s)
Yoshinori Okuji's avatar
Yoshinori Okuji committed
47

Aurel's avatar
Aurel committed
48 49
def bin(s):
    """Inverse of dump method."""
50 51
    if s is not None:
        return a2b_hex(s)
52

Aurel's avatar
Aurel committed
53

54
def makeChecksum(s):
55 56
    """Return a 20-byte checksum against a string."""
    return sha1(s).digest()
57

58 59

def resolve(hostname):
Grégory Wisniewski's avatar
Grégory Wisniewski committed
60 61 62 63 64 65 66 67 68
    """
        Returns the first IP address that match with the given hostname
    """
    try:
        # an IP resolves to itself
        _, _, address_list = socket.gethostbyname_ex(hostname)
    except socket.gaierror:
        return None
    return address_list[0]
69

70
def getAddressType(address):
Olivier Cros's avatar
Olivier Cros committed
71 72
    "Return the type (IPv4 or IPv6) of an ip"
    (host, port) = address
73

Olivier Cros's avatar
Olivier Cros committed
74 75 76 77 78 79 80
    for af_type in SOCKET_CONNECTORS_DICT.keys():
        try :
            socket.inet_pton(af_type, host)
        except:
            continue
        else:
            break
81 82
    else:
        raise ValueError("Unknown type of host", host)
Olivier Cros's avatar
Olivier Cros committed
83 84 85
    return af_type

def getConnectorFromAddress(address):
86
    address_type = getAddressType(address)
Olivier Cros's avatar
Olivier Cros committed
87
    return SOCKET_CONNECTORS_DICT[address_type]
88

Olivier Cros's avatar
Olivier Cros committed
89 90 91 92 93 94 95 96 97 98 99 100 101
def parseNodeAddress(address, port_opt=None):
    if ']' in address:
       (ip, port) = address.split(']')
       ip = ip.lstrip('[')
       port = port.lstrip(':')
       if port == '':
           port = port_opt
    elif ':' in address:
        (ip, port) = address.split(':')
        ip = resolve(ip)
    else:
        ip = address
        port = port_opt
102

Olivier Cros's avatar
Olivier Cros committed
103 104 105
    if port is None:
        raise ValueError
    return (ip, int(port))
106

107
def parseMasterList(masters, except_node=None):
Olivier Cros's avatar
Olivier Cros committed
108
    assert masters, 'At least one master must be defined'
109
    # load master node list
110
    socket_connector = None
111
    master_node_list = []
112
    for node in masters.split(' '):
Vincent Pelletier's avatar
Vincent Pelletier committed
113 114
        if not node:
            continue
Olivier Cros's avatar
Olivier Cros committed
115 116
        address = parseNodeAddress(node)

117 118
        if (address != except_node):
            master_node_list.append(address)
119

Olivier Cros's avatar
Olivier Cros committed
120
        socket_connector_temp = getConnectorFromAddress(address)
121
        if socket_connector is None:
Olivier Cros's avatar
Olivier Cros committed
122
            socket_connector = socket_connector_temp
123
        elif socket_connector != socket_connector_temp:
Vincent Pelletier's avatar
Vincent Pelletier committed
124 125
            raise TypeError("Wrong connector type : you're trying to use "
                "ipv6 and ipv4 simultaneously")
126
    return master_node_list, socket_connector
127

128 129 130 131 132 133
class Enum(dict):
    """
    Simulate an enumeration, define them as follow :
        class MyEnum(Enum):
          ITEM1 = Enum.Item(0)
          ITEM2 = Enum.Item(1)
134
    Enum items must be written in full upper case
135 136 137 138
    """

    class Item(int):

139 140 141
        _enum = None
        _name = None

142 143 144 145 146 147 148 149 150 151 152 153 154
        def __new__(cls, value):
            instance = super(Enum.Item, cls).__new__(cls, value)
            instance._enum = None
            instance._name = None
            return instance

        def __str__(self):
            return self._name

        def __repr__(self):
            return "<EnumItem %s (%d)>" % (self._name, self)

        def __eq__(self, other):
155 156
            if other is None:
                return False
157
            assert isinstance(other, (Enum.Item, int, float, long))
158 159
            if isinstance(other, Enum):
                assert self._enum == other._enum
160 161 162
            return int(self) == int(other)

    def __init__(self):
163
        dict.__init__(self)
164 165 166 167 168 169 170
        for name in dir(self):
            if not re.match('^[A-Z_]*$', name):
                continue
            item = getattr(self, name)
            item._name = name
            item._enum = self
            self[int(item)] = item
171

172 173 174
    def getByName(self, name):
        return getattr(self, name)

175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196

class ReadBuffer(object):
    """
        Implementation of a lazy buffer. Main purpose if to reduce useless
        copies of data by storing chunks and join them only when the requested
        size is available.
    """

    def __init__(self):
        self.size = 0
        self.content = deque()

    def append(self, data):
        """ Append some data and compute the new buffer size """
        size = len(data)
        self.size += size
        self.content.append((size, data))

    def __len__(self):
        """ Return the current buffer size """
        return self.size

197 198 199 200 201
    def read(self, size):
        """ Read and consume size bytes """
        if self.size < size:
            return None
        self.size -= size
202 203 204
        chunk_list = []
        pop_chunk = self.content.popleft
        append_data = chunk_list.append
205 206
        to_read = size
        chunk_len = 0
207
        # select required chunks
208
        while to_read > 0:
209
            chunk_size, chunk_data = pop_chunk()
210
            to_read -= chunk_size
211
            append_data(chunk_data)
212
        if to_read < 0:
213 214
            # too many bytes consumed, cut the last chunk
            last_chunk = chunk_list[-1]
215 216
            keep, let = last_chunk[:to_read], last_chunk[to_read:]
            self.content.appendleft((-to_read, let))
217 218
            chunk_list[-1] = keep
        # join all chunks (one copy)
219
        data = ''.join(chunk_list)
220 221 222 223 224 225 226 227
        assert len(data) == size
        return data

    def clear(self):
        """ Erase all buffer content """
        self.size = 0
        self.content.clear()