Merge pull request #12 from Solvik/feature/rack_awareness

Rack awareness

netbox_agent.yaml.example

@@ -9,4 +9,12 @@ datacenter_location:
 # regex: 'SysName: .*\.([A-Za-z0-9]+)'
 #
 # driver: "file:/tmp/datacenter"
+# regex: "(.*)"
+
+rack_location:
+# driver: 'cmd:lldpctl'
+# match SysName: sw-dist-a1.dc42
+# regex: 'SysName:[ ]+[A-Za-z]+-[A-Za-z]+-([A-Za-z0-9]+)'
+#
+# driver: "file:/tmp/datacenter"
 # regex: "(.*)"

netbox_agent/config.py

@@ -10,7 +10,21 @@ netbox_instance = pynetbox.api(
     token=config['netbox']['token']
 )
 
+DATACENTER_LOCATION_DRIVER_FILE = None
+DATACENTER_LOCATION = None
+DATACENTER_LOCATION_REGEX = None
+RACK_LOCATION_DRIVER_FILE = None
+RACK_LOCATION = None
+RACK_LOCATION_REGEX = None
 
-DATACENTER_LOCATION_DRIVER_FILE = config.get('datacenter_location').get('driver_file')
-DATACENTER_LOCATION = config.get('datacenter_location').get('driver')
-DATACENTER_LOCATION_REGEX = config.get('datacenter_location').get('regex')
+if config.get('datacenter_location'):
+    dc_loc = config.get('datacenter_location')
+    DATACENTER_LOCATION_DRIVER_FILE = dc_loc.get('driver_file')
+    DATACENTER_LOCATION = dc_loc.get('driver')
+    DATACENTER_LOCATION_REGEX = dc_loc.get('regex')
+
+if config.get('rack_location'):
+    rack_location = config['rack_location']
+    RACK_LOCATION_DRIVER_FILE = rack_location.get('driver_file')
+    RACK_LOCATION = rack_location.get('driver')
+    RACK_LOCATION_REGEX = rack_location.get('regex')

netbox_agent/datacenter.py

@@ -1,46 +0,0 @@
-import importlib
-import importlib.machinery
-
-from netbox_agent.config import DATACENTER_LOCATION, DATACENTER_LOCATION_DRIVER_FILE, \
-    DATACENTER_LOCATION_REGEX
-
-
-class Datacenter():
-    """
-    This class is used to guess the datacenter in order to push the information
-    in Netbox for a `Device`
-
-    A driver takes a `value` and evaluates a regex with a `capture group`.
-
-    There's embeded drivers such as `file` or `cmd` which read a file or return the
-    output of a file.
-
-    There's also a support for an external driver file outside of this project in case
-    the logic isn't supported here.
-    """
-    def __init__(self, *args, **kwargs):
-        self.driver = DATACENTER_LOCATION.split(':')[0]
-        self.driver_value = ':'.join(DATACENTER_LOCATION.split(':')[1:])
-        self.driver_file = DATACENTER_LOCATION_DRIVER_FILE
-
-        if self.driver_file:
-            try:
-                # FIXME: Works with Python 3.3+, support older version?
-                loader = importlib.machinery.SourceFileLoader('driver_file', self.driver_file)
-                self.driver = loader.load_module()
-            except ImportError:
-                raise ImportError("Couldn't import {} as a module".format(self.driver_file))
-        else:
-            try:
-                self.driver = importlib.import_module(
-                    'netbox_agent.drivers.datacenter_{}'.format(self.driver)
-                )
-            except ImportError:
-                raise ImportError("Driver {} doesn't exists".format(self.driver))
-
-    def get(self):
-        if not hasattr(self.driver, 'get'):
-            raise Exception(
-                "Your driver {} doesn't have a get() function, please fix it".format(self.driver)
-            )
-        return getattr(self.driver, 'get')(self.driver_value, DATACENTER_LOCATION_REGEX)

netbox_agent/location.py

@@ -0,0 +1,71 @@
+import importlib
+import importlib.machinery
+
+from netbox_agent.config import DATACENTER_LOCATION, DATACENTER_LOCATION_DRIVER_FILE, \
+    DATACENTER_LOCATION_REGEX, RACK_LOCATION, RACK_LOCATION_DRIVER_FILE, RACK_LOCATION_REGEX
+
+
+class LocationBase():
+    """
+    This class is used to guess the location in order to push the information
+    in Netbox for a `Device`
+
+    A driver takes a `value` and evaluates a regex with a `capture group`.
+
+    There's embeded drivers such as `file` or `cmd` which read a file or return the
+    output of a file.
+
+    There's also a support for an external driver file outside of this project in case
+    the logic isn't supported here.
+    """
+    def __init__(self, driver, driver_value, driver_file, regex, *args, **kwargs):
+        self.driver = driver
+        self.driver_value = driver_value
+        self.driver_file = driver_file
+        print(self.driver_file)
+        self.regex = regex
+
+        if self.driver_file:
+            print('if', self.driver_file)
+            try:
+                # FIXME: Works with Python 3.3+, support older version?
+                loader = importlib.machinery.SourceFileLoader('driver_file', self.driver_file)
+                self.driver = loader.load_module()
+            except ImportError:
+                raise ImportError("Couldn't import {} as a module".format(self.driver_file))
+        else:
+            if self.driver:
+                try:
+                    self.driver = importlib.import_module(
+                        'netbox_agent.drivers.{}'.format(self.driver)
+                    )
+                except ImportError:
+                    raise ImportError("Driver {} doesn't exists".format(self.driver))
+
+    def get(self):
+        if self.driver is None:
+            return None
+        if not hasattr(self.driver, 'get'):
+            raise Exception(
+                "Your driver {} doesn't have a get() function, please fix it".format(self.driver)
+            )
+        return getattr(self.driver, 'get')(self.driver_value, self.regex)
+
+
+class Datacenter(LocationBase):
+    def __init__(self):
+        driver = DATACENTER_LOCATION.split(':')[0] if DATACENTER_LOCATION else None
+        driver_value = ':'.join(DATACENTER_LOCATION.split(':')[1:]) if DATACENTER_LOCATION \
+            else None
+        driver_file = DATACENTER_LOCATION_DRIVER_FILE
+        regex = DATACENTER_LOCATION_REGEX
+        super().__init__(driver, driver_value, driver_file, regex)
+
+
+class Rack(LocationBase):
+    def __init__(self):
+        driver = RACK_LOCATION.split(':')[0] if RACK_LOCATION else None
+        driver_value = ':'.join(RACK_LOCATION.split(':')[1:]) if RACK_LOCATION else None
+        driver_file = RACK_LOCATION_DRIVER_FILE
+        regex = RACK_LOCATION_REGEX
+        super().__init__(driver, driver_value, driver_file, regex)

netbox_agent/server.py

@@ -2,8 +2,9 @@ from pprint import pprint
 import socket
 
 from netbox_agent.config import netbox_instance as nb
-from netbox_agent.location import Datacenter
+
 import netbox_agent.dmidecode as dmidecode
+from netbox_agent.location import Datacenter, Rack
 from netbox_agent.network import Network
 
 
@@ -28,6 +29,17 @@ class ServerBase():
         )
         return datacenter
 
+    def get_rack(self):
+        rack = Rack()
+        return rack.get()
+
+    def get_netbox_rack(self):
+        rack = nb.dcim.racks.get(
+            name=self.get_rack(),
+            datacenter=self.get_netbox_datacenter(),
+        )
+        return rack
+
     def get_product_name(self):
         """
         Return the Chassis Name from dmidecode info
@@ -217,6 +229,8 @@ class ServerBase():
         # FIXME: do something more generic by looping on every get_* methods
         print('Datacenter:', self.get_datacenter())
         print('Netbox Datacenter:', self.get_netbox_datacenter())
+        print('Rack:', self.get_rack())
+        print('Netbox Rack:', self.get_netbox_rack())
         print('Is blade:', self.is_blade())
         print('Product Name:', self.get_product_name())
         print('Chassis:', self.get_chassis())