From f1d6a0d4a012198917add825dcc61d19845b071e Mon Sep 17 00:00:00 2001
From: Artem Kotik <miaow2@yandex.ru>
Date: Tue, 9 Apr 2024 21:43:50 +0400
Subject: [PATCH] Handle junipers templates with set commands

---
 netbox_config_diff/configurator/base.py      |  2 +-
 netbox_config_diff/configurator/platforms.py | 73 +++++++++++++++++++-
 2 files changed, 71 insertions(+), 4 deletions(-)

diff --git a/netbox_config_diff/configurator/base.py b/netbox_config_diff/configurator/base.py
index a6635a2..38bc94b 100644
--- a/netbox_config_diff/configurator/base.py
+++ b/netbox_config_diff/configurator/base.py
@@ -126,7 +126,7 @@ async def _collect_one_diff(self, device: ConfiguratorDeviceDataClass) -> None:
                 )
                 device.rendered_config = rendered_config
             else:
-                actual_config = await conn.get_config()
+                actual_config = await conn.get_config(config_template=device.rendered_config)
             device.actual_config = conn.clean_config(actual_config.result)
 
             device.diff = get_unified_diff(device.rendered_config, device.actual_config, device.name)
diff --git a/netbox_config_diff/configurator/platforms.py b/netbox_config_diff/configurator/platforms.py
index d4ffe97..45eb5a2 100644
--- a/netbox_config_diff/configurator/platforms.py
+++ b/netbox_config_diff/configurator/platforms.py
@@ -1,5 +1,5 @@
 import re
-from typing import Pattern
+from typing import Any, Pattern
 
 from scrapli_cfg.exceptions import TemplateError
 from scrapli_cfg.platform.core.arista_eos import AsyncScrapliCfgEOS
@@ -94,13 +94,17 @@ async def render_substituted_config(
         """
         self.logger.info("fetching configuration and replacing with provided substitutes")
 
-        source_config = await self.get_config(source=source)
+        source_config = await self.get_config(config_template=config_template, source=source)
         return source_config, self._render_substituted_config(
             config_template=config_template,
             substitutes=substitutes,
             source_config=source_config.result,
         )
 
+    async def get_config(self, **kwargs) -> ScrapliCfgResponse:
+        kwargs.pop("config_template", None)
+        return await super().get_config(**kwargs)
+
 
 class CustomAsyncScrapliCfgEOS(CustomScrapliCfg, AsyncScrapliCfgEOS):
     pass
@@ -119,4 +123,67 @@ class CustomAsyncScrapliCfgNXOS(CustomScrapliCfg, AsyncScrapliCfgNXOS):
 
 
 class CustomAsyncScrapliCfgJunos(CustomScrapliCfg, AsyncScrapliCfgJunos):
-    pass
+    is_set_config = False
+
+    async def get_config(self, config_template: str, source: str = "running") -> ScrapliCfgResponse:
+        response = self._pre_get_config(source=source)
+
+        command = "show configuration"
+        if re.findall(r"^set\s+", config_template, flags=re.I | re.M):
+            self.is_set_config = True
+            command += " | display set"
+
+        if self._in_configuration_session is True:
+            config_result = await self.conn.send_config(config=f"run {command}")
+        else:
+            config_result = await self.conn.send_command(command=command)
+
+        return self._post_get_config(
+            response=response,
+            source=source,
+            scrapli_responses=[config_result],
+            result=config_result.result,
+        )
+
+    async def load_config(self, config: str, replace: bool = False, **kwargs: Any) -> ScrapliCfgResponse:
+        """
+        Load configuration to a device
+
+        Supported kwargs:
+            set: bool indicating config is a "set" style config (ignored if replace is True)
+
+        Args:
+            config: string of the configuration to load
+            replace: replace the configuration or not, if false configuration will be loaded as a
+                merge operation
+            kwargs: additional kwargs that the implementing classes may need for their platform,
+                see above for junos supported kwargs
+
+        Returns:
+            ScrapliCfgResponse: response object
+
+        Raises:
+            N/A
+
+        """
+        response = self._pre_load_config(config=config)
+
+        config = self._prepare_load_config(config=config, replace=replace)
+
+        config_result = await self.conn.send_config(config=config, privilege_level="root_shell")
+
+        if self.is_set_config is True:
+            load_config = f"load set {self.filesystem}{self.candidate_config_filename}"
+        else:
+            if self._replace is True:
+                load_config = f"load override {self.filesystem}{self.candidate_config_filename}"
+            else:
+                load_config = f"load merge {self.filesystem}{self.candidate_config_filename}"
+
+        load_result = await self.conn.send_config(config=load_config)
+        self._in_configuration_session = True
+
+        return self._post_load_config(
+            response=response,
+            scrapli_responses=[config_result, load_result],
+        )