#!/usr/bin/env python
"""Script to upgrade a Cisco ASA."""
from __future__ import print_function
from datetime import datetime
from getpass import getpass
from netmiko import ConnectHandler, FileTransfer


def asa_scp_handler(ssh_conn, cmd='ssh scopy enable', mode='enable'):
    """Enable/disable SCP on Cisco ASA."""
    if mode == 'disable':
        cmd = 'no ' + cmd
    return ssh_conn.send_config_set([cmd])


def main():
    """Script to upgrade a Cisco ASA."""
    try:
        ip_addr = raw_input("Enter ASA IP address: ")
    except NameError:
        ip_addr = input("Enter ASA IP address: ")
    my_pass = getpass()
    start_time = datetime.now()
    print(">>>> {}".format(start_time))

    net_device = {
        'device_type': 'cisco_asa',
        'ip': ip_addr,
        'username': 'admin',
        'password': my_pass,
        'secret': my_pass,
        'port': 22,
    }

    print("\nLogging in to ASA")
    ssh_conn = ConnectHandler(**net_device)
    print()

    # ADJUST TO TRANSFER IMAGE FILE
    dest_file_system = 'disk0:'
    source_file = 'test1.txt'
    dest_file = 'test1.txt'
    alt_dest_file = 'asa825-59-k8.bin'

    with FileTransfer(ssh_conn, source_file=source_file, dest_file=dest_file,
                      file_system=dest_file_system) as scp_transfer:

        if not scp_transfer.check_file_exists():
            if not scp_transfer.verify_space_available():
                raise ValueError("Insufficient space available on remote device")

            print("Enabling SCP")
            output = asa_scp_handler(ssh_conn, mode='enable')
            print(output)

            print("\nTransferring file\n")
            scp_transfer.transfer_file()

            print("Disabling SCP")
            output = asa_scp_handler(ssh_conn, mode='disable')
            print(output)

        print("\nVerifying file")
        if scp_transfer.verify_file():
            print("Source and destination MD5 matches")
        else:
            raise ValueError("MD5 failure between source and destination files")

    print("\nSending boot commands")
    full_file_name = "{}/{}".format(dest_file_system, alt_dest_file)
    boot_cmd = 'boot system {}'.format(full_file_name)
    output = ssh_conn.send_config_set([boot_cmd])
    print(output)

    print("\nVerifying state")
    output = ssh_conn.send_command('show boot')
    print(output)

    # UNCOMMENT TO PERFORM WR MEM AND RELOAD
    # print("\nWrite mem and reload")
    # output = ssh_conn.send_command_expect('write mem')
    # output += ssh_conn.send_command('reload')
    # output += ssh_conn.send_command('y')
    # print(output)

    print("\n>>>> {}".format(datetime.now() - start_time))
    print()


if __name__ == "__main__":
    main()
