From d1ee380ffb2ac552329b9b22afb7fd20f9a11120 Mon Sep 17 00:00:00 2001
From: Cyril LEVIS <cyril.levis@dailymotion.com>
Date: Wed, 16 Feb 2022 11:22:51 +0100
Subject: [PATCH] feat: add custom_fields suppport

```bash
$ netbox_agent -du --device.custom_fields="last_run=$(date),last_run_ts=$(date +'%s')"
```

obviously, custom_fields need to be create manually in netbox admin
---
 README.md                 |  2 +-
 netbox_agent.yaml.example |  5 +++--
 netbox_agent/config.py    |  2 ++
 netbox_agent/server.py    | 15 +++++++++++++++
 4 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 12eb900..af6d2b2 100644
--- a/README.md
+++ b/README.md
@@ -111,7 +111,7 @@ network:
 # blade_role: "Blade"
 # server_role: "Server"
 # tags: server, blade, ,just a comma,delimited,list
-#
+# custom_fields: field1=value1,field2=value2#
 #
 # Can use this to set the tenant
 #
diff --git a/netbox_agent.yaml.example b/netbox_agent.yaml.example
index a769fb4..a1a8e02 100644
--- a/netbox_agent.yaml.example
+++ b/netbox_agent.yaml.example
@@ -18,14 +18,15 @@ network:
 # blade_role: "Blade"
 # server_role: "Server"
 # tags: server, blade, ,just a comma,delimited,list
-
+# custom_fields: field1=value1,field2=value2
+#
 #
 # Use this to set the tenant
 #
 #tenant:
 # driver: "file:/tmp/tenant"
 # regex: "(.*)"
-
+
 datacenter_location:
  driver: "cmd:cat /etc/qualification | tr [A-Z] [a-z]"
  regex: "datacenter: (?P<datacenter>[A-Za-z0-9]+)"
diff --git a/netbox_agent/config.py b/netbox_agent/config.py
index 6913f65..d820836 100644
--- a/netbox_agent/config.py
+++ b/netbox_agent/config.py
@@ -45,6 +45,8 @@ def get_config():
                    help="Command to output hostname, used as Device's name in netbox")
     p.add_argument('--device.tags', default=r'',
                    help='tags to use for a host')
+    p.add_argument('--device.custom_fields', default=r'',
+                   help='custom_fields to use for a host, eg: field1=v1,field2=v2')
     p.add_argument('--device.blade_role', default=r'Blade',
                    help='role to use for a blade server')
     p.add_argument('--device.chassis_role', default=r'Server Chassis',
diff --git a/netbox_agent/server.py b/netbox_agent/server.py
index 24758ae..df9a5a5 100644
--- a/netbox_agent/server.py
+++ b/netbox_agent/server.py
@@ -32,6 +32,15 @@ class ServerBase():
             x.strip() for x in config.device.tags.split(',') if x.strip()
         ])) if config.device.tags else []
         self.nb_tags = list(create_netbox_tags(self.tags))
+        config_cf = set([
+            f.strip() for f in config.device.custom_fields.split(",")
+            if f.strip()
+        ])
+        self.custom_fields = {}
+        self.custom_fields.update(dict([
+            (k.strip(), v.strip()) for k, v in
+            [f.split("=", 1) for f in config_cf]
+        ]))
 
     def get_tenant(self):
         tenant = Tenant()
@@ -195,6 +204,7 @@ class ServerBase():
             tenant=tenant.id if tenant else None,
             rack=rack.id if rack else None,
             tags=[{'name': x} for x in self.tags],
+            custom_fields=self.custom_fields,
         )
         return new_chassis
 
@@ -217,6 +227,7 @@ class ServerBase():
             tenant=tenant.id if tenant else None,
             rack=rack.id if rack else None,
             tags=[{'name': x} for x in self.tags],
+            custom_fields=self.custom_fields,
         )
         return new_blade
 
@@ -438,6 +449,10 @@ class ServerBase():
             server.tags = [x.id for x in self.nb_tags]
             update += 1
 
+        if server.custom_fields != self.custom_fields:
+            server.custom_fields = self.custom_fields
+            update += 1
+
         if config.update_all or config.update_location:
             ret, server = self.update_netbox_location(server)
             update += ret