Refactoring code to adhere to persistence connection. (#44398)

* Refactoring code to adhere to persistence connection.

* Update cnos_rollback.py

* Review comment of Qalthos incorporated

* Updating license for the refactored method

* Update cnos_rollback.py

* Removing the BSD License as suggested by Legal

2 files changed, 108 insertions, 281 deletions

diff --git a/lib/ansible/module_utils/network/cnos/cnos.py b/lib/ansible/module_utils/network/cnos/cnos.py
index 921921c255..9b38489ba4 100644
--- a/lib/ansible/module_utils/network/cnos/cnos.py
+++ b/lib/ansible/module_utils/network/cnos/cnos.py
@@ -1327,204 +1327,6 @@ def interfaceLevel2Config(module, cmd, prompt, answer):
 # EOM
 
 
-def doStartUpConfigRollback(
-        protocol, timeout, confServerIp, confPath, confServerUser,
-        confServerPwd, obj):
-    # server = "10.241.105.214"
-    server = confServerIp
-
-    # username = "pbhosale"
-    username = confServerUser
-
-    # password = "Lab4man1"
-    password = confServerPwd
-
-    if((confPath is None) or (confPath is "")):
-        confPath = "cnos_config"
-
-    retVal = ""
-
-    # config backup command happens here
-    if(protocol == "ftp"):
-        command = "cp " + protocol + " " + protocol + "://" + username + \
-            "@" + server + "/" + confPath + " startup-config vrf management\n"
-        # debugOutput(command)
-        retVal = retVal + waitForDeviceResponse(command, "Password:", 3, obj)
-        # debugOutput(command)
-        retVal = retVal + waitForDeviceResponse(command, "[n]", timeout, obj)
-        command = "y\n"
-        retVal = retVal + waitForDeviceResponse(password, "#", timeout, obj)
-    elif(protocol == "tftp"):
-        command = "cp " + protocol + " " + protocol + "://" + \
-            server + "/" + confPath + " startup-config vrf management\n"
-        # debugOutput(command)
-        retVal = retVal + waitForDeviceResponse(command, "[n]", timeout, obj)
-        command = "y\n"
-        retVal = retVal + waitForDeviceResponse(command, "#", timeout, obj)
-    else:
-        return "Error-110"
-
-    return retVal
-# EOM
-
-# Utility Method to restore the start up config
-# This method supports only SCP or SFTP
-# Tuning of timeout parameter is pending
-
-
-def doSecureStartUpConfigRollback(
-        protocol, timeout, confServerIp, confPath, confServerUser,
-        confServerPwd, obj):
-    # server = "10.241.105.214"
-    server = confServerIp
-
-    # username = "pbhosale"
-    username = confServerUser
-
-    # password = "Lab4man1"
-    password = confServerPwd
-
-    if((confPath is None) or (confPath is "")):
-        confPath = "cnos_config"
-
-    retVal = ""
-
-    # config backup command happens here
-
-    # cp sftp sftp://root@10.241.106.118/cnos_config/running_config.conf
-    # startup-config vrf management
-    command = "cp " + protocol + " " + protocol + "://" + username + \
-        "@" + server + "/" + confPath + " startup-config vrf management \n"
-
-    # debugOutput(command)
-    response = waitForDeviceResponse(command, "(yes/no)", 3, obj)
-    if(response.lower().find('error-101')):
-        command = password + "\n"
-        retVal = retVal + waitForDeviceResponse(command, "[n]", timeout, obj)
-        command = "y\n"
-        retVal = retVal + waitForDeviceResponse(command, "#", timeout, obj)
-        return retVal
-    retVal = retVal + response
-    if(protocol == "scp"):
-        command = "yes \n"
-        # debugOutput(command)
-        retVal = retVal + waitForDeviceResponse(command, "timeout:", 3, obj)
-        command = "0\n"
-        # debugOutput(command)
-        retVal = retVal + waitForDeviceResponse(command, "Password:", 3, obj)
-    elif(protocol == "sftp"):
-        command = "yes \n"
-        # debugOutput(command)
-        retVal = retVal + waitForDeviceResponse(command, "Password:", 3, obj)
-    else:
-        return "Error-110"
-
-    # Password entry happens here
-    # debugOutput(command)
-    command = password + "\n"
-    retVal = retVal + waitForDeviceResponse(command, "[n]", timeout, obj)
-    command = "y\n"
-    retVal = retVal + waitForDeviceResponse(command, "#", timeout, obj)
-
-    return retVal
-# EOM
-
-
-def doRunningConfigRollback(
-        protocol, timeout, confServerIp, confPath, confServerUser,
-        confServerPwd, obj):
-    # server = "10.241.105.214"
-    server = confServerIp
-
-    # username = "pbhosale"
-    username = confServerUser
-
-    # password = "Lab4man1"
-    password = confServerPwd
-
-    if((confPath is None) or (confPath is "")):
-        confPath = "cnos_config"
-
-    retVal = ""
-
-    # config backup command happens here
-    if(protocol == "ftp"):
-        command = "cp " + protocol + " " + protocol + "://" + username + \
-            "@" + server + "/" + confPath + " running-config vrf management\n"
-        # debugOutput(command)
-        retVal = retVal + waitForDeviceResponse(command, "Password:", 3, obj)
-        # debugOutput(command)
-        retVal = retVal + waitForDeviceResponse(password, "#", timeout, obj)
-    elif(protocol == "tftp"):
-        command = "cp " + protocol + " " + protocol + "://" + \
-            server + "/" + confPath + " running-config vrf management\n"
-        # debugOutput(command)
-        retVal = retVal + waitForDeviceResponse(command, "#", timeout, obj)
-    else:
-        return "Error-110"
-
-    return retVal
-# EOM
-
-# Utility Method to restore the running config
-# This method supports only SCP or SFTP
-# Tuning of timeout parameter is pending
-
-
-def doSecureRunningConfigRollback(
-        protocol, timeout, confServerIp, confPath, confServerUser,
-        confServerPwd, obj):
-    # server = "10.241.105.214"
-    server = confServerIp
-
-    # username = "pbhosale"
-    username = confServerUser
-
-    # password = "Lab4man1"
-    password = confServerPwd
-
-    if((confPath is None) or (confPath is "")):
-        confPath = "cnos_config"
-
-    retVal = ""
-
-    # config backup command happens here
-
-    # cp sftp sftp://root@10.241.106.118/cnos_config/running_config.conf
-    # running-config vrf management
-    command = "cp " + protocol + " " + protocol + "://" + username + \
-        "@" + server + "/" + confPath + " running-config vrf management \n"
-
-    # debugOutput(command)
-    response = waitForDeviceResponse(command, "(yes/no)", 3, obj)
-    if(response.lower().find('error-101')):
-        command = password + "\n"
-        retVal = retVal + waitForDeviceResponse(command, "#", timeout, obj)
-        return retVal
-    retVal = retVal + response
-
-    if(protocol == "scp"):
-        command = "yes \n"
-        # debugOutput(command)
-        retVal = retVal + waitForDeviceResponse(command, "timeout:", 3, obj)
-        command = "0\n"
-        # debugOutput(command)
-        retVal = retVal + waitForDeviceResponse(command, "Password:", 3, obj)
-    elif(protocol == "sftp"):
-        command = "yes \n"
-        # debugOutput(command)
-        retVal = retVal + waitForDeviceResponse(command, "Password:", 3, obj)
-    else:
-        return "Error-110"
-
-    # Password entry happens here
-    # debugOutput(command)
-    command = password + "\n"
-    retVal = retVal + waitForDeviceResponse(command, "#", timeout, obj)
-
-    return retVal
-# EOM
-
 # Utility Method to download an image from FTP/TFTP server to device.
 # This method supports only FTP or TFTP
 # Tuning of timeout parameter is pending
diff --git a/lib/ansible/modules/network/cnos/cnos_rollback.py b/lib/ansible/modules/network/cnos/cnos_rollback.py
index 27f691f939..34bd8b9309 100644
--- a/lib/ansible/modules/network/cnos/cnos_rollback.py
+++ b/lib/ansible/modules/network/cnos/cnos_rollback.py
@@ -34,63 +34,74 @@ DOCUMENTATION = '''
 ---
 module: cnos_rollback
 author: "Anil Kumar Muraleedharan (@amuraleedhar)"
-short_description: Roll back the running or startup configuration from a remote server on devices running Lenovo CNOS
+short_description: Roll back the running or startup configuration from a remote
+ server on devices running Lenovo CNOS
 description:
-    - This module allows you to work with switch configurations. It provides a way to roll back configurations
-     of a switch from a remote server. This is achieved by using startup or running configurations of the target
-     device that were previously backed up to a remote server using FTP, SFTP, TFTP, or SCP.
-     The first step is to create a directory from where the remote server can be reached. The next step is to
-     provide the full file path of the backup configuration's location. Authentication details required by the
+    - This module allows you to work with switch configurations. It provides a
+     way to roll back configurations of a switch from a remote server. This is
+     achieved by using startup or running configurations of the target device
+     that were previously backed up to a remote server using FTP, SFTP, TFTP,
+     or SCP. The first step is to create a directory from where the remote
+     server can be reached. The next step is to provide the full file path of
+     he backup configuration's location. Authentication details required by the
      remote server must be provided as well.
-     By default, this method overwrites the switch's configuration file with the newly downloaded file.
-     This module uses SSH to manage network device configuration.
-     The results of the operation will be placed in a directory named 'results'
-     that must be created by the user in their local directory to where the playbook is run.
-     For more information about this module from Lenovo and customizing it usage for your
-     use cases, please visit U(http://systemx.lenovofiles.com/help/index.jsp?topic=%2Fcom.lenovo.switchmgt.ansible.doc%2Fcnos_rollback.html)
+     By default, this method overwrites the switch's configuration file with
+     the newly downloaded file. This module uses SSH to manage network device
+     configuration. The results of the operation will be placed in a directory
+     named 'results' that must be created by the user in their local directory
+     to where the playbook is run. For more information about this module from
+     Lenovo and customizing it usage for your use cases, please visit
+     U(http://systemx.lenovofiles.com/help/index.jsp?topic=%2Fcom.lenovo.switchmgt.ansible.doc%2Fcnos_rollback.html)
 version_added: "2.3"
 extends_documentation_fragment: cnos
 options:
    configType:
         description:
-            - This refers to the type of configuration which will be used for the rolling back process.
-             The choices are the running or startup configurations. There is no default value, so it will result
+            - This refers to the type of configuration which will be used for
+             the rolling back process. The choices are the running or startup
+             configurations. There is no default value, so it will result
              in an error if the input is incorrect.
         required: Yes
         default: Null
         choices: [running-config, startup-config]
    protocol:
         description:
-            - This refers to the protocol used by the network device to interact with the remote server from where to
-             download the backup configuration. The choices are FTP, SFTP, TFTP, or SCP. Any other protocols will result
-             in error. If this parameter is not specified, there is no default value to be used.
+            - This refers to the protocol used by the network device to
+              interact with the remote server from where to download the backup
+              configuration. The choices are FTP, SFTP, TFTP, or SCP. Any other
+              protocols will result in error. If this parameter is not
+              specified, there is no default value to be used.
         required: Yes
         default: Null
         choices: [SFTP, SCP, FTP, TFTP]
    rcserverip:
         description:
-            - This specifies the IP Address of the remote server from where the backup configuration will be downloaded.
+            - This specifies the IP Address of the remote server from where the
+             backup configuration will be downloaded.
         required: Yes
         default: Null
    rcpath:
         description:
-            - This specifies the full file path of the configuration file located on the remote server. In case the relative
-             path is used as the variable value, the root folder for the user of the server needs to be specified.
+            - This specifies the full file path of the configuration file
+             located on the remote server. In case the relative path is used as
+             the variable value, the root folder for the user of the server
+             needs to be specified.
         required: Yes
         default: Null
    serverusername:
         description:
-            - Specify the username for the server relating to the protocol used.
+            - Specify username for the server relating to the protocol used.
         required: Yes
         default: Null
    serverpassword:
         description:
-            - Specify the password for the server relating to the protocol used.
+            - Specify password for the server relating to the protocol used.
         required: Yes
         default: Null
 '''
 EXAMPLES = '''
-Tasks : The following are examples of using the module cnos_rollback. These are written in the main.yml file of the tasks directory.
+Tasks : The following are examples of using the module cnos_rollback.
+ These are written in the main.yml file of the tasks directory.
 ---
 
 - name: Test Rollback of config - Running config
@@ -163,17 +174,13 @@ msg:
 '''
 
 import sys
-try:
-    import paramiko
-    HAS_PARAMIKO = True
-except ImportError:
-    HAS_PARAMIKO = False
 import time
 import socket
 import array
 import json
 import time
 import re
+import os
 try:
     from ansible.module_utils.network.cnos import cnos
     HAS_LIB = True
@@ -183,6 +190,70 @@ from ansible.module_utils.basic import AnsibleModule
 from collections import defaultdict
 
 
+# Utility Method to rollback the running config or start up copnfig
+# This method supports only SCP or SFTP or FTP or TFTP
+def doConfigRollBack(module, prompt, answer):
+    host = module.params['host']
+    server = module.params['serverip']
+    username = module.params['serverusername']
+    password = module.params['serverpassword']
+    protocol = module.params['protocol'].lower()
+    rcPath = module.params['rcpath']
+    configType = module.params['configType']
+    confPath = rcPath
+    retVal = ''
+
+    command = "copy " + protocol + " " + protocol + "://"
+    command = command + username + "@" + server + "/" + confPath
+    command = command + " " + configType + " vrf management\n"
+    cnos.debugOutput(command + "\n")
+    # cnos.checkForFirstTimeAccess(module, command, 'yes/no', 'yes')
+    cmd = []
+    if(protocol == "scp"):
+        scp_cmd1 = [{'command': command, 'prompt': 'timeout:', 'answer': '0'}]
+        scp_cmd2 = [{'command': '\n', 'prompt': 'Password:',
+                     'answer': password}]
+        cmd.extend(scp_cmd1)
+        cmd.extend(scp_cmd2)
+        if(configType == 'startup-config'):
+            scp_cmd3 = [{'command': 'y', 'prompt': None, 'answer': None}]
+            cmd.extend(scp_cmd3)
+        retVal = retVal + str(cnos.run_cnos_commands(module, cmd))
+    elif(protocol == "sftp"):
+        sftp_cmd = [{'command': command, 'prompt': 'Password:',
+                     'answer': password}]
+        cmd.extend(sftp_cmd)
+        # cnos.debugOutput(configType + "\n")
+        if(configType == 'startup-config'):
+            sftp_cmd2 = [{'command': 'y', 'prompt': None, 'answer': None}]
+            cmd.extend(sftp_cmd2)
+        retVal = retVal + str(cnos.run_cnos_commands(module, cmd))
+    elif(protocol == "ftp"):
+        ftp_cmd = [{'command': command, 'prompt': 'Password:',
+                    'answer': password}]
+        cmd.extend(ftp_cmd)
+        if(configType == 'startup-config'):
+            ftp_cmd2 = [{'command': 'y', 'prompt': None, 'answer': None}]
+            cmd.extend(ftp_cmd2)
+        retVal = retVal + str(cnos.run_cnos_commands(module, cmd))
+    elif(protocol == "tftp"):
+        command = "copy " + protocol + " " + protocol
+        command = command + "://" + server + "/" + confPath
+        command = command + " " + configType + " vrf management\n"
+        cnos.debugOutput(command)
+        tftp_cmd = [{'command': command, 'prompt': None, 'answer': None}]
+        cmd.extend(tftp_cmd)
+        if(configType == 'startup-config'):
+            tftp_cmd2 = [{'command': 'y', 'prompt': None, 'answer': None}]
+            cmd.extend(tftp_cmd2)
+        retVal = retVal + str(cnos.run_cnos_commands(module, cmd))
+    else:
+        return "Error-110"
+
+    return retVal
+# EOM
+
+
 def main():
     module = AnsibleModule(
         argument_spec=dict(
@@ -200,67 +271,21 @@ def main():
             serverpassword=dict(required=False, no_log=True),),
         supports_check_mode=False)
 
-    username = module.params['username']
-    password = module.params['password']
-    enablePassword = module.params['enablePassword']
     outputfile = module.params['outputfile']
-    host = module.params['host']
-    deviceType = module.params['deviceType']
-    configType = module.params['configType']
     protocol = module.params['protocol'].lower()
-    rcserverip = module.params['serverip']
-    rcpath = module.params['rcpath']
-    serveruser = module.params['serverusername']
-    serverpwd = module.params['serverpassword']
-    output = ""
-    timeout = 90
-    tftptimeout = 450
-    if not HAS_PARAMIKO:
-        module.fail_json(msg='paramiko is required for this module')
-
-    # Create instance of SSHClient object
-    remote_conn_pre = paramiko.SSHClient()
-
-    # Automatically add untrusted hosts (make sure okay for security policy in your environment)
-    remote_conn_pre.set_missing_host_key_policy(paramiko.AutoAddPolicy())
-
-    # initiate SSH connection with the switch
-    remote_conn_pre.connect(host, username=username, password=password)
-    time.sleep(2)
-
-    # Use invoke_shell to establish an 'interactive session'
-    remote_conn = remote_conn_pre.invoke_shell()
-    time.sleep(2)
-
-    # Enable and then send command
-    output = output + cnos.waitForDeviceResponse("\n", ">", 2, remote_conn)
-
-    output = output + cnos.enterEnableModeForDevice(enablePassword, 3, remote_conn)
-
-    # Make terminal length = 0
-    output = output + cnos.waitForDeviceResponse("terminal length 0\n", "#", 2, remote_conn)
-
-    # Invoke method for Config transfer from server
-    if(configType == 'running-config'):
-        if(protocol == "tftp" or protocol == "ftp"):
-            transfer_status = cnos.doRunningConfigRollback(protocol, tftptimeout, rcserverip, rcpath, serveruser, serverpwd, remote_conn)
-        elif(protocol == "sftp" or protocol == "scp"):
-            transfer_status = cnos.doSecureRunningConfigRollback(protocol, timeout, rcserverip, rcpath, serveruser, serverpwd, remote_conn)
-        else:
-            transfer_status = "Invalid Protocol option"
-    elif(configType == 'startup-config'):
-        if(protocol == "tftp" or protocol == "ftp"):
-            transfer_status = cnos.doStartUpConfigRollback(protocol, tftptimeout, rcserverip, rcpath, serveruser, serverpwd, remote_conn)
-        elif(protocol == "sftp" or protocol == "scp"):
-            transfer_status = cnos.doSecureStartUpConfigRollback(protocol, timeout, rcserverip, rcpath, serveruser, serverpwd, remote_conn)
-        else:
-            transfer_status = "Invalid Protocol option"
+    output = ''
+    if protocol in ('tftp', 'ftp', 'sftp', 'scp'):
+        transfer_status = doConfigRollBack(module, None, None)
     else:
-        transfer_status = "Invalid configType Option"
-
+        transfer_status = 'Invalid Protocol option'
     output = output + "\n Config Transfer status \n" + transfer_status
 
     # Save it into the file
+    if '/' in outputfile:
+        path = outputfile.rsplit('/', 1)
+        # cnos.debugOutput(path[0])
+        if not os.path.exists(path[0]):
+            os.makedirs(path[0])
     file = open(outputfile, "a")
     file.write(output)
     file.close()