def test_wrong_class(self): ''' Try to register a class of the wrong type ''' class WrongType: pass with self.assertRaisesRegexp(ValueError, 'subclass'): type_registry.register_type_class(WrongType)
def test_wrong_message_type(self): ''' Try to register a class of the wrong type ''' class WrongType(ControlMessage): message_type = 16 with self.assertRaisesRegexp(ValueError, 'message.type'): type_registry.register_type_class(WrongType)
def test_register_type_class(self): ''' Normal registration of a type class ''' class ValidType(ControlMessage): message_type = 1 type_registry.register_type_class(ValidType) the_class = type_registry.get_type_class(1) self.assertEqual(the_class, ValidType)
def test_duplicate_registration_of_same_class(self): ''' Try to register the same class twice ''' class ValidType(ControlMessage): message_type = 1 type_registry.register_type_class(ValidType) the_class = type_registry.get_type_class(1) self.assertEqual(the_class, ValidType) type_registry.register_type_class(ValidType) the_class = type_registry.get_type_class(1) self.assertEqual(the_class, ValidType)
def test_bad_duplicate_registration(self): ''' Try to register two classes for the same type ''' class ValidType1(ControlMessage): message_type = 1 class ValidType2(ControlMessage): message_type = 1 type_registry.register_type_class(ValidType1) the_class = type_registry.get_type_class(1) self.assertEqual(the_class, ValidType1) with self.assertRaisesRegexp(ValueError, 'bound'): type_registry.register_type_class(ValidType2) the_class = type_registry.get_type_class(1) self.assertEqual(the_class, ValidType1)
# Add the nonce bitstream += BitArray(bytes=self.nonce) # Add the key-id and authentication data bitstream += BitArray('uint:16=%d, uint:16=%d, hex=%s' % (self.key_id, len(self.authentication_data), self.authentication_data.encode('hex'))) # Add the TTL bitstream += BitArray('uint:32=%d' % self.ttl) # Add reserved bits bitstream += self._reserved2 # Add the EID prefix mask length bitstream += BitArray('uint:8=%d' % self.eid_prefix.prefixlen) # Add the EID prefix bitstream += get_bitstream_for_afi_address(self.eid_prefix) # Add the reply bitstream += get_bitstream_for_afi_address(self.reply) return bitstream.bytes # Register this class in the registry type_registry.register_type_class(InfoMessage)
r''' Create bytes from properties >>> message = EncapsulatedControlMessage(payload='Dummy') >>> message.to_bytes() '\x80\x00\x00\x00Dummy' ''' # Verify that properties make sense self.sanitize() # Start with the type bitstream = BitArray('uint:4=%d' % self.message_type) # Add the flags bitstream += BitArray('bool=%d, bool=%d' % (self.security, self.ddt_originated)) # Add padding bitstream += self._reserved1 # Determine payload payload = self.payload if hasattr(payload, 'to_bytes'): payload = payload.to_bytes() return bitstream.bytes + payload # Register this class in the registry type_registry.register_type_class(EncapsulatedControlMessage)
# Add reserved bits bitstream += self._reserved1 # Add record count bitstream += BitArray('uint:8=%d' % len(self.records)) # Add the nonce bitstream += BitArray(bytes=self.nonce) # Add the key-id and authentication data bitstream += BitArray('uint:16=%d, uint:16=%d, hex=%s' % (self.key_id, len(self.authentication_data), self.authentication_data.encode('hex'))) # Add the map-reply records for record in self.records: bitstream += record.to_bitstream() # Add xTR-ID and site-ID if we said we would if has_xtr_site_id: bitstream += BitArray('uint:128=%d, uint:64=%d' % (self.xtr_id, self.site_id)) return bitstream.bytes # Register this class in the registry type_registry.register_type_class(MapNotifyMessage)
bitstream += BitArray(bytes=self.nonce) # Add the source EID bitstream += get_bitstream_for_afi_address(self.source_eid) # Add the ITR RLOCs for itr_rloc in self.itr_rlocs: bitstream += get_bitstream_for_afi_address(itr_rloc) # Add the EIDs for eid_prefix in self.eid_prefixes: # Make sure it is a network eid_prefix = ip_network(eid_prefix) # Add padding and prefix length bitstream += BitArray('uint:8=0, ' 'uint:8=%d' % eid_prefix.prefixlen) # Add the address bitstream += get_bitstream_for_afi_address(eid_prefix) # Add the map-reply record if present if self.map_reply is not None: bitstream += self.map_reply.to_bitstream() return bitstream.bytes # Register this class in the registry type_registry.register_type_class(MapRequestMessage)
return packet def to_bytes(self): ''' Create bytes from properties ''' # Verify that properties make sense self.sanitize() # Start with the type bitstream = BitArray('uint:4=%d' % self.message_type) # Add padding bitstream += self._reserved1 # Add the record count bitstream += BitArray('uint:8=%d' % len(self.records)) # Add the nonce bitstream += BitArray(bytes=self.nonce) # Add the records for record in self.records: bitstream += record.to_bitstream() return bitstream.bytes # Register this class in the registry type_registry.register_type_class(MapReferralMessage)
# Add the rest of the flags bitstream += BitArray('bool=%d' % self.want_map_notify) # Add record count bitstream += BitArray('uint:8=%d' % len(self.records)) # Add the nonce bitstream += BitArray(bytes=self.nonce) # Add the key-id and authentication data bitstream += BitArray('uint:16=%d, uint:16=%d, hex=%s' % (self.key_id, len(self.authentication_data), self.authentication_data.encode('hex'))) # Add the map-reply records for record in self.records: bitstream += record.to_bitstream() # Add xTR-ID and site-ID if we said we would if has_xtr_site_id: bitstream += BitArray('uint:128=%d, uint:64=%d' % (self.xtr_id, self.site_id)) return bitstream.bytes # Register this class in the registry type_registry.register_type_class(MapRegisterMessage)