From 8d686a98fb6881a346fa6dc83fcc64266ccd13d6 Mon Sep 17 00:00:00 2001
From: Artem Kotik <miaow2@yandex.ru>
Date: Mon, 23 Oct 2023 14:48:38 +0400
Subject: [PATCH] Add status field to ConfigDiffScript

---
 docs/colliecting-diffs.md              |  2 +-
 netbox_config_diff/compliance/base.py  | 16 +++++++++-------
 netbox_config_diff/compliance/utils.py | 11 +++++++++++
 3 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/docs/colliecting-diffs.md b/docs/colliecting-diffs.md
index 35d3057..20e3d30 100644
--- a/docs/colliecting-diffs.md
+++ b/docs/colliecting-diffs.md
@@ -32,7 +32,7 @@ In the script, you can define a site, on which devices run compliance, or device
  If you define both fields, script will run only on devices from `Devices` field
 
 !!! warning
-    Script runs only on devices with status `Active`, assigned Primary IP, Config Template, Platform and PlatformSetting
+    Script runs only on devices with assigned Primary IP, Config Template, Platform and PlatformSetting
 
 If you have configs in NetBox DataSource, you can define it, the script instead of connecting to devices will find configs in DataSource by device's names.
 
diff --git a/netbox_config_diff/compliance/base.py b/netbox_config_diff/compliance/base.py
index f902019..e45022f 100644
--- a/netbox_config_diff/compliance/base.py
+++ b/netbox_config_diff/compliance/base.py
@@ -18,27 +18,29 @@
 
 from .models import DeviceDataClass
 from .secrets import SecretsMixin
-from .utils import PLATFORM_MAPPING, exclude_lines, get_unified_diff
+from .utils import PLATFORM_MAPPING, CustomChoiceVar, exclude_lines, get_unified_diff
 
 
 class ConfigDiffBase(SecretsMixin):
     site = ObjectVar(
         model=Site,
         required=False,
-        description="Run compliance for devices (with status Active, "
-        "primary IP, platform and config template) in this site",
+        description="Run compliance for devices (with primary IP, platform and config template) in this site",
     )
     devices = MultiObjectVar(
         model=Device,
         required=False,
         query_params={
-            "status": DeviceStatusChoices.STATUS_ACTIVE,
             "has_primary_ip": True,
             "platform_id__n": "null",
             "config_template_id__n": "null",
         },
         description="If you define devices in this field, the Site field will be ignored",
     )
+    status = CustomChoiceVar(
+        choices=DeviceStatusChoices,
+        default=DeviceStatusChoices.STATUS_ACTIVE,
+    )
     data_source = ObjectVar(
         model=DataSource,
         required=False,
@@ -66,7 +68,7 @@ def validate_data(self, data: dict) -> Iterable[Device]:
             devices = (
                 data["devices"]
                 .filter(
-                    status=DeviceStatusChoices.STATUS_ACTIVE,
+                    status=data["status"],
                     platform__platform_setting__isnull=False,
                     config_template__isnull=False,
                 )
@@ -77,7 +79,7 @@ def validate_data(self, data: dict) -> Iterable[Device]:
         else:
             devices = Device.objects.filter(
                 site=data["site"],
-                status=DeviceStatusChoices.STATUS_ACTIVE,
+                status=data["status"],
                 platform__platform_setting__isnull=False,
                 config_template__isnull=False,
             ).exclude(
@@ -90,7 +92,7 @@ def validate_data(self, data: dict) -> Iterable[Device]:
 
         if not devices:
             self.log_warning(
-                "No matching devices found, devices must have status `Active`, primary IP, platform and platformsetting"
+                "No matching devices found, devices must have status primary IP, platform and platformsetting"
             )
         else:
             self.log_info(f"Working with device(s): {', '.join(d.name for d in devices)}")
diff --git a/netbox_config_diff/compliance/utils.py b/netbox_config_diff/compliance/utils.py
index d3729d5..c0aa4af 100644
--- a/netbox_config_diff/compliance/utils.py
+++ b/netbox_config_diff/compliance/utils.py
@@ -1,6 +1,9 @@
 import re
 from difflib import unified_diff
 
+from django.forms import ChoiceField
+from extras.scripts import ScriptVariable
+
 PLATFORM_MAPPING = {
     "arista_eos": "arista_eos",
     "cisco_aireos": "cisco_aireos",
@@ -17,6 +20,14 @@
 }
 
 
+class CustomChoiceVar(ScriptVariable):
+    form_field = ChoiceField
+
+    def __init__(self, choices, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.field_attrs["choices"] = choices
+
+
 def get_unified_diff(rendered_config: str, actual_config: str, device: str) -> str:
     diff = unified_diff(
         rendered_config.strip().splitlines(),