dnsanycast

#!/usr/bin/python2.5

import socket
import sys
import os

from twisted.internet import reactor, protocol
from twisted.python import log

from pybgp import speaker, pathattr, proto

import tcpcheck

# This script is a simple example of the pybgp library. The setup below
# assumes the following config:
#
# router 192.168.1.1 AS# 65000
#  |
# us     192.168.1.2 AS# 65001
#
# ...and virtual IPs bound to the loopback interface, on which a DNS
# server will be listening - e.g. under linux:
#
# ip addr add 192.168.2.1/32 dev lo
# ip addr add 192.168.2.2/32 dev lo
# rndc reload
#
# As long as the virtual IPs continue to answer TCP port 53 connects the
# VIPs will be advertised by eBGP to the router

class Checker(tcpcheck.Checker):
    def change(self, old, new):
        if new=='up':
            self.bgp.advertise('%s/32' % (self.host,))
        else:
            self.bgp.withdraw('%s/32' % (self.host,))

bgp_connect = protocol.ClientCreator(reactor, speaker.BGP)

class Peer:
    def __init__(self, process, ip, remoteas):
        self.proto = None
        self.host = ip
        self.asnum = remoteas
        self.process = process

        self.reinit()

    def reinit(self, reconnect=False):
        self.state = 'idle'
        self.routes = None
        if reconnect:
            reactor.callLater(5, self.start)

    def start(self):
        if self.state!='idle':
            raise Exception('already connected')

        self.routes = set()
        self.state = 'connecting'

        d = bgp_connect.connectTCP(self.host, 179)
        d.addCallbacks(self.ok, self.err)

    def closed(self, reason):
        log.msg("connection to", self.host, "lost", reason)
        self.reinit(True)

    def ok(self, protocol):
        self.proto = protocol
        self.proto.closed = self.closed
        self.state = 'connected'

        log.msg("peer connected", self.host)

        # hook up the message callback to ourselves
        self.proto.handle_msg = self.msg

        # send an open
        self.proto.open(self.process.asnum, self.process.bgpid)

    def err(self, reason):
        log.msg("connection to", self.host, "failed", reason)

        self.reinit()
        reactor.callLater(5, self.start)

    def msg(self, msg):
        if msg.kind=='open':
            if self.state!='connected':
                raise Exception('open at invalid time')

            log.msg("peer opened", self.host)

            self.state = 'open'
            self.proto.start_timer(msg.holdtime)

            reactor.callLater(0.5, self.send_updates)

        elif msg.kind=='notification':
            log.msg("notification from", self.host, msg)
            self.proto.transport.loseConnection()

            self.reinit(True)

    def send_updates(self):
        if self.state!='open':
            return

        current = self.process.routes()
        to_send = []
        to_withdraw = []

        for route in current:
            if not route in self.routes:
                self.routes.add(route)
                to_send.append(route)

        for route in self.routes:
            if not route in current:
                to_withdraw.append(route)

        for route in to_withdraw:
            self.routes.remove(route)

        if to_send:
            up = proto.Update(
                    pathattr.Origin('incomplete'),
                    pathattr.Med(0),
                    pathattr.AsPath([[self.process.asnum]]),
                    pathattr.NextHop(self.process.bgpid),
                    nlri=to_send,
                    withdraw=to_withdraw,
                    )
            log.msg("sending update to", self.host, up)
            self.proto.send(up)
            

        elif to_withdraw:
            up = proto.Update(withdraw=to_withdraw)
            log.msg("sending update to", self.host, up)
            self.proto.send(up)

class BgpProcess:
    def __init__(self, asnum, bgpid):
        self.asnum = asnum
        self.bgpid = bgpid

        self.peers = {}

        self._routes = {}
        self.timer = None

    def peer(self, ip, remoteas):
        if ip in self.peers:
            raise Exception('connection already established')

        p = Peer(self, ip, remoteas)
        p.start()

        self.peers[ip] = p

    def advertise(self, prefix):
        self._routes[prefix] = True

        if self.timer:
            log.msg('advertising', prefix, '(coalescing)')
            self.timer.delay(0.1)
        else:
            log.msg('advertising', prefix)
            self.timer = reactor.callLater(0.1, self.update)

    def withdraw(self, prefix):
        if not prefix in self._routes:
            return
        del self._routes[prefix]

        if self.timer:
            log.msg('withdrawing', prefix, '(coalescing)')
            self.timer.delay(0.1)
        else:
            log.msg('withdrawing', prefix)
            self.timer = reactor.callLater(0.1, self.update)

    def update(self):
        self.timer = None
        for ip, peer in self.peers.items():
            log.msg('asking peer', ip, 'to update')
            peer.send_updates()

    def routes(self):
        return self._routes.keys()
        
def main(bgpid):
    b = BgpProcess(65001, bgpid)
    b.peer('192.168.1.254', 65000)
    b.peer('192.168.1.253', 65000)

    for ip in ('192.168.2.1', '192.168.2.2'):
        c = Checker(ip, 53, '127.0.0.1')
        c.bgp = b
        c.start(5)

if __name__=='__main__':
    myip = socket.gethostbyname(os.uname()[1])

    log.startLogging(sys.stderr)

    reactor.callWhenRunning(main, myip)
    reactor.run()