ctl.py 10.1 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165
import logging, socket, struct
from collections import namedtuple
from . import utils

uint16 = struct.Struct("!H")
header = struct.Struct("!HI")

class Struct(object):

    def __init__(self, format, *args):
        if args:
            t = namedtuple(*args)
        if isinstance(format, str):
            s = struct.Struct("!" + format)
            def encode(buffer, value):
                buffer += s.pack(*value)
            def decode(buffer, offset=0):
                return offset + s.size, t(*s.unpack_from(buffer, offset))
        else:
            def encode(buffer, value):
                for f, value in zip(format, value):
                    f.encode(buffer, value)
            def decode(buffer, offset=0):
                r = []
                for f in format:
                    offset, x = f.decode(buffer, offset)
                    r.append(x)
                return offset, t(*r)
        self.encode = encode
        self.decode = decode

class Array(object):

    def __init__(self, item):
        self._item = item

    def encode(self, buffer, value):
        buffer += uint16.pack(len(value))
        encode = self._item.encode
        for value in value:
            encode(buffer, value)

    def decode(self, buffer, offset=0):
        r = []
        o = offset + 2
        decode = self._item.decode
        for i in xrange(*uint16.unpack_from(buffer, offset)):
            o, x = decode(buffer, o)
            r.append(x)
        return o, r

class String(object):

    @staticmethod
    def encode(buffer, value):
        buffer += value + "\0"

    @staticmethod
    def decode(buffer, offset=0):
        i = buffer.index("\0", offset)
        return i + 1, buffer[offset:i]


class Buffer(object):

    def __init__(self):
        self._buf = bytearray()
        self._r = self._w = 0


    def __iadd__(self, value):
        self._buf += value
        return self

    def __len__(self):
        return len(self._buf)

    def _seek(self, r):
        n = len(self._buf)
        if r < n:
            self._r = r
        else:
            self._w -= n
            del self._buf[:]
            self._r = 0

    # reading

    @property
    def ready(self):
        return self._w <= len(self._buf)

    def want(self, n):
        self._w = self._r + n

    def unpack_from(self, struct):
        r = self._r
        value = struct.unpack_from(self._buf, r)
        self._seek(r + struct.size)
        return value

    def decode(self, decode):
        r, value = decode(self._buf, self._r)
        self._seek(r)
        return value

    try: # BBB: Python < 2.7.4 (http://bugs.python.org/issue10212)
        uint16.unpack_from(bytearray(uint16.size))
    except TypeError:
        def unpack_from(self, struct):
            r = self._r
            x = r + struct.size
            value = struct.unpack(buffer(self._buf)[r:x])
            self._seek(x)
            return value
        def decode(self, decode):
            r = self._r
            size, value = decode(buffer(self._buf)[r:])
            self._seek(r + size)
            return value

    # writing

    def send(self, socket, *args):
        r = self._r
        self._seek(r + socket.send(self._buf[r:], *args))

    def pack_into(self, struct, offset, *args):
        struct.pack_into(self._buf, offset, *args)


class Packet(object):

    response_dict = {}

    def __new__(cls, id, request, response=None):
        if response:
            cls.response_dict[id] = response.decode
        if request:
            def packet(*args):
                self = object.__new__(cls)
                self.id = id
                self.args = args
                self.request = request
                return self
            return packet

    def write(self, buffer):
        logging.trace('send %s%r', self.__class__.__name__,
                                   (self.id,) + self.args)
        offset = len(buffer)
        buffer += '\0' * header.size
        r = self.request
        if isinstance(r, Struct):
            r.encode(buffer, self.args)
        else:
            r.encode(buffer, *self.args)
        buffer.pack_into(header, offset, self.id,
                         len(buffer) - header.size - offset)


Dump = Packet(1,
  Struct("B"),
  Struct((
    Array(Struct((Struct("I", "index", "index"), String), "interface", "index name")),
166
    Array(Struct("16sIHHHHHiHH", "neighbour", "address ifindex reach rxcost txcost rtt rttcost channel if_up cost_multiplier")),
167 168 169 170
    Array(Struct("16sBH", "xroute", "prefix plen metric")),
    Array(Struct("16sBHHH8siiI16s16sB", "route", "prefix plen metric smoothed_metric refmetric id seqno age ifindex neigh_address nexthop flags")),
    ), "dump", "interfaces neighbours xroutes routes"))

171 172 173 174
SetCostMultiplier = Packet(2,
  Struct("16sIH"),
  Struct("B", "set_cost_multiplier", "flags"))

175

176 177 178 179
class BabelException(Exception): pass


class ConnectionClosed(BabelException):
180 181 182 183 184

    def __str__(self):
        return "connection to babeld closed (%s)" % self.args


185 186 187 188 189
class Babel(object):

    _decode = None

    def __init__(self, socket_path, handler, network):
190 191
        self.socket_path = socket_path
        self.handler = handler
192
        self.network = network
193 194 195 196 197 198 199 200
        self.locked = set()
        self.reset()

    def reset(self):
        try:
            del self.socket, self.request_dump
        except AttributeError:
            pass
201 202 203
        self.write_buffer = Buffer()
        self.read_buffer = Buffer()
        self.read_buffer.want(header.size)
204 205
        s = socket.socket(socket.AF_UNIX,
            socket.SOCK_STREAM | socket.SOCK_CLOEXEC)
206 207
        def select(*args):
            try:
208 209
                s.connect(self.socket_path)
            except socket.error, e:
210
                logging.debug("Can't connect to %r (%r)", self.socket_path, e)
211
                return e
212 213
            s.send("\1")
            s.setblocking(0)
214
            del self.select
215 216 217
            self.socket = s
            return self.select(*args)
        self.select = select
218
        self.close = s.close
219 220 221 222 223 224 225 226

    def request_dump(self):
        if self.select({}, {}, ()):
            self.handle_dump((), (), (), ())
        else:
            # interfaces + neighbours + installed routes
            self.request_dump = lambda: self.send(Dump(11))
            self.request_dump()
227 228 229 230 231 232 233 234 235 236 237 238 239

    def send(self, packet):
        packet.write(self.write_buffer)

    def select(self, r, w, t):
        s = self.socket
        r[s] = self._read
        if self.write_buffer:
            w[s] = self._write

    def _read(self):
        d = self.socket.recv(65536)
        if not d:
240
            raise ConnectionClosed(self.socket_path)
241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264
        b = self.read_buffer
        b += d
        while b.ready:
            if self._decode:
                packet = b.decode(self._decode)
                self._decode = None
                b.want(header.size)
                name = packet.__class__.__name__
                logging.trace('recv %r', packet)
                try:
                    h = getattr(self, "handle_" + name)
                except AttributeError:
                    h = getattr(self.handler, "babel_" + name)
                h(*packet)
            else:
                packet_type, size = b.unpack_from(header)
                self._decode = Packet.response_dict[packet_type]
                b.want(size)

    def _write(self):
        self.write_buffer.send(self.socket)

    def handle_dump(self, interfaces, neighbours, xroutes, routes):
        # neighbours = {neigh_prefix: (neighbour, {dst_prefix: route})}
265
        n = {(n.address, n.ifindex): (n, {}) for n in neighbours}
266 267 268 269 270
        unidentified = set(n)
        self.neighbours = neighbours = {}
        a = len(self.network)
        for route in routes:
            assert route.flags & 1, route # installed
271 272
            if route.prefix.startswith('\0\0\0\0\0\0\0\0\0\0\xff\xff'):
                continue
273 274 275 276 277 278 279 280
            assert route.neigh_address == route.nexthop, route
            address = route.neigh_address, route.ifindex
            neigh_routes = n[address]
            ip = utils.binFromRawIp(route.prefix)
            if ip[:a] == self.network:
                prefix = ip[a:route.plen]
                if prefix and not route.refmetric:
                    neighbours[prefix] = neigh_routes
281 282 283 284 285 286 287 288 289 290 291 292 293 294
                    # XXX: Temporary logging to understand when a KeyError
                    #      happens. Then, we'll problably replace 'remove' by
                    #      'discard'.
                    try:
                        unidentified.remove(address)
                    except KeyError as e:
                        logging.warning("address: %s; prefix: %s",
                                        address, prefix)
                        logging.warning("neighbours: %r", neighbours)
                        logging.warning("routes: %r", routes)
                        try:
                            tm.sendto(tm.cache.registry_prefix, '\7%s' % e)
                        except AttributeError:
                            pass
295 296 297
            else:
                prefix = None
            neigh_routes[1][prefix] = route
298
        self.locked.clear()
299 300 301
        if unidentified:
            routes = {}
            for address in unidentified:
302 303 304 305
                neigh, r = n[address]
                if not neigh.cost_multiplier:
                    self.locked.add(address)
                routes.update(r)
306 307 308 309
            if routes:
                neighbours[None] = None, routes
                logging.trace("Routes via unidentified neighbours. %r",
                              neighbours)
310
        self.interfaces = {i.index: name for i, name in interfaces}
311
        self.handler.babel_dump()
312 313 314

    def handle_set_cost_multiplier(self, flags):
        pass
315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335


class iterRoutes(object):

    _waiting = True

    def __new__(cls, control_socket, network):
        self = object.__new__(cls)
        c = Babel(control_socket, self, network)
        c.request_dump()
        while self._waiting:
            args = {}, {}, ()
            c.select(*args)
            utils.select(*args)
        return (prefix
            for neigh_routes in c.neighbours.itervalues()
            for prefix in neigh_routes[1]
            if prefix)

    def babel_dump(self):
        self._waiting = False