#!/usr/bin/python
#
# This file is part of Ansible
#
# Ansible is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Ansible is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Ansible.  If not, see <http://www.gnu.org/licenses/>.
#
from __future__ import absolute_import, division, print_function


__metaclass__ = type


DOCUMENTATION = """
module: nxos_install_os
extends_documentation_fragment:
- cisco.nxos.nxos
short_description: Set boot options like boot, kickstart image and issu.
description:
- Install an operating system by setting the boot options like boot image and kickstart
  image and optionally select to install using ISSU (In Server Software Upgrade).
version_added: 1.0.0
notes:
- Tested against the following platforms and images - N9k 7.0(3)I4(6), 7.0(3)I5(3),
  7.0(3)I6(1), 7.0(3)I7(1), 7.0(3)F2(2), 7.0(3)F3(2) - N3k 6.0(2)A8(6), 6.0(2)A8(8),
  7.0(3)I6(1), 7.0(3)I7(1) - N7k 7.3(0)D1(1), 8.0(1), 8.1(1), 8.2(1)
- Tested against Cisco MDS NX-OS 9.2(1)
- This module requires both the ANSIBLE_PERSISTENT_CONNECT_TIMEOUT and ANSIBLE_PERSISTENT_COMMAND_TIMEOUT
  timers to be set to 600 seconds or higher. The module will exit if the timers are
  not set properly.
- When using connection local, ANSIBLE_PERSISTENT_CONNECT_TIMEOUT and ANSIBLE_PERSISTENT_COMMAND_TIMEOUT
  can only be set using ENV variables or the ansible.cfg file.
- Do not include full file paths, just the name of the file(s) stored on the top level
  flash directory.
- This module attempts to install the software immediately, which may trigger a reboot.
- In check mode, the module will indicate if an upgrade is needed and whether or not
  the upgrade is disruptive or non-disruptive(ISSU).
author:
- Jason Edelman (@jedelman8)
- Gabriele Gerbibo (@GGabriele)
options:
  system_image_file:
    description:
    - Name of the system (or combined) image file on flash.
    required: true
    type: str
  kickstart_image_file:
    description:
    - Name of the kickstart image file on flash. (Not required on all Nexus platforms)
    type: str
  issu:
    description:
    - Upgrade using In Service Software Upgrade (ISSU). (Supported on N5k, N7k, N9k
      platforms)
    - Selecting 'required' or 'yes' means that upgrades will only proceed if the switch
      is capable of ISSU.
    - Selecting 'desired' means that upgrades will use ISSU if possible but will fall
      back to disruptive upgrade if needed.
    - Selecting 'no' means do not use ISSU. Forced disruptive.
    choices:
    - "required"
    - "desired"
    - "yes"
    - "no"
    default: "no"
    type: str
"""

EXAMPLES = """
- name: Install OS on N9k
  check_mode: false
  cisco.nxos.nxos_install_os:
    system_image_file: nxos.7.0.3.I6.1.bin
    issu: desired

- name: Wait for device to come back up with new image
  wait_for:
    port: 22
    state: started
    timeout: 500
    delay: 60
    host: '{{ inventory_hostname }}'

- name: Check installed OS for newly installed version
  nxos_command:
    commands: [show version | json]
  register: output
- assert:
    that:
    - output['stdout'][0]['kickstart_ver_str'] == '7.0(3)I6(1)'
"""

RETURN = """
install_state:
    description: Boot and install information.
    returned: always
    type: dict
    sample: {
    "install_state": [
        "Compatibility check is done:",
        "Module  bootable          Impact  Install-type  Reason",
        "------  --------  --------------  ------------  ------",
        "     1       yes  non-disruptive         reset  ",
        "Images will be upgraded according to following table:",
        "Module       Image                  Running-Version(pri:alt)           New-Version  Upg-Required",
        "------  ----------  ----------------------------------------  --------------------  ------------",
        "     1        nxos                               7.0(3)I6(1)           7.0(3)I7(1)           yes",
        "     1        bios                        v4.4.0(07/12/2017)    v4.4.0(07/12/2017)            no"
    ],
    }
"""


import re

from time import sleep

from ansible.module_utils.basic import AnsibleModule

from ansible_collections.cisco.nxos.plugins.module_utils.network.nxos.nxos import (
    load_config,
    run_commands,
)


# Output options are 'text' or 'json'
def execute_show_command(module, command, output="text"):
    cmds = [{"command": command, "output": output}]

    return run_commands(module, cmds)


def get_platform(module):
    """Determine platform type"""
    data = execute_show_command(module, "show inventory", "json")
    pid = data[0]["TABLE_inv"]["ROW_inv"][0]["productid"]

    if re.search(r"N3K", pid):
        type = "N3K"
    elif re.search(r"N5K", pid):
        type = "N5K"
    elif re.search(r"N6K", pid):
        type = "N6K"
    elif re.search(r"N7K", pid):
        type = "N7K"
    elif re.search(r"N9K", pid):
        type = "N9K"
    else:
        type = "unknown"

    return type


def parse_show_install(data):
    """Helper method to parse the output of the 'show install all impact' or
        'install all' commands.

    Sample Output:

    Installer will perform impact only check. Please wait.

    Verifying image bootflash:/nxos.7.0.3.F2.2.bin for boot variable "nxos".
    [####################] 100% -- SUCCESS

    Verifying image type.
    [####################] 100% -- SUCCESS

    Preparing "bios" version info using image bootflash:/nxos.7.0.3.F2.2.bin.
    [####################] 100% -- SUCCESS

    Preparing "nxos" version info using image bootflash:/nxos.7.0.3.F2.2.bin.
    [####################] 100% -- SUCCESS

    Performing module support checks.
    [####################] 100% -- SUCCESS

    Notifying services about system upgrade.
    [####################] 100% -- SUCCESS



    Compatibility check is done:
    Module  bootable          Impact  Install-type  Reason
    ------  --------  --------------  ------------  ------
         8       yes      disruptive         reset  Incompatible image for ISSU
        21       yes      disruptive         reset  Incompatible image for ISSU


    Images will be upgraded according to following table:
    Module       Image  Running-Version(pri:alt)    New-Version   Upg-Required
    ------  ----------  ----------------------------------------  ------------
         8       lcn9k                7.0(3)F3(2)    7.0(3)F2(2)           yes
         8        bios                     v01.17         v01.17            no
        21       lcn9k                7.0(3)F3(2)    7.0(3)F2(2)           yes
        21        bios                     v01.70         v01.70            no
    """
    if len(data) > 0:
        data = massage_install_data(data)
    ud = {"raw": data}
    ud["processed"] = []
    ud["disruptive"] = False
    ud["upgrade_needed"] = False
    ud["error"] = False
    ud["invalid_command"] = False
    ud["install_in_progress"] = False
    ud["server_error"] = False
    ud["upgrade_succeeded"] = False
    ud["use_impact_data"] = False

    # Check for server errors
    if isinstance(data, int):
        if data == -1:
            ud["server_error"] = True
        elif data >= 500:
            ud["server_error"] = True
        elif data == -32603:
            ud["server_error"] = True
        elif data == 1:
            ud["server_error"] = True
        return ud
    else:
        ud["list_data"] = data.split("\n")

    for x in ud["list_data"]:
        # Check for errors and exit if found.
        if re.search(r"Pre-upgrade check failed", x):
            ud["error"] = True
            break
        if re.search(r"[I|i]nvalid command", x):
            ud["invalid_command"] = True
            ud["error"] = True
            break
        if re.search(r"No install all data found", x):
            ud["error"] = True
            break

        # Check for potentially transient conditions
        if re.search(r"Another install procedure may\s*be in progress", x):
            ud["install_in_progress"] = True
            break
        if re.search(r"Backend processing error", x):
            ud["server_error"] = True
            break
        if re.search(r"timed out", x):
            ud["server_error"] = True
            break
        if re.search(r"^(-1|5\d\d)$", x):
            ud["server_error"] = True
            break

        # Check for messages indicating a successful upgrade.
        if re.search(r"Finishing the upgrade", x):
            ud["upgrade_succeeded"] = True
            break
        if re.search(r"Install has been successful", x):
            ud["upgrade_succeeded"] = True
            break
        if re.search(r"Switching over onto standby", x):
            ud["upgrade_succeeded"] = True
            break

        # We get these messages when the upgrade is non-disruptive and
        # we loose connection with the switchover but far enough along that
        # we can be confident the upgrade succeeded.
        if re.search(r"timeout .*trying to send command: install", x):
            ud["upgrade_succeeded"] = True
            ud["use_impact_data"] = True
            break
        if re.search(r"[C|c]onnection failure: timed out", x):
            ud["upgrade_succeeded"] = True
            ud["use_impact_data"] = True
            break

        # Begin normal parsing.
        if re.search(r"----|Module|Images will|Compatibility", x):
            ud["processed"].append(x)
            continue
        # Check to see if upgrade will be disruptive or non-disruptive and
        # build dictionary of individual modules and their status.
        # Sample Line:
        #
        # Module  bootable      Impact  Install-type  Reason
        # ------  --------  ----------  ------------  ------
        #     8        yes  disruptive         reset  Incompatible image
        rd = r"(\d+)\s+(\S+)\s+(disruptive|non-disruptive)\s+(\S+)"
        mo = re.search(rd, x)
        if mo:
            ud["processed"].append(x)
            key = "m%s" % mo.group(1)
            field = "disruptive"
            if mo.group(3) == "non-disruptive":
                ud[key] = {field: False}
            else:
                ud[field] = True
                ud[key] = {field: True}
            field = "bootable"
            if mo.group(2) == "yes":
                ud[key].update({field: True})
            else:
                ud[key].update({field: False})
            continue

        # Check to see if switch needs an upgrade and build a dictionary
        # of individual modules and their individual upgrade status.
        # Sample Line:
        #
        # Module  Image  Running-Version(pri:alt)    New-Version  Upg-Required
        # ------  -----  ----------------------------------------  ------------
        # 8       lcn9k                7.0(3)F3(2)    7.0(3)F2(2)           yes
        mo = re.search(r"(\d+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(yes|no)", x)
        if mo:
            ud["processed"].append(x)
            key = "m%s_%s" % (mo.group(1), mo.group(2))
            field = "upgrade_needed"
            if mo.group(5) == "yes":
                ud[field] = True
                ud[key] = {field: True}
            else:
                ud[key] = {field: False}
            continue

    return ud


def massage_install_data(data):
    # Transport cli returns a list containing one result item.
    # Transport nxapi returns a list containing two items.  The second item
    # contains the data we are interested in.
    default_error_msg = "No install all data found"
    if len(data) == 1:
        result_data = data[0]
    elif len(data) == 2:
        result_data = data[1]
    else:
        result_data = default_error_msg

    # Further processing may be needed for result_data
    if len(data) == 2 and isinstance(data[1], dict):
        if "clierror" in data[1].keys():
            result_data = data[1]["clierror"]
        elif "code" in data[1].keys() and data[1]["code"] == "500":
            # We encountered a backend processing error for nxapi
            result_data = data[1]["msg"]
        else:
            result_data = default_error_msg
    return result_data


def build_install_cmd_set(issu, image, kick, type, force=True):
    commands = ["terminal dont-ask"]

    # Different NX-OS platforms behave differently for
    # disruptive and non-disruptive upgrade paths.
    #
    # 1) Combined kickstart/system image:
    #    * Use option 'non-disruptive' for issu.
    #    * Omit option 'non-disruptive' for disruptive upgrades.
    # 2) Separate kickstart + system images.
    #    * Omit hidden 'force' option for issu.
    #    * Use hidden 'force' option for disruptive upgrades.
    #    * Note: Not supported on all platforms
    if re.search(r"required|desired|yes", issu):
        if kick is None:
            issu_cmd = "non-disruptive"
        else:
            issu_cmd = ""
    else:
        if kick is None:
            issu_cmd = ""
        else:
            issu_cmd = "force" if force else ""

    if type == "impact":
        rootcmd = "show install all impact"
        # The force option is not available for the impact command.
        if kick:
            issu_cmd = ""
    else:
        rootcmd = "install all"
    if kick is None:
        commands.append("%s nxos %s %s" % (rootcmd, image, issu_cmd))
    else:
        commands.append("%s %s system %s kickstart %s" % (rootcmd, issu_cmd, image, kick))

    return commands


def parse_show_version(data):
    version_data = {"raw": data[0].split("\n")}
    version_data["version"] = ""
    version_data["error"] = False
    for x in version_data["raw"]:
        mo = re.search(r"(kickstart|system|NXOS):\s+version\s+(\S+)", x)
        if mo:
            version_data["version"] = mo.group(2)
            continue

    if version_data["version"] == "":
        version_data["error"] = True

    return version_data


def check_mode_legacy(module, issu, image, kick=None):
    """Some platforms/images/transports don't support the 'install all impact'
    command so we need to use a different method."""
    current = execute_show_command(module, "show version", "json")[0]
    # Call parse_show_data on empty string to create the default upgrade
    # data structure dictionary
    data = parse_show_install("")
    upgrade_msg = "No upgrade required"

    # Process System Image
    data["error"] = False
    tsver = "show version image bootflash:%s" % image
    data["upgrade_cmd"] = [tsver]
    target_image = parse_show_version(execute_show_command(module, tsver))
    if target_image["error"]:
        data["error"] = True
        data["raw"] = target_image["raw"]
    if current["kickstart_ver_str"] != target_image["version"] and not data["error"]:
        data["upgrade_needed"] = True
        data["disruptive"] = True
        upgrade_msg = "Switch upgraded: system: %s" % tsver

    # Process Kickstart Image
    if kick is not None and not data["error"]:
        tkver = "show version image bootflash:%s" % kick
        data["upgrade_cmd"].append(tsver)
        target_kick = parse_show_version(execute_show_command(module, tkver))
        if target_kick["error"]:
            data["error"] = True
            data["raw"] = target_kick["raw"]
        if current["kickstart_ver_str"] != target_kick["version"] and not data["error"]:
            data["upgrade_needed"] = True
            data["disruptive"] = True
            upgrade_msg = upgrade_msg + " kickstart: %s" % tkver

    data["list_data"] = data["raw"]
    data["processed"] = upgrade_msg
    return data


def check_mode_nextgen(module, issu, image, kick=None):
    """Use the 'install all impact' command for check_mode"""
    opts = {"ignore_timeout": True}
    commands = build_install_cmd_set(issu, image, kick, "impact")
    data = parse_show_install(load_config(module, commands, True, opts))
    # If an error is encountered when issu is 'desired' then try again
    # but set issu to 'no'
    if data["error"] and issu == "desired":
        issu = "no"
        commands = build_install_cmd_set(issu, image, kick, "impact")
        # The system may be busy from the previous call to check_mode so loop
        # until it's done.
        data = check_install_in_progress(module, commands, opts)
    if data["server_error"]:
        data["error"] = True
    data["upgrade_cmd"] = commands
    return data


def check_install_in_progress(module, commands, opts):
    for attempt in range(20):
        data = parse_show_install(load_config(module, commands, True, opts))
        if data["install_in_progress"]:
            sleep(1)
            continue
        break
    return data


def check_mode(module, issu, image, kick=None):
    """Check switch upgrade impact using 'show install all impact' command"""
    data = check_mode_nextgen(module, issu, image, kick)
    if data["server_error"]:
        # We encountered an unrecoverable error in the attempt to get upgrade
        # impact data from the 'show install all impact' command.
        # Fallback to legacy method.
        data = check_mode_legacy(module, issu, image, kick)
    if data["invalid_command"]:
        # If we are upgrading from a device running a separate kickstart and
        # system image the impact command will fail.
        # Fallback to legacy method.
        data = check_mode_legacy(module, issu, image, kick)
    return data


def do_install_all(module, issu, image, kick=None):
    """Perform the switch upgrade using the 'install all' command"""
    impact_data = check_mode(module, issu, image, kick)
    if module.check_mode:
        # Check mode set in the playbook so just return the impact data.
        msg = "*** SWITCH WAS NOT UPGRADED: IMPACT DATA ONLY ***"
        impact_data["processed"].append(msg)
        return impact_data
    if impact_data["error"]:
        # Check mode discovered an error so return with this info.
        return impact_data
    elif not impact_data["upgrade_needed"]:
        # The switch is already upgraded.  Nothing more to do.
        return impact_data
    else:
        # If we get here, check_mode returned no errors and the switch
        # needs to be upgraded.
        if impact_data["disruptive"]:
            # Check mode indicated that ISSU is not possible so issue the
            # upgrade command without the non-disruptive flag unless the
            # playbook specified issu: yes/required.
            if issu == "yes":
                msg = "ISSU/ISSD requested but impact data indicates ISSU/ISSD is not possible"
                module.fail_json(msg=msg, raw_data=impact_data["list_data"])
            else:
                issu = "no"

        commands = build_install_cmd_set(issu, image, kick, "install")
        opts = {"ignore_timeout": True}
        # The system may be busy from the call to check_mode so loop until
        # it's done.
        upgrade = check_install_in_progress(module, commands, opts)
        if upgrade["invalid_command"] and "force" in commands[1]:
            # Not all platforms support the 'force' keyword.  Check for this
            # condition and re-try without the 'force' keyword if needed.
            commands = build_install_cmd_set(issu, image, kick, "install", False)
            upgrade = check_install_in_progress(module, commands, opts)
        upgrade["upgrade_cmd"] = commands

        # Special case:  If we encounter a server error at this stage
        # it means the command was sent and the upgrade was started but
        # we will need to use the impact data instead of the current install
        # data.
        if upgrade["server_error"]:
            upgrade["upgrade_succeeded"] = True
            upgrade["use_impact_data"] = True

        if upgrade["use_impact_data"]:
            if upgrade["upgrade_succeeded"]:
                upgrade = impact_data
                upgrade["upgrade_succeeded"] = True
            else:
                upgrade = impact_data
                upgrade["upgrade_succeeded"] = False

        if not upgrade["upgrade_succeeded"]:
            upgrade["error"] = True
    return upgrade


def main():
    argument_spec = dict(
        system_image_file=dict(required=True),
        kickstart_image_file=dict(required=False),
        issu=dict(choices=["required", "desired", "no", "yes"], default="no"),
    )

    module = AnsibleModule(argument_spec=argument_spec, supports_check_mode=True)

    warnings = list()

    # Get system_image_file(sif), kickstart_image_file(kif) and
    # issu settings from module params.
    sif = module.params["system_image_file"]
    kif = module.params["kickstart_image_file"]
    issu = module.params["issu"]

    if re.search(r"(yes|required)", issu):
        issu = "yes"

    if kif == "null" or kif == "":
        kif = None

    install_result = do_install_all(module, issu, sif, kick=kif)
    if install_result["error"]:
        cmd = install_result["upgrade_cmd"]
        msg = "Failed to upgrade device using command: %s" % cmd
        module.fail_json(msg=msg, raw_data=install_result["list_data"])

    state = install_result["processed"]
    changed = install_result["upgrade_needed"]
    module.exit_json(changed=changed, install_state=state, warnings=warnings)


if __name__ == "__main__":
    main()