From cc27e3f1b4d6ac647f3d0a4b974a589885207858 Mon Sep 17 00:00:00 2001
From: Omar Al-Ghussein <omaral-ghussein@Omars-MacBook-Air.local>
Date: Tue, 30 Jul 2019 14:08:34 +0300
Subject: [PATCH] fix juniper CLI mode

---
 README.rst                                    |  4 ++--
 aionet/vendors/devices/base_junos.py          | 14 +++++++------
 .../vendors/devices/juniper/juniper_junos.py  | 20 +++++++++++++++++--
 aionet/vendors/terminal_modes/juniper.py      |  9 +++++++++
 4 files changed, 37 insertions(+), 10 deletions(-)

diff --git a/README.rst b/README.rst
index c3e54c5..aac0107 100644
--- a/README.rst
+++ b/README.rst
@@ -59,12 +59,12 @@ Example of interacting with Cisco IOS devices:
         dev1 = { 'username' : 'user',
                  'password' : 'pass',
                  'device_type': 'cisco_ios',
-                 'host': 'ip address',
+                 'ip': 'ip address',
         }
         dev2 = { 'username' : 'user',
                  'password' : 'pass',
                  'device_type': 'cisco_ios',
-                 'host': 'ip address',
+                 'ip': 'ip address',
         }
         devices = [dev1, dev2]
         tasks = [task(dev) for dev in devices]
diff --git a/aionet/vendors/devices/base_junos.py b/aionet/vendors/devices/base_junos.py
index 0e74d6a..2f81e02 100644
--- a/aionet/vendors/devices/base_junos.py
+++ b/aionet/vendors/devices/base_junos.py
@@ -72,15 +72,17 @@ async def _set_base_prompt(self):
         # Strip off trailing terminator
         if "@" in prompt:
             prompt = prompt.split("@")[1]
-        self._base_prompt = prompt
+        self.device_prompt = prompt
         delimiters = map(re.escape, type(self)._delimiter_list)
         delimiters = r"|".join(delimiters)
-        base_prompt = re.escape(self._base_prompt[:12])
+        base_prompt = re.escape(self.device_prompt[:12])
         pattern = type(self)._pattern
-        self._base_pattern = pattern.format(delimiters=delimiters)
-        self._logger.debug("Base Prompt: %s" % self._base_prompt)
-        self._logger.debug("Base Pattern: %s" % self._base_pattern)
-        return self._base_prompt
+        base_pattern = pattern.format(delimiters=delimiters)
+        self._logger.debug("Base Prompt: %s" % self.device_prompt)
+        self._logger.debug("Base Pattern: %s" % base_pattern)
+        self.prompt_pattern = base_pattern
+        self._conn.set_base_pattern(base_pattern)
+        return self.device_prompt
 
     async def send_config_set(
             self,
diff --git a/aionet/vendors/devices/juniper/juniper_junos.py b/aionet/vendors/devices/juniper/juniper_junos.py
index d42beb7..0408bcb 100644
--- a/aionet/vendors/devices/juniper/juniper_junos.py
+++ b/aionet/vendors/devices/juniper/juniper_junos.py
@@ -1,4 +1,4 @@
-from aionet.vendors.terminal_modes.juniper import CliMode
+from aionet.vendors.terminal_modes.juniper import CliMode,ConfigMode
 from aionet.vendors.devices.base_junos import BAseJunOSDevice
 
 
@@ -15,12 +15,28 @@ def __init__(self, *args, **kwargs):
             device=self
         )
 
+        self.config_mode = ConfigMode(
+            enter_command=type(self)._config_command,
+            check_string=type(self)._config_check,
+            exit_command='exit',
+            device=self,
+            parent=self.cli_mode
+
+        )
+
     _cli_check = ">"
     """Checking string for shell mode"""
 
     _cli_command = "cli"
     """Command for entering to cli mode"""
 
+    _config_check = "#"
+    """Checking string for shell mode"""
+
+    _config_command = "configure"
+    """Command for entering to cli mode"""
+
     async def _session_preparation(self):
-        await self.cli_mode()
         await super()._session_preparation()
+        await self.cli_mode()
+
diff --git a/aionet/vendors/terminal_modes/juniper.py b/aionet/vendors/terminal_modes/juniper.py
index 4749fc6..9ad98c4 100644
--- a/aionet/vendors/terminal_modes/juniper.py
+++ b/aionet/vendors/terminal_modes/juniper.py
@@ -6,6 +6,15 @@
 
 
 class ConfigMode(CiscoConfigMode):
+    async def check(self, force=False):
+        """Check if are in configuration mode. Return boolean"""
+        if self.device.current_terminal is not None and not force:
+            if self.device.current_terminal == self:
+                return True
+        await self.device.send_new_line()
+        output = await self.device.send_new_line()
+
+        return self._check_string in output
     pass