Add pfx_table wrapper class.

michael-koeller · mroethke · commit d2c5d6eaa1c8 · 2021-01-04T10:39:35.000+01:00
diff --git a/rtrlib.cdef b/rtrlib.cdef
@@ -457,6 +457,16 @@ void pfx_table_free(struct pfx_table *pfx_table);
  */
 int pfx_table_add(struct pfx_table *pfx_table, const struct pfx_record *pfx_record);
 
+/**
+ * @brief Removes a pfx_record from a pfx_table.
+ * @param[in] pfx_table pfx_table to use.
+ * @param[in] pfx_record Record that will be removed.
+ * @return PFX_SUCCESS On success.
+ * @return PFX_ERROR On error.
+ * @return PFX_RECORD_NOT_FOUND If pfx_records couldn't be found.
+ */
+int pfx_table_remove(struct pfx_table *pfx_table, const struct pfx_record *pfx_record);
+
 /**
  * @brief Validates the origin of a BGP-Route.
  * @param[in] pfx_table pfx_table to use.
diff --git a/rtrlib/__init__.py b/rtrlib/__init__.py
@@ -9,6 +9,7 @@
 
 import logging
 
+from .pfx_table import PfxTable
 from .rtr_manager import RTRManager, PfxvState
 from .manager_group import ManagerGroupStatus
 
diff --git a/rtrlib/pfx_table.py b/rtrlib/pfx_table.py
@@ -0,0 +1,174 @@
+# -*- coding: utf8 -*-
+"""
+rtrlib.rtr_manager
+------------------
+
+"""
+
+from __future__ import absolute_import, unicode_literals
+
+
+from .exceptions import PFXException
+from .rtr_manager import ValidationResult
+from .util import ip_str_to_addr
+
+from _rtrlib import ffi, lib
+
+
+class PfxTable(object):
+    """
+    Wrapper class around pfx_table.
+
+    """
+
+    def __init__(self):
+        # allocate pfx_table
+        self.pfx_table = ffi.new('struct pfx_table *')
+        # initialize it
+        lib.pfx_table_init(self.pfx_table,
+                           ffi.NULL)
+        self.closed = False
+
+    @staticmethod
+    def _create_pfx_record(asn, ip, min_length, max_length):
+        record = ffi.new('struct pfx_record *')
+        lib.lrtr_ip_str_to_addr(ip.encode('ascii'), ffi.addressof(record, 'prefix'))
+
+        record.asn = asn
+        record.min_len = min_length
+        record.max_len = max_length
+        record.socket = ffi.NULL
+
+        return record
+
+    def add_record(self, asn, ip, min_length, max_length):
+        """
+        Add a BGP prefix to the table.
+
+        :param asn: autonomous system number
+        :type asn: int
+
+        :param ip: ip address
+        :type ip: str
+
+        :param min_length: minimum length of the subnet mask
+        :type min_length: int
+
+        :param max_length: minimum length of the subnet mask
+        :type max_length: int
+        """
+
+        record = self._create_pfx_record(asn, ip, min_length, max_length)
+
+        lib.pfx_table_add(self.pfx_table, record)
+
+    def remove_record(self, asn, ip, min_length, max_length):
+        """
+        Add a BGP prefix to the table.
+
+        :param asn: autonomous system number
+        :type asn: int
+
+        :param ip: ip address
+        :type ip: str
+
+        :param min_length: minimum length of the subnet mask
+        :type min_length: int
+
+        :param max_length: minimum length of the subnet mask
+        :type max_length: int
+        """
+
+        record = self._create_pfx_record(asn, ip, min_length, max_length)
+
+        lib.pfx_table_remove(self.pfx_table, record)
+
+    def validate(self, asn, prefix, mask_len):
+        """
+        Validate BGP prefix and returns state as ValidationResult object.
+        The reason list in the returned result will be empty.
+
+        :param asn: autonomous system number
+        :type asn: int
+
+        :param prefix: ip address
+        :type prefix: str
+
+        :param mask_len: length of the subnet mask
+        :type mask_len: int
+
+        :rtype: ValidationResult
+        """
+
+        result = ffi.new('enum pfxv_state *')
+
+        ret = lib.pfx_table_validate(self.pfx_table,
+                                     asn,
+                                     ip_str_to_addr(prefix),
+                                     mask_len,
+                                     result)
+
+        if ret == lib.PFX_ERROR:
+            raise PFXException("An error occurred during validation")
+
+        return ValidationResult(prefix,
+                                mask_len,
+                                asn,
+                                result[0])
+
+    def validate_r(self, asn, prefix, mask_len):
+        """
+        Validate BGP prefix and returns state as ValidationResult object.
+        The reason list in the returned result will contain a list of Reason objects.
+
+        :param asn: autonomous system number
+        :type asn: int
+
+        :param prefix: ip address
+        :type prefix: str
+
+        :param mask_len: length of the subnet mask
+        :type mask_len: int
+
+        :rtype: ValidationResult
+        """
+
+        result = ffi.new('enum pfxv_state *')
+
+        reason = ffi.new('struct pfx_record **')
+        reason[0] = ffi.NULL
+        reason_length = ffi.new('unsigned int *')
+        reason_length[0] = 0
+
+        ret = lib.pfx_table_validate_r(self.pfx_table,
+                                       reason,
+                                       reason_length,
+                                       asn,
+                                       ip_str_to_addr(prefix),
+                                       mask_len,
+                                       result)
+
+        if ret == lib.PFX_ERROR:
+            raise PFXException("An error occurred during validation")
+
+        return ValidationResult(prefix,
+                                mask_len,
+                                asn,
+                                result[0],
+                                reason,
+                                reason_length[0])
+
+    def close(self):
+        if not self.closed:
+            lib.pfx_table_free(self.pfx_table)
+            self.closed = True
+
+    def __enter__(self):
+        return self
+
+    def __del__(self):
+        self.close()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+        return False
