#!/usr/bin/python2 -s
# Copyright (C) 2016 Red Hat, Inc.
# This file is part of libsan.
#
# libsan 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.
#
# libsan 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 libsan.  If not, see <http://www.gnu.org/licenses/>.

"""
Tool to perform basic commands on SAN
"""
import libsan.sanmgmt as sanmgmt
import libsan.host.fcoe as fcoe
import libsan.host.linux as linux
import libsan.host.scsi as scsi
import libsan.host.net as net
import libsan.host.dt as dt
import libsan.host.fio as fio
import libsan.host.lio as lio
import libsan.misc.time as time
import sys
import argparse


OBJ_SANMGMT = None

def show_host_cfg():
    """
    Show SAN configuration of the host
    """
    global OBJ_SANMGMT
    if not OBJ_SANMGMT:
        #Initialize SAN obj based on all information configured on /etc/san_top.cfg
        #based on HW information
        OBJ_SANMGMT = sanmgmt.create_sanmgmt_for_mpath()
        if not OBJ_SANMGMT:
            print "FAIL: Could not create LibSAN object"
            sys.exit(1)

    nics = net.get_nics()
    macs = []
    if nics:
        for nic in nics:
            macs.append(net.get_mac_of_nic(nic))
    if macs:
        print "Host has the following MACs:"
        for mac in macs:
            sw_port = "Not configured on SAN top"
            connected_sw = OBJ_SANMGMT.get_sw_self(mac)
            if connected_sw:
                sw_port = connected_sw["port_of"][mac]
            print "\t%s - SW port: %s" % (mac, sw_port)

    iscsi_macs = OBJ_SANMGMT.macs()
    if iscsi_macs:
        print "Host has the following iSCSI MACs:"
        for mac in iscsi_macs:
            sw_port = "Not configured on SAN top"
            connected_sw = OBJ_SANMGMT.get_sw_self(mac)
            if connected_sw:
                sw_port = connected_sw["port_of"][mac]
            print "\t%s - SW port: %s" % (mac, sw_port)


    h_wwpns = OBJ_SANMGMT.h_wwpns()
    if h_wwpns:
        print "Host has the following h_wwpns configured:"
        for wwpn in h_wwpns:
            sw_port = "Not configured on SAN top"
            connected_sw = OBJ_SANMGMT.get_sw_self(wwpn)
            if connected_sw:
                sw_port = connected_sw["port_of"][wwpn]
            print "\t%s - SW port: %s" % (wwpn, sw_port)

    t_wwpns = OBJ_SANMGMT.t_wwpns()
    if t_wwpns:
        print "Host has the following t_wwpns configured:"
        for wwpn in t_wwpns:
            connected_sw = OBJ_SANMGMT.get_sw_self(wwpn)
            if connected_sw:
                sw_port = connected_sw["port_of"][wwpn]
                print "\t%s - SW port: %s" % (wwpn, sw_port)

def check_ports_ready():
    """
    Check if all ports need for test are up
    """
    global OBJ_SANMGMT
    if not OBJ_SANMGMT:
        #Initialize SAN obj based on all information configured on /etc/san_top.cfg
        #based on HW information
        OBJ_SANMGMT = sanmgmt.create_sanmgmt_for_mpath()
        if not OBJ_SANMGMT:
            print "FAIL: Could not create LibSAN object"
            sys.exit(1)
    return OBJ_SANMGMT.check_ports_ready()

def rcsn_enable(addr):
    """
    Enable RCSN
    """
    global OBJ_SANMGMT
    if not OBJ_SANMGMT:
        #Initialize SAN obj based on all information configured on /etc/san_top.cfg
        #based on HW information
        OBJ_SANMGMT = sanmgmt.create_sanmgmt_for_mpath()
        if not OBJ_SANMGMT:
            print "FAIL: Could not create LibSAN object"
            sys.exit(1)

    if not addr:
        print "FAIL: rcsn_enable() - requires addr as parameter"
        return False

    connected_sw = OBJ_SANMGMT.get_sw_self(addr)
    if not connected_sw:
        print "FAIL: %s is not managed by libsan" % addr
        return False
    return OBJ_SANMGMT.rcsn_trigger(action="ENABLE", addr=addr)

def rcsn_disable(addr):
    """
    Disable RCSN
    """
    global OBJ_SANMGMT
    if not OBJ_SANMGMT:
        #Initialize SAN obj based on all information configured on /etc/san_top.cfg
        #based on HW information
        OBJ_SANMGMT = sanmgmt.create_sanmgmt_for_mpath()
        if not OBJ_SANMGMT:
            print "FAIL: Could not create LibSAN object"
            sys.exit(1)

    if not addr:
        print "FAIL: rcsn_disable() - requires addr as parameter"
        return False

    connected_sw = OBJ_SANMGMT.get_sw_self(addr)
    if not connected_sw:
        print "FAIL: %s is not managed by libsan" % addr
        return False
    return OBJ_SANMGMT.rcsn_trigger(action="DISABLE", addr=addr)


def port_connect(addr):
    """
    Enable switch port
    """
    global OBJ_SANMGMT
    if not OBJ_SANMGMT:
        #Initialize SAN obj based on all information configured on /etc/san_top.cfg
        #based on HW information
        OBJ_SANMGMT = sanmgmt.create_sanmgmt_for_mpath()
        if not OBJ_SANMGMT:
            print "FAIL: Could not create LibSAN object"
            sys.exit(1)

    if not addr:
        print "FAIL: port_connect() - requires addr as parameter"
        return False

    connected_sw = OBJ_SANMGMT.get_sw_self(addr)
    if not connected_sw:
        print "FAIL: %s is not managed by libsan" % addr
        return False
    return OBJ_SANMGMT.link_trigger(action="UP", addr=addr)

def port_disconnect(addr):
    """
    Disable switch port
    """
    global OBJ_SANMGMT
    if not OBJ_SANMGMT:
        #Initialize SAN obj based on all information configured on /etc/san_top.cfg
        #based on HW information
        OBJ_SANMGMT = sanmgmt.create_sanmgmt_for_mpath()
        if not OBJ_SANMGMT:
            print "FAIL: Could not create LibSAN object"
            sys.exit(1)

    if not addr:
        print "FAIL: port_disconnect() - requires addr as parameter"
        return False

    connected_sw = OBJ_SANMGMT.get_sw_self(addr)
    if not connected_sw:
        print "FAIL: %s is not managed by libsan" % addr
        return False
    return OBJ_SANMGMT.link_trigger(action="DOWN", addr=addr)

def link_torture(addr, runtime, interval):
    """
    Bring switch port down and up for specific time
    """
    global OBJ_SANMGMT
    if not OBJ_SANMGMT:
        #Initialize SAN obj based on all information configured on /etc/san_top.cfg
        #based on HW information
        OBJ_SANMGMT = sanmgmt.create_sanmgmt_for_mpath()
        if not OBJ_SANMGMT:
            print "FAIL: Could not create LibSAN object"
            sys.exit(1)

    if not addr or not runtime:
        print "FAIL: link_torture() - requires addr, runtime as parameteres"
        return False

    if interval is None:
        interval = 30

    interval = time.time_2_sec(interval)
    if interval is None:
        print "FAIL: link_torture() - Could not convert interval to seconds"
        return False

    time_in_sec = time.time_2_sec(runtime)
    if not time_in_sec:
        return True

    current_time = time.get_time(in_seconds=True)
    if not current_time:
        print "FAIL: link_torture() - Could not get current time"
        return False

    end_time = current_time + time_in_sec

    error = False
    while current_time < end_time:
        if not port_disconnect(addr):
            error = True
        print "INFO: sleeping for %d seconds" % interval
        linux.sleep(interval)
        if not port_connect(addr):
            error = True
        if error:
            break
        print "INFO: sleeping for %d seconds" % interval
        linux.sleep(interval)
        current_time = time.get_time(in_seconds=True)

    if error:
        #Try to bring port up
        print "Trying to bring %s back up" % addr
        port_connect(addr)
        return False

    return True

def lun_create(size, lun_name=None):
    """
    Create a LUN on array
    """
    global OBJ_SANMGMT
    if not OBJ_SANMGMT:
        #Initialize SAN obj based on all information configured on /etc/san_top.cfg
        #based on HW information
        OBJ_SANMGMT = sanmgmt.create_sanmgmt_for_mpath()
        if not OBJ_SANMGMT:
            print "FAIL: Could not create LibSAN object"
            sys.exit(1)

    if not size:
        print "FAIL: lun_create() - requires size as parameter"
        return None

    new_lun_name = OBJ_SANMGMT.lun_create(size, lun_name=lun_name)
    if new_lun_name:
        print "INFO: created LUN %s" % new_lun_name
        return new_lun_name
    print "FAIL: Could not create LUN"
    return None

def lun_create_and_map(size, lun_name=None):
    """
    Create a LUN to array and map it to server
    """
    global OBJ_SANMGMT
    if not OBJ_SANMGMT:
        #Initialize SAN obj based on all information configured on /etc/san_top.cfg
        #based on HW information
        OBJ_SANMGMT = sanmgmt.create_sanmgmt_for_mpath()
        if not OBJ_SANMGMT:
            print "FAIL: lun_create_and_map() - Could not create LibSAN object"
            sys.exit(1)

    if not size:
        print "FAIL: lun_create_and_map() - requires size as parameter"
        return None

    new_lun_name = OBJ_SANMGMT.lun_create_and_map(size, lun_name=lun_name, rescan=True)
    if new_lun_name:
        print "INFO: created and mapped LUN %s" % new_lun_name
        return new_lun_name
    print "FAIL: Could not create LUN"
    return None

def lun_remove(lun_name):
    """
    Remove LUN from array
    """
    global OBJ_SANMGMT
    if not OBJ_SANMGMT:
        #Initialize SAN obj based on all information configured on /etc/san_top.cfg
        #based on HW information
        OBJ_SANMGMT = sanmgmt.create_sanmgmt_for_mpath()
        if not OBJ_SANMGMT:
            print "FAIL: Could not create LibSAN object"
            sys.exit(1)

    if not lun_name:
        print "FAIL: lun_remove() - requires lun_name as parameter"
        return False

    error = 0
    names = lun_name.split(",")
    for name in names:
        if not OBJ_SANMGMT.lun_remove(name):
            print "FAIL: Could not delete %s" % name
            error += 1
    if error > 0:
        return False
    return True

def lun_map(lun_name, t_addr=None, i_addr=None):
    global OBJ_SANMGMT
    if not OBJ_SANMGMT:
        #Initialize SAN obj based on all information configured on /etc/san_top.cfg
        #based on HW information
        OBJ_SANMGMT = sanmgmt.create_sanmgmt_for_mpath()
        if not OBJ_SANMGMT:
            print "FAIL: Could not create LibSAN object"
            sys.exit(1)

    if not lun_name:
        print "FAIL: lun_map() - requires lun_name as parameter"
        return False

    return OBJ_SANMGMT.lun_map(lun_name, t_addr=t_addr, i_addr=i_addr, rescan=True)

def lun_unmap(lun_name, t_addr=None, i_addr=None):
    global OBJ_SANMGMT
    if not OBJ_SANMGMT:
        #Initialize SAN obj based on all information configured on /etc/san_top.cfg
        #based on HW information
        OBJ_SANMGMT = sanmgmt.create_sanmgmt_for_mpath()
        if not OBJ_SANMGMT:
            print "FAIL: Could not create LibSAN object"
            sys.exit(1)

    if not lun_name:
        print "FAIL: lun_unmap() - requires lun_name as parameter"
        return False

    return OBJ_SANMGMT.lun_unmap(lun_name, t_addr=t_addr, i_addr=i_addr)

def lun_info(lun_name):
    global OBJ_SANMGMT
    if not OBJ_SANMGMT:
        #Initialize SAN obj based on all information configured on /etc/san_top.cfg
        #based on HW information
        OBJ_SANMGMT = sanmgmt.create_sanmgmt_for_mpath()
        if not OBJ_SANMGMT:
            print "FAIL: Could not create LibSAN object"
            sys.exit(1)

    if not lun_name:
        print "FAIL: lun_info() - requires lun_name as parameter"
        return False

    info = OBJ_SANMGMT.lun_info(lun_name)
    if info:
        print info
        return True
    print "FAIL: Could not get info for %s" % lun_name
    return False

def lun_show():
    global OBJ_SANMGMT
    if not OBJ_SANMGMT:
        #Initialize SAN obj based on all information configured on /etc/san_top.cfg
        #based on HW information
        OBJ_SANMGMT = sanmgmt.create_sanmgmt_for_mpath()
        if not OBJ_SANMGMT:
            print "FAIL: Could not create LibSAN object"
            sys.exit(1)

    luns = OBJ_SANMGMT.lun_query()
    if not luns:
        return False
    for lun_name in luns:
        print "%s" % (lun_name)
    return True

def remove_device_wwid(wwid):
    return linux.remove_device_wwid(wwid)

def rescan_host(hostid):
    return scsi.rescan_host(hostid)

def dt_stress(of, log=None, thread=None, runtime=None):
    return dt.dt_stress(of, log=log, thread=thread, time=runtime)

def dt_stress_background(of, log=None, thread=None, runtime=None):
    return dt.dt_stress_background(of, log=log, thread=thread, time=runtime)

def fio_stress(of, threads=None, runtime=None, verify=None):
    ret, output = fio.fio_stress(of, threads=threads, runtime=runtime,
                                 verify=verify, return_output=True)
    if output:
        print output
    return ret

def fio_stress_background(of, threads=None, runtime=None, verify=None):
    return fio.fio_stress_background(of, threads=threads, runtime=runtime, verify=verify)

def main():
    parser = argparse.ArgumentParser(description='LibSAN tool')
    #parser.add_argument('command', help='Subcommand to run')

    subparsers = parser.add_subparsers(help="Valid commands", dest="command")
    #parser_show_host_cfg = subparsers.add_parser('show_host_config')
    #parser_show_host_cfg.set_defaults(func=show_host_cfg)
    subparsers.add_parser('show_host_config')

    subparsers.add_parser('setup_soft_fcoe')

    subparsers.add_parser('create_basic_iscsi_target')

    parser_lun_create = subparsers.add_parser('lun_create')
    parser_lun_create.add_argument("--lun_name", "-n", required=False,
                                   dest="lun_name", metavar="lun_name", help="Lun Name.")
    parser_lun_create.add_argument("--size", "-s", required=True,
                                   dest="lun_size", metavar="size",
                                   help="Lun Size. Eg. 2MiB, 2GiB,<F12> 2TiB")

    parser_lun_create_and_map = subparsers.add_parser('lun_create_and_map')
    parser_lun_create_and_map.add_argument("--lun_name", "-n", required=False,
                                           dest="lun_name", metavar="lun_name", help="Lun Name.")
    parser_lun_create_and_map.add_argument("--size", "-s", required=True,
                                           dest="lun_size", metavar="size",
                                           help="Lun Size. Eg. 2MiB, 2GiB, 2TiB")

    parser_lun_remove = subparsers.add_parser('lun_remove')
    parser_lun_remove.add_argument("--lun_name", "-n", required=True,
                                   dest="lun_name", metavar="lun_name", help="Lun Name.")

    parser_lun_map = subparsers.add_parser('lun_map')
    parser_lun_map.add_argument("--lun_name", "-n", required=True,
                                dest="lun_name", metavar="lun_name", help="Lun Name.")
    parser_lun_map.add_argument("--i_addr", "-i", required=False,
                                dest="i_addr", metavar="init_addr",
                                help="Initiator address. IQN/WWPN")
    parser_lun_map.add_argument("--t_addr", "-t", required=False,
                                dest="t_addr", metavar="target_addr",
                                help="Target address. IQN/WWPN")

    parser_lun_unmap = subparsers.add_parser('lun_unmap')
    parser_lun_unmap.add_argument("--lun_name", "-n", required=True,
                                  dest="lun_name", metavar="lun_name", help="Lun Name.")
    parser_lun_unmap.add_argument("--i_addr", "-i", required=False,
                                  dest="i_addr", metavar="init_addr",
                                  help="Initiator address. IQN/WWPN")
    parser_lun_unmap.add_argument("--t_addr", "-t", required=False,
                                  dest="t_addr", metavar="target_addr",
                                  help="Target address. IQN/WWPN")

    parser_lun_info = subparsers.add_parser('lun_info')
    parser_lun_info.add_argument("--lun_name", "-n", required=True,
                                 dest="lun_name", metavar="lun_name", help="Lun Name.")

    subparsers.add_parser('lun_show')

    subparsers.add_parser('check_ports_ready')

    parser_rcsn_enable = subparsers.add_parser('rcsn_enable')
    parser_rcsn_enable.add_argument("--addr", "-a", required=True,
                                    dest="port_addr", metavar="addr",
                                    help="WWPN/MAC of port to enable RCSN")

    parser_rcsn_disable = subparsers.add_parser('rcsn_disable')
    parser_rcsn_disable.add_argument("--addr", "-a", required=True,
                                     dest="port_addr", metavar="addr",
                                     help="WWPN/MAC of port to enable RCSN")

    parser_port_connect = subparsers.add_parser('port_connect')
    parser_port_connect.add_argument("--addr", "-a", required=True,
                                     dest="port_addr", metavar="addr",
                                     help="WWPN/MAC connected to the port")

    parser_port_disconnect = subparsers.add_parser('port_disconnect')
    parser_port_disconnect.add_argument("--addr", "-a", required=True,
                                        dest="port_addr", metavar="addr",
                                        help="WWPN/MAC connected to the port")

    parser_link_torture = subparsers.add_parser('link_torture')
    parser_link_torture.add_argument("--addr", "-a", required=True,
                                     dest="addr", metavar="address",
                                     help="WWPN/MAC connected to switch")
    parser_link_torture.add_argument("--runtime", "-t", required=True,
                                     dest="runtime", metavar="runtime",
                                     help="For how long it should run")
    parser_link_torture.add_argument("--interval", "-i", required=False,
                                     dest="interval", metavar="interval",
                                     help="How long switch port stays on state before changing")

    parser_remove_device_wwid = subparsers.add_parser('remove_device_wwid')
    parser_remove_device_wwid.add_argument("--wwid", "-w", required=True,
                                           dest="wwid", metavar="wwid",
                                           help="SCSI device WWID")

    parser_rescan_host = subparsers.add_parser('rescan_host')
    parser_rescan_host.add_argument("--hostid", "-H", required=False,
                                    dest="hostid", metavar="hostid", help="Host ID")

    parser_dt_stress = subparsers.add_parser('dt_stress')
    parser_dt_stress.add_argument("--of", "-o", required=True,
                                  dest="of", metavar="output",
                                  help="output, can be a device or file")
    parser_dt_stress.add_argument("--time", "-t", required=True,
                                  dest="time", metavar="time",
                                  help="For how long it should run")
    parser_dt_stress.add_argument("--log", "-l", required=False,
                                  dest="log", metavar="log", help="File to store log")
    parser_dt_stress.add_argument("--thread", "-T", required=False,
                                  dest="threads", metavar="threads", help="Number of threads")

    parser_dt_stress_background = subparsers.add_parser('dt_stress_background')
    parser_dt_stress_background.add_argument("--of", "-o", required=True,
                                             dest="of", metavar="output",
                                             help="output, can be a device or file")
    parser_dt_stress_background.add_argument("--time", "-t", required=True,
                                             dest="time", metavar="time",
                                             help="For how long it should run")
    parser_dt_stress_background.add_argument("--log", "-l", required=False,
                                             dest="log", metavar="log", help="File to store log")
    parser_dt_stress_background.add_argument("--threads", "-T", required=False,
                                             dest="threads", metavar="threads",
                                             help="Number of threads")

    parser_fio_stress = subparsers.add_parser('fio_stress')
    parser_fio_stress.add_argument("--of", "-o", required=True,
                                   dest="of", metavar="output",
                                   help="output, can be a device or file")
    parser_fio_stress.add_argument("--time", "-t", required=False,
                                   dest="time", metavar="time",
                                   help="For how long it should run")
    parser_fio_stress.add_argument("--verify", "-V", required=False,
                                   dest="verify", metavar="type", help="crc32c, md5")
    parser_fio_stress.add_argument("--threads", "-T", required=False,
                                   dest="threads", metavar="threads", help="Number of threads")

    parser_fio_stress_background = subparsers.add_parser('fio_stress_background')
    parser_fio_stress_background.add_argument("--of", "-o", required=True,
                                              dest="of", metavar="output",
                                              help="output, can be a device or file")
    parser_fio_stress_background.add_argument("--time", "-t", required=False,
                                              dest="time", metavar="time",
                                              help="For how long it should run")
    parser_fio_stress_background.add_argument("--verify", "-V", required=False,
                                              dest="verify", metavar="type", help="crc32c, md5")
    parser_fio_stress_background.add_argument("--threads", "-T", required=False,
                                              dest="threads", metavar="threads",
                                              help="Number of threads")

    args = parser.parse_args()

    #if args.show_host_config:
    if args.command == "show_host_config":
        if show_host_cfg():
            sys.exit(0)
        sys.exit(1)

    #if args.setup_soft_fcoe:
    if args.command == "setup_soft_fcoe":
        if fcoe.setup_soft_fcoe():
            sys.exit(0)
        sys.exit(1)

    if args.command == "create_basic_iscsi_target":
        if lio.lio_setup_iscsi_target():
            sys.exit(0)
        sys.exit(1)

    if args.command == "lun_create":
        if lun_create(size=args.lun_size, lun_name=args.lun_name):
            sys.exit(0)
        sys.exit(1)

    if args.command == "lun_create_and_map":
        if lun_create_and_map(size=args.lun_size, lun_name=args.lun_name):
            sys.exit(0)
        sys.exit(1)

    if args.command == "lun_remove":
        if lun_remove(args.lun_name):
            sys.exit(0)
        sys.exit(1)

    if args.command == "lun_map":
        if lun_map(args.lun_name, t_addr=args.t_addr, i_addr=args.i_addr):
            sys.exit(0)
        sys.exit(1)

    if args.command == "lun_unmap":
        if lun_unmap(args.lun_name, t_addr=args.t_addr, i_addr=args.i_addr):
            sys.exit(0)
        sys.exit(1)

    if args.command == "lun_info":
        if lun_info(args.lun_name):
            sys.exit(0)
        sys.exit(1)

    if args.command == "lun_show":
        if lun_show():
            sys.exit(0)
        sys.exit(1)

    if args.command == "check_ports_ready":
        if check_ports_ready():
            sys.exit(0)
        sys.exit(1)

    if args.command == "rcsn_enable":
        if rcsn_enable(args.port_addr):
            sys.exit(0)
        sys.exit(1)

    if args.command == "rcsn_disable":
        if rcsn_disable(args.port_addr):
            sys.exit(0)
        sys.exit(1)

    if args.command == "port_connect":
        if port_connect(args.port_addr):
            sys.exit(0)
        sys.exit(1)

    if args.command == "port_disconnect":
        if port_disconnect(args.port_addr):
            sys.exit(0)
        sys.exit(1)

    if args.command == "link_torture":
        if link_torture(args.addr, args.runtime, args.interval):
            sys.exit(0)
        sys.exit(1)

    if args.command == "remove_device_wwid":
        if remove_device_wwid(args.wwid):
            sys.exit(0)
        sys.exit(1)

    if args.command == "rescan_host":
        if rescan_host(args.hostid):
            sys.exit(0)
        sys.exit(1)

    if args.command == "dt_stress":
        if dt_stress(args.of, log=args.log, thread=args.threads, runtime=args.time):
            sys.exit(0)
        sys.exit(1)

    if args.command == "dt_stress_background":
        if dt_stress_background(args.of, log=args.log, thread=args.threads,
                                runtime=args.time):
            sys.exit(0)
        sys.exit(1)

    if args.command == "fio_stress":
        if fio_stress(args.of, verify=args.verify, threads=args.threads, runtime=args.time):
            sys.exit(0)
        sys.exit(1)

    if args.command == "fio_stress_background":
        pid = fio_stress_background(args.of, verify=args.verify, threads=args.threads,
                                    runtime=args.time)
        if not pid:
            sys.exit(0)
        sys.exit(pid)
    print "FAIL: Unsupported command: %s" % args.command
    sys.exit(1)


main()
