add exception hierarchy, refactor parsing code, add tests for speaker code

Author: philmayers

pybgp/exceptions.py

@@ -0,0 +1,38 @@
+import struct
+
+class BgpExc(Exception):
+    data = None
+
+    send_error = True
+
+    def __str__(self):
+        return '<%s %r>' % (self.__class__.__name__, self.data)
+
+class NotSync(BgpExc):
+    code = 1
+    subcode = 1
+
+
+class BadLen(BgpExc):
+    code = 1
+    subcode = 2
+
+    def __init__(self, msg, len):
+        self.msg = msg
+        self.len = len
+        self.data = struct.pack('!H', len)
+
+    def __str__(self):
+        return '<BadLen %d msgtype=%d>' % (self.len,self.msg)
+
+class BadMsg(BgpExc):
+    code = 1
+    subcode = 3
+
+    def __init__(self, msg):
+        self.msg = msg
+        self.data = struct.pack('B', msg)
+
+    def __str__(self):
+        return '<BadMsg %d>' % (self.msg,)
+

pybgp/proto.py

@@ -5,7 +5,7 @@
 from odict import OrderedDict as OD
 
 
-from pybgp import nlri, pathattr
+from pybgp import nlri, pathattr, exceptions
 
 class Open:
     kind = 'open'
@@ -28,8 +28,6 @@ def from_bytes(cls, bytes):
 
         self.version, self.asnum, self.holdtime, bgpid, paramlen = struct.unpack_from('!BHH4sB', bytes)
         self.bgpid = socket.inet_ntoa(bgpid)
-        if paramlen==0:
-            return
 
         assert len(bytes) == 10 + paramlen, "message too short?"
 
@@ -133,6 +131,11 @@ class Notification:
     kind = 'notification'
     number = 3
 
+    def __init__(self, code, subcode, data=''):
+        self.code = code
+        self.subcode = subcode
+        self.data = data
+
     REASONS = {
         (1,1): 'not synchronised',
         (1,2): 'bad message len',
@@ -157,12 +160,10 @@ class Notification:
     }
 
     def from_bytes(cls, bytes):
-        self = cls()
+        code, subcode = struct.unpack_from('BB', bytes)
+        data = bytes[2:]
 
-        self.code, self.subcode = struct.unpack_from('BB', bytes)
-        self.data = bytes[2:]
-
-        return self
+        return cls(code, subcode, data)
     from_bytes = classmethod(from_bytes)
 
     def __str__(self):
@@ -171,6 +172,13 @@ def __str__(self):
             return 'Notification "%s" params %r' % (self.REASONS[c,s], self.data)
         return 'Notification message code=%d subcode=%d params=%r' % (self.code, self.subcode, self.data)
 
+    def encode(self):
+        v = struct.pack('BB', self.code, self.subcode)
+        if self.data:
+            v += self.data
+        return v
+
+
 class Update:
     kind = 'update'
     number = 2
@@ -266,3 +274,31 @@ def encode(self):
 
         return v
 
+
+class ProtoBase:
+    def parse_payload(self, type, payload):
+        length = len(payload)
+        totlen = length + 19
+
+        if type==1:
+            if length<10:
+                raise exceptions.BadLen(type, totlen)
+            return Open.from_bytes(payload)
+
+        elif type==2:
+            if length<4:
+                raise exceptions.BadLen(type, totlen)
+            return Update.from_bytes(payload)
+
+        elif type==3:
+            if length<2:
+                raise exceptions.BadLen(type, totlen)
+            return Notification.from_bytes(payload)
+
+        elif type==4:
+            if length:
+                print "keepalive with data? %r" % (payload,)
+            return Keepalive.from_bytes(payload)
+
+        raise exceptions.BadMsg(type)
+

pybgp/speaker.py

@@ -4,9 +4,9 @@
 from twisted.internet import reactor, protocol, task
 from twisted.python import log
 
-from pybgp import nlri, pathattr, proto
+from pybgp import nlri, pathattr, proto, exceptions
 
-class BGP(protocol.Protocol):
+class BGP(protocol.Protocol, proto.ProtoBase):
 
     def connectionMade(self):
         self.buffer = ''
@@ -37,35 +37,33 @@ def dataReceived(self, data):
             if len(self.buffer) < 19:
                 return
             auth, length, type = struct.unpack('!16sHB', self.buffer[:19])
+
+            if auth!='\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff\xff':
+                return self.notify(exceptions.NotSync())
+
+            if length < 19 or length > 4096:
+                return self.notify(exceptions.BadLen(type, length))
+
             if len(self.buffer) < length:
                 return
 
-            msg = self.buffer[:length]
+            payload = self.buffer[19:length]
             self.buffer = self.buffer[length:]
 
-            msg = self.parse_msg(msg)
+            try:
+                msg = self.parse_payload(type, payload)
+            except exceptions.BgpExc, ex:
+                return self.notify(ex)
 
             self._handle_msg(msg)
 
-    def parse_msg(self, msg):
-        auth, length, type = struct.unpack('!16sHB', msg[:19])
-        payload = msg[19:]
-        del msg
-
-        if type==1:
-            return proto.Open.from_bytes(payload)
-
-        elif type==2:
-            return proto.Update.from_bytes(payload)
-
-        elif type==3:
-            return proto.Notification.from_bytes(payload)
-
-        elif type==4:
-            return proto.Keepalive.from_bytes(payload)
-
-        else:
-            raise Exception('invalid message')
+    def notify(self, ex):
+        if ex.send_error:
+            log.msg("sending notify", ex)
+            notify = proto.Notification(ex.code, ex.subcode, ex.data)
+            self.send(notify)
+        log.msg("disconnecting")
+        self.transport.loseConnection()
 
     def send(self, msg):
         body = msg.encode()

pybgp/test/test_proto.py

@@ -2,7 +2,7 @@
 
 import unittest
 
-from pybgp import proto, pathattr, nlri
+from pybgp import proto, pathattr, nlri, exceptions
 
 class TestOpen(unittest.TestCase):
     def test_decode(self):