#! /usr/bin/python2
#------------------------------------------------------------------------------
#
# Copyright (C) 2020-2021 NVIDIA Corporation. ALL RIGHTS RESERVED.
# Copyright 2013,2014,2015,2016,2017,2018,2019 Cumulus Networks Inc.
# All rights reserved.
#
#------------------------------------------------------------------------------
#
#------------------------------------------------------------------------------
#

### Imports

import argparse
import ConfigParser
import contextlib  # for set_temp_host
import copy
import datetime
import fcntl
import json
import hashlib
import os
import shutil
import signal
import subprocess
import sys
import syslog
import tempfile  # for str_to_tmp_file
import time
import urllib2
import re
import glob
from collections import OrderedDict


### Variables

DEVNULL = open(os.devnull, 'wb')
date = datetime.datetime
log_handler = None
osrelease_file = '/etc/os-release'
imgrelease_file = '/etc/mlnx-release'
save_script = True
script_flag = 'BLUEFIELD-AUTOPROVISIONING'
cascade_flag = 'BLUEFIELD-AUTOPROVISIONING-CASCADE'
front_panel_flag = 'BLUEFIELD-AUTOPROVISIONING-FRONT-PANEL'
allow_cascade = False
cascade_count = 0
allow_front_panel = False
state_dir = "/var/lib/bluefield/ztp"
state_date = date.utcnow().strftime("%c UTC")
state_file = state_dir + "/ztp_state.log"
state = ConfigParser.ConfigParser()
verbose = False
version = '1.0'
waterfall = False
ztp_dhcp = "/var/run/ztp.dhcp"
ztp_script = "/var/lib/bluefield/ztp/ztp_script-{}-{}"
ztp_lock = "/var/run/ztp.lock"
sha_file = "/var/lib/bluefield/ztp/ztp_state.sha"
ztp_sha_lock = "/var/run/ztp_sha.lock"
ztp_run = "/var/run/ztp.run"
cleanup_signal = None
ztp_in_progress = False
mgmt_port = 'oob_net0'

### Logging

syslog_priority_map = {"crit": syslog.LOG_CRIT,
                       "error": syslog.LOG_ERR,
                       "info": syslog.LOG_INFO,
                       "warn": syslog.LOG_WARNING,
                       "debug": syslog.LOG_DEBUG}


stdout_priority_map = {"crit": "error",
                       "error": "error",
                       "info": "",
                       "warn": "warning",
                       "debug": "debug"}


def log_init():
    global log_handler
    if verbose is True:
        log_handler = log_handler_stdout
    else:
        syslog.openlog("ztp ",
            syslog.LOG_CONS | syslog.LOG_PID, syslog.LOG_DAEMON)
        log_handler = log_handler_syslog


def log_handler_stdout(priority, buf):
    p = stdout_priority_map.get(priority, "")
    if p != "":
        p = p + ': '
    print(p + buf)
    sys.stdout.flush()


def log_handler_syslog(priority, buf):
    syslog.syslog(syslog_priority_map.get(priority, syslog.LOG_INFO), buf)


def log_error(*args, **kwargs):
    log_handler("error", ''.join(args))


def log_warn(*args, **kwargs):
    log_handler("warn", ''.join(args))


def log_debug(*args, **kwargs):
    # Eat debug messages, unless running with --verbose
    if verbose is True:
        log_handler("debug", ''.join(args))


def log_info(*args, **kargs):
    log_handler("info", ''.join(args))


def write_wall_msg(msg, logfn=None):
    try:
        if logfn:
            logfn(msg)

        msg = 'ZTP: {}'.format(msg)
        w = subprocess.Popen('/usr/bin/wall', stdin=subprocess.PIPE,
                             stderr=subprocess.PIPE)
        w.communicate(input=msg)
        w.stdin.close()
        w.stderr.close()
        w.wait()
    except:
        pass


### State Setup

def setup(isboot):
    do_init = False
    if os.path.exists(state_dir) is False:
        log_debug('%s: State Directory does not exist. Creating it...'
                                                                % state_dir)
        try:
            os.makedirs(state_dir, 0o755)
        except OSError:
            log_error('%s: Could not create create directory.' % state_dir)
            sys.exit(1)
        if isboot is True:
            do_init = True
    elif isboot is True and not os.path.isfile(sha_file):
        do_init = True

    if do_init is True:
	    log_debug('Boot mode with no state, initializing state')
	    # first run after install or after manual remove, so initialize
	    initialize_config()

    if not os.path.isfile(ztp_lock):
        log_debug('%s: Lock File does not exist. Creating it...' % ztp_lock)
        try:
            with open(ztp_lock, 'wb') as LOCKFILE:
                LOCKFILE.write("")
        except:
            log_error('%s: Could not create file.' % ztp_lock)
            sys.exit(1)


def state_file_setup():
    success = False
    lock = exclusive_lock(ztp_lock)
    if os.path.isfile(state_file):
        state.read(state_file)
    else:
        log_debug('%s: State File does not exist. Creating it...' % state_file)
        state.add_section("STATUS")

    if not state.has_section("Most Recent"):
        state.add_section("Most Recent")
    else:
        try:
            new_section = state.get("Most Recent", "DATE")
        except:
            log_error("Missing section or option in State File")
            log_warn("Erasing Most Recent Section")
            state.remove_section("Most Recent")
            state.add_section("Most Recent")
        else:
            state.add_section(new_section)
            for option, value in state.items("Most Recent"):
                state.set(new_section, option, value)
                state.remove_option("Most Recent", option)
    try:
        with open(state_file, 'wb') as STATEFILE:
            state.write(STATEFILE)
        release_lock(lock)
    except:
        release_lock(lock)
        return False
    return True


def on_air():
    global ztp_in_progress

    try:
        lock = exclusive_lock(ztp_lock)
        with open(ztp_run, 'wb') as f:
            f.write(str(os.getpid()))
    except:
        log_error('{}: Could not create file.'.format(ztp_run))
    else:
        ztp_in_progress = True
    release_lock(lock)


def is_on_air():
    try:
        if os.path.isfile(ztp_run):
            with open(ztp_run, 'r') as f:
                pid = f.read().strip('\n')
                with open('/proc/{}/comm'.format(pid), 'r') as p:
                    name = p.read().strip('\n')
                    if name == 'ztp':
                        return True
    except:
        log_info('Cannot verify ztp pid')

    return False


def off_air():
    global ztp_in_progress

    if ztp_in_progress is True and os.path.exists(ztp_run):
        os.remove(ztp_run)
    ztp_in_progress = False


### HTTP Headers

def mgmt_mac_address(intf):
    try:
        with open('/sys/class/net/%s/address' % intf,'r') as address:
            mac = address.read().strip()
        with open('/sys/class/net/%s/carrier' % intf,'r') as carrier:
            active = int(carrier.read())
        if active == 1:
            return mac
    except:
        log_error("Problems getting port MAC: ", intf)

    return None


def mgmt_ip_address(intf):
    cmd = ['ip', '-j', 'addr', 'show', intf]
    try:
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE)
        ipaddr = json.loads(proc.stdout.read())[0]
        ip = [ ip.get('local') for ip in ipaddr.get('addr_info', {}) ]
        if ipaddr.get('operstate') == 'UP':
            return ",".join(ip)
    except subprocess.CalledProcessError:
        log_error("Problems getting port IP: ", intf)

    return None


def bmc_addresses():
    mac, ip = None, None
    cmd = ['ipmitool', 'lan', 'print']
    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE,
                                 stderr=subprocess.STDOUT)
    output = proc.communicate()[0]
    res = re.search(r'^MAC Address\s*:\s(?P<mac>.+)$', output, re.MULTILINE)
    if res:
        mac = res.groupdict().get('mac')
    res = re.search(r'^IP Address\s*:\s(?P<ip>.+)$', output, re.MULTILINE)
    if res:
        ip = res.groupdict().get('ip')
    return mac, ip


def bmc_firmware():
    fwver = None
    cmd = ['ipmitool', 'mc', 'info']
    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE,
                                 stderr=subprocess.STDOUT)
    output = proc.communicate()[0]
    res = re.search(r'^Firmware Revision\s*:\s(?P<rev>.+)$', output, re.MULTILINE)
    if res:
        fwver = res.groupdict().get('rev')
    else:
        return None
    res = re.search(r'Aux Firmware Rev Info\s*:\s(?P<aux>.+)', output, re.DOTALL)
    if res:
        fwaux = res.groupdict().get('aux', '').split()
        if fwver and len(fwaux) > 1:
            fwver = '{}.{:02d}-{:02d}'.format(fwver, int(fwaux[0], 16), int(fwaux[1], 16))
    return fwver

def osrelease_val(field):
    for line in open(osrelease_file, 'r'):
        field2 = field + "="
        pos = line.find(field2)
        if pos > -1:
            pos2 = pos + len(field) + 1
            return line[pos2:len(line)].strip().replace("\"", "")
    return ""


def uname_arch():
    cmd = ['/bin/uname', '-m']
    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE)
    for line in proc.stdout:
        return line.strip()
    return ""

def ofed_ver():
    cmd = ['/usr/bin/ofed_info', '-s']
    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE)
    for line in proc.stdout:
        return line.strip()
    return ""

def imgrelease_val():
    try:
        with open(imgrelease_file) as f:
            ver = f.read().strip()
        return ver
    except:
        pass

    return ""

def get_vpd():
    vpd = {}
    cmd = [ 'mstvpd', '03:00.0' ]
    try:
        output = subprocess.check_output(cmd).strip()
        vpd = {k:v.strip() for k,v in re.findall(r'(?P<k>\w+):(?P<v>.+)', output, re.MULTILINE)}
    except subprocess.CalledProcessError, cpe:
        log_error("Problems reading VPD: %s" % cpe)
    except OSError, ose:
        log_error("Problems running mstvpd: %s" % ose)
    return vpd

def add_bluefield_headers(req):
    mfgr = "NVIDIA"
    prodname = "unknown"
    serial = "unknown"

    mac = mgmt_mac_address(mgmt_port)
    if mac is None:
        mac = 'unknown'
    ip = mgmt_ip_address(mgmt_port)
    if not ip:
        ip = 'unknown'

    # get adapter VPD information
    vpd = get_vpd()
    prodname = vpd.get('ID', 'unknown')
    partnum = vpd.get('PN', 'unknown')
    serial = vpd.get('SN', 'unknown')
    revision = vpd.get('EC', 'unknown')

    # get onboard BMC information
    bmc_mac, bmc_ip = bmc_addresses()
    bmc_fw = bmc_firmware()
    if bmc_mac is None:
        bmc_mac = 'unknown'
    if bmc_ip is None:
        bmc_ip = 'unknown'
    if bmc_fw is None:
        bmc_fw = 'unknown'

    # assemble headers to send the server
    pfx = "BLUEFIELD"
    req.add_header('User-agent', "BlueField-AutoProvision/%s" % (version))
    req.add_header(pfx + '-ARCH', uname_arch())
    req.add_header(pfx + '-BUILD', osrelease_val("VERSION"))
    req.add_header(pfx + '-MANUFACTURER', mfgr)
    req.add_header(pfx + '-PRODUCTNAME', prodname)
    req.add_header(pfx + '-PARTNUMBER', partnum)
    req.add_header(pfx + '-REVISION', revision)
    req.add_header(pfx + '-SERIAL', serial)
    req.add_header(pfx + '-MGMT-MAC', mac)
    req.add_header(pfx + '-MGMT-IP', ip)
    req.add_header(pfx + '-OS-VERSION', osrelease_val("VERSION_ID"))
    req.add_header(pfx + '-IMG-RELEASE', imgrelease_val())
    req.add_header(pfx + '-OFED-VERSION', ofed_ver())
    req.add_header(pfx + '-BMC-MAC', bmc_mac)
    req.add_header(pfx + '-BMC-IP', bmc_ip)
    req.add_header(pfx + '-BMC-FIRMWARE', bmc_fw)
    return req


### Lock Functions

def shared_lock(file_lock):
    lock = None
    try:
        lock = open(file_lock, 'r')
        fcntl.lockf(lock, fcntl.LOCK_SH | fcntl.LOCK_NB)
    except IOError:
        lock = None
        log_error("Another instance is running. Exiting...")
    if lock is None:
        sys.exit(1)
    return lock


def exclusive_lock(file_lock):
    lock = None
    try:
        lock = open(file_lock, 'w')
        fcntl.lockf(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except IOError:
        lock = None
        log_error("Another instance is running. Exiting...")
    if lock is None:
        sys.exit(1)
    return lock


def release_lock(lock):
    try:
        lock.close()
    except:
        log_warn("Could not release the lock correctly")
    return


### Read and Write Functions


def systemctl_call(argument):

    cmd = ['/bin/systemctl', argument, 'ztp.service']

    if argument == "is-enabled":
        try:
            sub = subprocess.Popen(cmd, stdout=subprocess.PIPE)
            status = sub.communicate()[0].replace("\n", "")
            sub.wait()
        except subprocess.CalledProcessError:
            status = "systemctl error"
        return status
    elif argument == "stop":
        try:
            subprocess.check_output('/bin/systemctl stop ztp.service'.split())
            return True
        except subprocess.CalledProcessError:
            log_error("Could not stop ZTP")

        return False
    else:
        status = argument + "d"
        current_state = systemctl_call("is-enabled")
        if status != current_state:
            try:
                subprocess.check_output(cmd)
            except subprocess.CalledProcessError:
                log_error("Could not %s ZTP." % argument)
                return False
            else:
                log_info("ZTP " + argument + "d")
                lock = shared_lock(ztp_lock)
                if os.path.isfile(state_file):
                    try:
                        state.read(state_file)
                    except ConfigParser.Error:
                        log_error("Could not read State File correctly")
                        return False
                    else:
                        current_state = systemctl_call("is-enabled")
                if state.has_option("STATUS", "ZTP") \
                and state.get("STATUS", "ZTP") != current_state:
                    release_lock(lock)
                    lock = exclusive_lock(ztp_lock)
                    state.set("STATUS", "ZTP", current_state)
                    with open(state_file, 'wb') as STATEFILE:
                        state.write(STATEFILE)
                        success = True
                    release_lock(lock)
        return True


def print_latest_ztp_logs():
    cmd = ['/usr/bin/journalctl', '/usr/sbin/ztp', '-n', '20']
    try:
        output = subprocess.check_output(cmd)
        print('Recent ZTP logs:\n')
        print(output)
    except subprocess.CalledProcessError, error:
        log_error('Could not retrieve ZTP logs')


def print_json_state():
    print_state("json")


def print_state(json_arg):
    current_state = systemctl_call("is-enabled")
    output_dict = {'state': current_state, 'version': version}
    active_state = is_on_air()
    if active_state:
        if json_arg is None:
            print('\nZTP in progress...\n')
            print_latest_ztp_logs()
        else:
            output_dict['state'] = 'In Progress'
            print(json.dumps(output_dict, sort_keys=False, indent=4))
        sys.exit(0)

    lock = shared_lock(ztp_lock)
    if os.path.isfile(state_file):
        try:
            state.read(state_file)
        except ConfigParser.Error:
            log_error("Could not read State File correctly")
        else:
            try:
                state.items("Most Recent")
            except ConfigParser.Error:
                log_error("Missing section in State File")
                return

            result_dict = OrderedDict()
            for (name, val) in state.items("Most Recent"):
                output_dict[name] = val
                i = re.search(" \d", name).group(0) if re.search(" \d", name) else None
                if i and json_arg:
                    if i not in result_dict:
                        result_dict[i] = {}
                    result_dict[i][name.split(i)[0]] = val
                    del output_dict[name]
            if result_dict:
                output_dict["ztp results"] = list(result_dict.values())

            if json_arg is not None:
                print(json.dumps(output_dict, sort_keys=False, indent=4))
            else:
                print("Last ZTP INFO:\n")
                print("\n".join("{}: {}".format(k, v) for k, v in sorted(output_dict.items())))
    else:
        if json_arg is not None:
            print(json.dumps(output_dict, sort_keys=False, indent=4))
        else:
            print("\nZTP INFO:\n")
            print("\n".join("{}: {}".format(k, v) for k, v in sorted(output_dict.items())))

    release_lock(lock)


def write_state(ztp_return, method, url):
    global allow_cascade
    global cascade_count
    success = False
    if os.path.isfile(state_file):
        try:
            state.read(state_file)
        except ConfigParser.Error:
            log_error("Could not read State File correctly")
        else:
            if state.has_section("STATUS") \
            and state.has_section("Most Recent"):
                state.set("STATUS", "ZTP", systemctl_call("is-enabled"))
                state.set("Most Recent", "VERSION", version)
                if allow_cascade or cascade_count != 0:
                    cascade_count += 1
                    cascade_str = " {}".format(str(cascade_count))
                else:
                    cascade_str = ""
                state.set("Most Recent", "Result{}".format(cascade_str), ztp_return)
                state.set("Most Recent", "METHOD{}".format(cascade_str), method)
                state.set("Most Recent", "URL{}".format(cascade_str), url)
                state.set("Most Recent", "DATE{}".format(cascade_str), date.utcnow().strftime("%c UTC"))
                state.set("Most Recent", "DATE", date.utcnow().strftime("%c UTC"))
                with open(state_file, 'wb') as STATEFILE:
                    state.write(STATEFILE)
                    success = True
            else:
                log_error("Missing section in State File")
    else:
        log_error("State File not found")
    if success is False:
        log_error("Could not write in the State File")
    return success


def sha_list_generate(check_list):
    sha_list = []
    for file in check_list:
        if not os.path.isfile(file):
            sha_list.append("")
        else:
            with open(file) as file_to_check:
                data = file_to_check.read()
                sha = hashlib.sha256(data).hexdigest()
                sha_list.append(sha)
    return sha_list

def manual_config_check_list():
    '''
    Generate list of files to check if user has manually configured
    the system.
    '''

    file_list = [
        '/etc/openvswitch/conf.db',
        '/etc/network/interfaces', 
        '/etc/frr/frr.conf',
        '/etc/passwd',
        '/etc/group',
        '/etc/shadow'
        ]
    try:
        if os.path.isdir('/etc/netplan'):
            files = os.listdir('/etc/netplan')
            file_list.extend(map(lambda n: '/etc/netplan/' + n,
                                 filter(lambda x: x.endswith('.yaml'), files)))
        if os.path.isdir('/etc/network/interfaces.d'):
            files = os.listdir('/etc/network/interfaces.d')
            file_list.extend(map(lambda n: '/etc/network/interfaces.d/' + n,
                                 filter(lambda x: x.endswith('.intf'), files)))
    except:
        pass

    return file_list

def initialize_config():
    '''
    Initialize the config database.  Record the current state as the
    the (fresh install) sha256 of various configuration files.
    May happen multiple times while people are testing, so not necessarily
    fresh install state.
    '''

    shaconf = ConfigParser.ConfigParser()
    check_list = manual_config_check_list()

    if not os.path.isfile(ztp_sha_lock):
        with open(ztp_sha_lock, 'wb') as LOCKFILE:
            LOCKFILE.write("")
    lock = exclusive_lock(ztp_sha_lock)
    sha_list = sha_list_generate(check_list)

    # it used to allow writing once to this file. With the need to
    # rebase the hash, relaxing the restriction
    # //if not os.path.isfile(sha_file):
    shaconf.add_section("SHASUM")
    for key, value in zip(check_list, sha_list):
        shaconf.set("SHASUM", key, value)
    with open(sha_file, 'wb') as SHA1FILE:
        shaconf.write(SHA1FILE)
    release_lock(lock)

def check_initialized_config():
    '''
    Compare the current sha256 hashes of various configuration files
    to the hashes captured during the initial boot of the system.

    Return True if any of the hashes are different, False otherwise.

    '''

    shaconf = ConfigParser.ConfigParser()
    check_list = manual_config_check_list()

    if not os.path.isfile(sha_file):
        log_error("ZTP not initialized correctly, run with -R")
        ztp_exit(False)
    sha_list = sha_list_generate(check_list)
    try:
        shaconf.read(sha_file)
    except ConfigParser.Error:
        log_error("Could not read sha File correctly")
        ztp_exit(False)
    for key, value in zip(check_list, sha_list):
        v = None
        try:
            v = shaconf.get("SHASUM", key)
        except:
            pass

        if v != value:
            write_wall_msg("DPU has already been configured. "
                           + key + " was modified", log_warn)
            write_wall_msg("ZTP will not run", log_warn)
            lock = exclusive_lock(ztp_lock)
            # we have both "failure" and "failed" status, should
            # have removed one of them but that may break existing scripts,
            # hence leaving it unchanged
            write_state("failed", "DPU manually configured", "None")
            release_lock(lock)
            return True

    # No changes detected
    return False

@contextlib.contextmanager
def str_to_tmp_file(data):
    """ Write a string to a temp file
    file is deleted after exiting the with block.

    >>> with str_to_tmp_file("foo") as tmpfoo:
    ...     print 'tempfilename:', tmpfoo
    ...     print open(tmpfoo).read()
    """
    file_fd, filename = tempfile.mkstemp(prefix='/root/ztp',text=True)
    os.write(file_fd, data)
    os.close(file_fd)
    yield filename
    try:
        os.unlink(filename)
    except:
        pass


def get_platform_info():
    vpd = get_vpd()
    vendor = 'nvidia'
    model = vpd.get('PN', 'unknown')
    revision = vpd.get('EC', 'unknown')
    arch = subprocess.check_output(['/bin/uname', '-m']).rstrip()

    return vendor, model, revision, arch


def waterfall_search(method, directory, partition=None):
    global waterfall
    if waterfall is False:
        vendor, model, revision, arch = get_platform_info()
        waterfall_parts = ['bluefield-ztp', '-', arch, '-', vendor,
                        '_', model, '-r', revision]
        waterfall_parts_len = len(waterfall_parts)
        waterfall = [''.join(waterfall_parts[:i])
                     for i in xrange(1, waterfall_parts_len + 1, 2)]
        waterfall = waterfall[::-1]

    if waterfall is False:
        return False

    for filename in waterfall:
        script = directory + '/' + filename
        log_debug(method + ': Waterfall search for ' + script)
        if os.path.isfile(script):
            log_info(method + ': Found matching name: ' + script)
            script = "file://" + script
            return tryurl(script, True, method, partition)
    return False


def ztp_local_event():
    log_debug("ZTP LOCAL: Looking for ZTP local Script")
    return waterfall_search("ZTP LOCAL", state_dir)


def enable_dhclient(port):
    try:
        cmd = ['dhclient', '-nw', '-pf', '/run/dhclient.{}.pid'.format(port),
               '-lf', '/var/lib/dhcp/dhclient.{}.leases'.format(port), port]
        output = subprocess.check_output(cmd)
    except subprocess.CalledProcessError, error:
        log_error('ZTP failed to start dhclients %s' % error.returncode)


def disable_dhclients():
    # diable all the dhclients started by ztp
    up_ports, down_ports = get_ports()
    try:
        cmd = ['dhclient', '-x', '-pf', '/run/dhclient.ztp.pid', mgmt_port]
        cmd.extend(up_ports)
        output = subprocess.check_output(cmd)

        # need a robust way of disabling dhclients started by netplan/networkd
    except subprocess.CalledProcessError, error:
        log_warn('ZTP disabling of dhclients failed')


def renew_dhcp_leases():
    global allow_front_panel

    # Force renew leases with dhclient -x, this includes timed out cases.
    log_info('Restarting dhclients')
    up_ports, down_ports = get_ports()
    try:
        # # stop dhclient that is started by ztp
        cmd = ['dhclient', '-x', '-pf', '/run/dhclient.ztp.pid', mgmt_port]
        cmd.extend(up_ports)
        output = subprocess.check_output(cmd)
    except:
        log_warn('ZTP failed to renew lease on %s (%s)' % (mgmt_port, error.returncode))

    # need a robust way of renewing leases started by netplan/networkd

    if allow_front_panel:
        try:
            # start dhclients on front ports
            intfs = ' '.join(p for p in up_ports)
            if len(intfs) > 0:
                cmd = 'dhclient -nw -pf /run/dhclient.ztp.pid {}'.format(intfs)
                output = subprocess.check_output(cmd.split())
        except subprocess.CalledProcessError, error:
            log_warn('ZTP failed to renew lease on %s (%s)' % (intfs, error.returncode))


def get_ports():
    uplink_ports = ['p0', 'p1']
    up_ports = []
    down_ports = []
    for port in uplink_ports:
        try:
            with open('/sys/class/net/%s/carrier' % port,'r') as carrier:
                active = int(carrier.read())
            if active == 1:
                up_ports.append(port)
            else:
                down_ports.append(port)
        except:
            pass

    return up_ports, down_ports


def enable_front_panel_ports(eval_breakout=False):
    global breakouts
    try:
        up_ports_0, down_ports = get_ports()
        for p in down_ports:
            cmd = ['/usr/sbin/ip', 'link', 'set', 'up', p]
            subprocess.check_output(cmd)

        time.sleep(10)

        up_ports, down_ports = get_ports()

        return up_ports
    except subprocess.CalledProcessError, error:
        log_error('ZTP cannot enumerate front panel ports %s'
                                                % error.returncode)
        return []


def get_vrf(interface):
    try:
        cmd = ['/usr/sbin/ip', '-j', 'link', 'show', interface]
        output = subprocess.check_output(cmd)
        iplink = json.loads(output)[0]
        vrf = iplink.get('master', 'default')
    except subprocess.CalledProcessError, error:
        log_error('ZTP cannot determine dhcp interface %s'
                                                % error.returncode)
        vrf = 'mgmt'

    return vrf


def ztp_dhcp_event():
    log_debug("ZTP DHCP: Looking for ZTP Script provided by DHCP")
    try:
        fd = os.open(ztp_dhcp, os.O_RDWR)
        intf_url = os.read(fd, 1024).replace("\n", "").split('|', 1)
        if len(intf_url) == 2:
            interface = intf_url[0]
            url = intf_url[1]
            vrf = get_vrf(interface)
        else:
            url = intf_url[0]
            #vrf = 'mgmt'
            vrf = 'default'
    except IOError:
        log_error("Could not open %s" % ztp_dhcp)
        return False
    os.close(fd)

    return tryurl(url, True, "ZTP DHCP")


# Returns True if there is no url error and script execution is successful.
# Returns False for any other failures.  Caller will decide what to do next.
#
def tryurl(url, doexec, method, partition=None):
    success = False
    provision_msg = 'Attempting to provision via {} from {}'.format(method, url)

    write_wall_msg(provision_msg, log_info)
    lock = exclusive_lock(ztp_lock)
    if url.lower().startswith("http://") or \
            url.lower().startswith("https://") or \
            url.lower().startswith("ftp://") or \
            url.lower().startswith("/") or \
            url.lower().startswith("file://"):

        if url.lower().startswith("/"):
            url = "file://" + url
        req = urllib2.Request(url)
        if url.lower().startswith("http"):
            add_bluefield_headers(req)
        for retry in range(1, 4):
            try:
                resp = urllib2.urlopen(req)
            except urllib2.URLError, error:
                write_wall_msg('{}: URL Error {}'.format(method, error.reason),
                               log_error)
            except urllib2.HTTPError, error:
                write_wall_msg('{}: HTTP Error {}'.format(method, error.code),
                               log_error)
            else:
                if url.lower().startswith("file://"):
                    url = url.replace("file://", "")
                scriptcontent = resp.read()
                log_info(method + ': URL response code %s' % resp.code)
                success = processweb(url, scriptcontent, doexec, method, partition)
                if success is False:
                    success = "script_failure"
                else :
                    write_wall_msg('Provisioning via %s complete' % method,
                                   log_info)
                    break
            if (retry != 4):
                time.sleep(5 * retry)
                log_info(method + ': Retrying {}'.format(retry))

    elif url.lower().startswith("tftp://"):
        try:
            cmd = ['/usr/bin/curl', '-s', url]
            scriptcontent = subprocess.check_output(cmd)
        except subprocess.CalledProcessError, error:
            err_msg = '{}: cURL error code {}'.format(method, error.returncode)
            write_wall_msg(err_msg, log_error)
        else:
            success = processweb(url, scriptcontent, doexec, method, partition)
            if success is False:
                success = "script_failure"
    else:
        err_msg = "{}: url's format not recognized".format(method)
        write_wall_msg(err_msg, log_error)

    if success is True:
        write_state("success", method, url)
    elif success is False:
        # false actually means url failure. Let the caller
        # decide whether to retry url acquisition (e.g. let dhclient get
        # a new url) or not
        write_state("failure", method, url)
    else:
        # scipt execution error, let caller decide what to do next
        write_wall_msg("Script returned failure", log_error)
        write_state("Script Failure", method, url)
    release_lock(lock)

    return success

# Check to see if the script has windows line endings, or
# if it has non-ASCII characters and warn accordingly, so
# the likely failures will not be surprising.   Not a
# fatal error, because there are valid cases for both
def check_script(file):
    ofd = -1
    ofile = file + '.cmp'
    try:
        ofd = os.open(ofile, os.O_WRONLY|os.O_TRUNC|os.O_CREAT)
        args = ['sed', 's/\\r$//g', file]
        subprocess.call(args, stdout=ofd)
        os.close(ofd)
        args = ['cmp', '-s', file, ofile]
        success = subprocess.call(args)
        if success != 0:
            log_warn('ZTP:' + file + ' appears to have CR LF line endings')
        ofd = os.open(ofile, os.O_WRONLY|os.O_TRUNC|os.O_CREAT)
        args = ['iconv', '-c', '-f', 'utf-8', '-t', 'ascii', '-o', ofile, file]
        subprocess.call(args)
        args = ['cmp', '-s', file, ofile]
        success = subprocess.call(args)
        if success != 0:
            log_warn('ZTP:' + file + ' has non-ASCII characters')
        os.remove(ofile)
    except OSError:
        log_info('Errors while checking ZTP script:' + file)
        if (ofd != -1):
                os.close(ofd)
                os.remove(ofile)

def save_ztp_script(script_file, method):
    script_date = date.utcnow().strftime('%Y%m%d%H%M')
    if method == "ZTP LOCAL":
        shutil.copyfile(script_file, ztp_script.format("local", script_date))
    else:
        shutil.copyfile(script_file, ztp_script.format("dhcp", script_date))

def processweb(url, scriptcontent, doexec, method, partition=None):
    global allow_cascade
    global allow_front_panel

    success = False
    if scriptcontent.find(cascade_flag) > -1:
        log_info(method + ": Found Marker " + cascade_flag)
        allow_cascade = True

    if scriptcontent.find(front_panel_flag) > -1:
        log_info(method + ": Found Marker " + front_panel_flag)
        allow_front_panel = True

    if scriptcontent.find(script_flag) > -1:
        log_info(method + ": Found Marker " + script_flag)
        log_info("{}: Executing {}".format(method, url))
        with str_to_tmp_file(scriptcontent) as script_file:
            os.chmod(script_file, 0o700)
            if save_script is True:
                save_ztp_script(script_file, method)
            ztp_env = os.environ.copy()
            ztp_env['ZTP_URL'] = url

            check_script(script_file)

            if method == "ZTP LOCAL":
                # Allowing time for dhcp to initialize the management ethernet
                time.sleep(5)
            try:
                if verbose is True:
                    subprocess.check_call(script_file, env=ztp_env)
                else:
                    proc = subprocess.Popen(script_file, env=ztp_env,
                                                stdout=subprocess.PIPE,
                                                stderr=subprocess.STDOUT)
                    output = proc.communicate()[0]
                    if output:
                        for line in output.split('\n'):
                            if line:
                                log_info(line)
                    if proc.returncode != 0:
                        raise subprocess.CalledProcessError(proc.returncode,
                                                            script_file,
                                                            output)
            except subprocess.CalledProcessError, cpe:
                err_msg = '{}: Payload returned code {}'.format(method, cpe.returncode)
                write_wall_msg(err_msg, log_error)
            except OSError, ose:
                if ose.errno == 2:
                    write_wall_msg(method + ": Could not find referenced "
                                   "script/interpreter in downloaded payload.",
                                   log_error)
                elif ose.errno == 8:
                    write_wall_msg(method + ": Could not find interpreter "
                                        "line (#!<path_to_interpreter>) "
                                        "in downloaded payload.", log_error)
                else:
                    err_msg = '{}: Unexpected OS error: {}'.format(method, str(ose))
                    write_wall_msg(err_msg, log_error)
            except SystemExit, se:
                if se.code == 0 and cleanup_signal == signal.SIGTERM:
                    write_wall_msg(method + ': user ztp script initiated system shut down.',
                                   log_info)
                    success = True
                else:
                    err_msg = '{}: Unexpected SystemExit code: {}'.format(method, str(se))
                    write_wall_msg(err_msg, log_error)
                    raise
            except BaseException, be:
                err_msg = '{}: Unexpected error: {}'.format(method, str(be))
                write_wall_msg(err_msg, log_error)
            else:
                log_info(method + ": Script returned success")
                success = True
    else:
        # code ok, but no markers
        err_msg = "{}: No marker '{}' found".format(method, script_flag)
        write_wall_msg(err_msg, log_error)
    return success


def cleanup(signum, frame):
    global cleanup_signal
    cleanup_signal = signum
    if signum == signal.SIGTERM:
        # Normal shutdown
        ztp_exit("terminate")
    else:
        ztp_exit(False)

def ztp_exit(success):
    global allow_front_panel

    off_air()

    # kill all dhclients that are not started by netplan/networkd
    if allow_front_panel:
        disable_dhclients()

    if success is True:
        log_info("ZTP service completes. Exiting...")
        systemctl_call("disable")
        sys.exit(0)
    elif success == "enabled_service":
        sys.exit(0)
    elif success == "script_failure":
        log_error("ZTP script failed. Exiting...")
        sys.exit(1)
    elif success == "missing_state":
        # This is a workaround when the ztp state could not be written
        # or is deleted by the user's script using the "ztp -R" command:
        # If the file is not here, we have no idea about the current state
        # of ZTP so we should keep it enabled and allow it to be run on the
        # next boot. Exit with a warning so the user is aware that the file
        # is missing and that ZTP will run on the next boot.
        log_warn("ZTP is still enabled and will run on next boot...")
        sys.exit(0)
    elif success == "systemd_exit":
        # This is a "successful" exit from the perspective of systemd,
        # though ztp itself did not run.  For example, this is the
        # exit state when ztp determines that the user has manually
        # configured the system.
        log_warn("Exiting and disabling ZTP...")
        systemctl_call("disable")
        sys.exit(2)
    elif success == "terminate":
        # This is termination by systemd, e.g. on reboot.  It should not
        # disable the service.
        log_info("ZTP service terminated")
        sys.exit(0)
    else:
        log_error("ZTP failed. Exiting and disabling ZTP...")
        systemctl_call("disable")
        sys.exit(1)


def main():
    """ main function """

    global allow_cascade
    global allow_front_panel

    if os.geteuid() != 0:
        print("ERROR: {fname} must be run as root".format(fname=__file__))
        sys.exit(1)

    os.umask (0o22)

    signal.signal(signal.SIGINT, cleanup)
    signal.signal(signal.SIGTERM, cleanup)
    signal.signal(signal.SIGQUIT, cleanup)

    parser = argparse.ArgumentParser(description="ZTP v%s" % version)

    # Command line arg parser

    group = parser.add_mutually_exclusive_group(required=False)

    group.add_argument('-b', '--boot', dest='boot',
            help="startup discovery", action="store_true")
    group.add_argument('-d', '--disable', dest='disable',
            help="disable ZTP", action="store_true")
    group.add_argument('-e', '--enable', dest='enable',
            help="Enable ZTP", action="store_true")
    parser.add_argument('-f', '--front', dest='front_panel',
            help="Enable front panel ports", action="store_true")
    group.add_argument('-j', '--json', dest='json',
            help="display json formatted details", action="store_true")
    group.add_argument('-q', '--queue', dest='qurl',
            help="DHCP queue", action="store")
    group.add_argument('-r', '--run', dest='url',
            help="run ZTP with an url", action="store")
    group.add_argument('-R', '--reset', dest='reset',
            help="reset ZTP", action="store_true")
    group.add_argument('-s', '--status', dest='status',
            help="display ZTP information, enable state, and results status",
            action="store_true")
    parser.add_argument('-u', '--unsaved', dest='unsaved',
            help='do not save the ZTP script in the State directory',
            action="store_true")
    group.add_argument('-V', '--version', action="store_true",
            help="display ZTP version", dest='version')
    parser.add_argument('-v', '--verbose', dest='log',
            help="enable verbosity", action="store_true")

    args = parser.parse_args()

    if args.unsaved is True:
        if len(sys.argv) == 2 \
        or (len(sys.argv) == 3 and args.log is True):
            sys.exit(0)
        global save_script
        save_script = False

    if args.log is True:
        global verbose
        verbose = True

    log_init()

    if args.qurl is not None:
        log_debug("Found ZTP DHCP Request", log_info)
        with str_to_tmp_file(args.qurl) as tmp_file:
            shutil.move(tmp_file, ztp_dhcp)

    setup(args.boot or args.reset)

    if not len(sys.argv) > 1 or args.status is True \
    or (len(sys.argv) == 2 and args.log is True):
        print_state(None)

    if args.version is True:
        print("ZTP version %s" % version)

    if args.disable is True:
        if systemctl_call("disable") is True:
            systemctl_call("stop")
            ztp_exit(True)
        else:
            sys.exit(1)

    if args.enable is True:
        if systemctl_call("enable") is True:
            ztp_exit("enabled_service")
        else:
            sys.exit(1)

    if args.reset is True:
        if systemctl_call("enable") is True:
            shutil.rmtree(state_dir)
            if os.path.exists(state_dir) is True:
                log_error("Could not remove State Directory")
                sys.exit(1)
            os.makedirs(state_dir, 0o755)
            if os.path.isfile(ztp_dhcp):
                os.remove(ztp_dhcp)
            initialize_config()
        else:
            sys.exit(1)

    if args.json is True:
        print_json_state()

    if args.boot is True or args.url is not None:
        if state_file_setup() is False:
            ztp_exit(False)

    if args.boot is True:
        if is_on_air():
            print('Another instance of ZTP is in progress.')
            sys.exit(1)

        allow_front_panel = args.front_panel

        # Check if the box has been manually configured
        if check_initialized_config() is True:
            print('box has been manually configured.')
            ztp_exit("systemd_exit")

        on_air()

        # Check for local ztp script file once at boot
        rc = ztp_local_event()
        if rc is True:
            if not allow_cascade:
                ztp_exit(True)
            write_wall_msg('Provisioning continue with the next method',
                           log_info)
            initialize_config()
            allow_cascade = False
        elif rc == 'script_failure':
            ztp_exit(rc)

        attempts = 3

        # Check for DHCP ztp events forever, unless the
        # machine is manually configured.
        while True:
            if check_initialized_config() is True:
                ztp_exit("systemd_exit")

            if os.path.isfile(ztp_dhcp):
                rc = ztp_dhcp_event()
                if rc is True:
                    ztp_exit(True)
                elif rc == 'script_failure':
                    # rebase the initial config so that ztp
                    # doesn't exit due to failed config
                    initialize_config()

            # if there is no url or execution failed, try again
            # starting with option 239 acquisition. Try front panel ports
            # as well
            if attempts == 3:
                if allow_front_panel:
                    enable_front_panel_ports(True)
                    # for now only renew leases from front panel ports
                    renew_dhcp_leases()
                attempts = 0
            else:
                attempts += 1
            time.sleep(60)

    if args.url is not None:
        if is_on_air():
            print('Another instance of ZTP is in progress.')
            sys.exit(1)

        if args.url.lower().startswith("."):
            print("Please, enter the complete file path")
            log_error("Incomplete path")
            sys.exit(1)

        on_air()
        if tryurl(args.url, True, "ZTP Manual") is True:
            ztp_exit(True)
        else:
            ztp_exit(False)


if __name__ == "__main__":
    sys.exit(main())
