#!/usr/bin/python
#------------------------------------------------------------------------------
#
# 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


### Variables

DEVNULL = open(os.devnull, 'wb')
date = datetime.datetime
license_check = "/var/lib/cumulus/ztp/ztp_license_check"
log_handler = None
MOUNT_POINT = None
osrelease_file = '/etc/os-release'
save_script = True
script_date = date.utcnow().strftime('%Y%m%d%H%M')
script_flag = 'CUMULUS-AUTOPROVISIONING'
state_dir = "/var/lib/cumulus/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/cumulus/ztp/ztp_script-" + script_date
ztp_lock = "/var/run/ztp.lock"
sha_file = "/var/lib/cumulus/ztp/ztp_state.sha"
ztp_sha_lock = "/var/run/ztp_sha.lock"
cleanup_signal = None

### 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))


### 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


### 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:
        pass

    return None

def license_installed():
    cmd = ['/usr/sbin/switchd', '-lic']
    try:
        ret = subprocess.call(cmd, stdout=DEVNULL, stderr=DEVNULL)
    except:
        return 1

    if ret == 0 or ret == 99:
        return 0
    else:
        return 1


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 add_cumulus_headers(req):
    mfgr = "unknown"
    prodname = "unknown"
    serial = "unknown"

    for intf in ['eth0','eth1']:
        mac = mgmt_mac_address(intf)
        if mac:
            break

    if mac is None:
        mac = 'unknown'

    cmd = [ '/usr/cumulus/bin/decode-syseeprom', '--json' ]
    try:
        output = subprocess.check_output(cmd).strip()
    except subprocess.CalledProcessError:
        log_error("Problems reading EEPROM")
        output = '{}'

    try:
        eeprom = json.loads(output)
    except ValueError:
        log_error("The json from '%s' is invalid%s" % (' '.join(cmd), output))
        eeprom = {}

    if 'tlv' in eeprom:
        tlvs = eeprom['tlv']
        if 'Vendor Name' in tlvs:
            mfgr = tlvs['Vendor Name']['value']
        elif 'Manufacturer' in tlvs:
            mfgr = tlvs['Manufacturer']['value']

        if 'Product Name' in tlvs:
            prodname = tlvs['Product Name']['value']
        elif 'Part Number' in tlvs:
            prodname = tlvs['Part Number']['value']

        if 'Serial Number' in tlvs:
            serial = tlvs['Serial Number']['value']
        elif 'Service Tag' in tlvs:
            serial = tlvs['Service Tag']['value']


    # assemble headers to send the server
    pfx = "CUMULUS"
    req.add_header('User-agent', "CumulusLinux-AutoProvision/%s" % (version))
    req.add_header(pfx + '-ARCH', uname_arch())
    req.add_header(pfx + '-BUILD', osrelease_val("VERSION"))
    req.add_header(pfx + '-LICENSE-INSTALLED', license_installed())
    req.add_header(pfx + '-MANUFACTURER', mfgr)
    req.add_header(pfx + '-PRODUCTNAME', prodname)
    req.add_header(pfx + '-SERIAL', serial)
    req.add_header(pfx + '-MGMT-MAC', mac)
    req.add_header(pfx + '-VERSION', osrelease_val("VERSION_ID"))
    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
    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_json_state():
    print_state("json")


def print_state(json_arg):
    lock = shared_lock(ztp_lock)
    current_state = systemctl_call("is-enabled")
    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

            ztp_info = [['State', current_state], ['Version', ''],
                        ['Result', ''], ['Date', ''],
                        ['Method', ''], ['URL', '']]

            if state.has_option("Most Recent", "ZTP"):
                ztp_info[1][1] = state.get("Most Recent", "VERSION")
                ztp_info[2][1] = state.get("Most Recent", "ZTP")
                ztp_info[3][1] = state.get("Most Recent", "DATE")
                ztp_info[4][1] = state.get("Most Recent", "METHOD")
                ztp_info[5][1] = state.get("Most Recent", "URL")

            if json_arg is not None:
                json_output = []

                if state.has_section("Most Recent"):
                    json_output.append({ztp_info[0][0]: ztp_info[0][1],
                                        ztp_info[1][0]: ztp_info[1][1],
                                        ztp_info[2][0]: ztp_info[2][1],
                                        ztp_info[3][0]: ztp_info[3][1],
                                        ztp_info[4][0]: ztp_info[4][1],
                                        ztp_info[5][0]: ztp_info[5][1]})
                else:
                    json_output.append({ztp_info[0][0]: ztp_info[0][1]})
                print json.dumps(json_output, sort_keys=False, indent=4)
            else:
                print "\nZTP INFO:\n"
                col_width = max(len(word) for row in ztp_info
                                        for word in row) + 2
                for row in ztp_info:
                    print "".join(word.ljust(col_width) for word in row)
    else:
        ztp_info = [['State', current_state], ['Version', version]]

        if json_arg is not None:
            json_output = []
            json_output.append({ztp_info[0][0]: ztp_info[0][1],
                                ztp_info[1][0]: ztp_info[1][1]})
            print json.dumps(json_output, sort_keys=False, indent=4)
        else:
            print "\nZTP INFO:\n"
            col_width = max(len(word) for row in ztp_info
                                        for word in row) + 2
            for row in ztp_info:
                print "".join(word.ljust(col_width) for word in row)

    release_lock(lock)


def write_state(ztp_return, method, url):
    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"):
                if ztp_return != "Script Failure":
                    if systemctl_call("disable") is False:
                        ztp_exit(False)

                status = systemctl_call("is-enabled")
                state.set("STATUS", "ZTP", status)
                state.set("Most Recent", "ZTP", ztp_return)
                state.set("Most Recent", "VERSION", version)
                state.set("Most Recent", "DATE",
                        date.utcnow().strftime("%c UTC"))
                state.set("Most Recent", "METHOD", method)
                state.set("Most Recent", "URL", url)
                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")
        ztp_exit("missing_state")
    if success is False:
        log_error("Could not write in the State File")
        ztp_exit(False)
    return success


def check_license():
    if not os.path.isfile(license_check):
        with open(license_check, 'wb') as FILE:
            FILE.write("")

    if license_installed() == 0:
        with open(license_check, 'wb') as FILE:
            FILE.write("license installed = yes")
    else:
        with open(license_check, 'wb') as FILE:
            FILE.write("license installed = no")


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.
    '''

    passwd_file = "/etc/passwd"
    if_file = "/etc/network/interfaces"
    group_file = "/etc/group"
    shadow_file = "/etc/shadow"
    dpkg_file = "/var/log/dpkg.log"

    return [passwd_file, if_file, group_file, shadow_file,
            dpkg_file, license_check]

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()

    check_license()
    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)
    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)
    check_license()
    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):
        if shaconf.get("SHASUM", key) != value:
            log_warn("Switch has already been configured. "
                     + key + " was modified")
            log_warn("ZTP will not run")
            lock = exclusive_lock(ztp_lock)
            write_state("failed", "Switch 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 waterfall_search(method, directory, partition=None):
    global waterfall
    if waterfall is False:
        platform = subprocess.check_output(['/usr/bin/platform-detect',
                                            '--all'])
        platform_fields = platform.replace(',', ' ').split()
        vendor = platform_fields[0]
        model = platform_fields[1]
        revision = platform_fields[2]
        arch = subprocess.check_output(['/bin/uname', '-m']).rstrip()

        waterfall_parts = ['cumulus-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_info(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 ztp_usb_event():
    log_debug("ZTP USB: Looking for unmounted USB devices")
    return parse_partitions()


def ztp_dhcp_event():
    log_debug("ZTP DHCP: Looking for ZTP Script provided by DHCP")
    try:
        fd = os.open(ztp_dhcp, os.O_RDWR)
        url = os.read(fd, 1024).replace("\n", "")
    except IOError:
        log_error("Could not open %s" % ztp_dhcp)
        return False
    os.close(fd)
    return tryurl(url, True, "ZTP DHCP")


def tryurl(url, doexec, method, partition=None):
    success = False
    provision_msg = ('Attempting to provision via %s from %s'
                    % (method, url))

    log_info(provision_msg)
    wall_msg = 'ZTP: ' + provision_msg
    w = subprocess.Popen('/usr/bin/wall', stdin=subprocess.PIPE,
                         stderr=subprocess.PIPE)
    w.communicate(input=wall_msg)
    w.stdin.close()
    w.stderr.close()
    w.wait()
    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_cumulus_headers(req)
        for retry in range(1, 4):
            try:
                resp = urllib2.urlopen(req)
            except urllib2.URLError, error:
                log_error(method + ': URL Error %s' % error)
            except urllib2.HTTPError, error:
                log_error(method + ': HTTP Error %s' % error.code)
            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 :
                    break
            if (retry != 4):
                time.sleep(5 * retry)
                log_info(method + ': Retrying ')

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

    if success is True:
        write_state("success", method, url)
    elif success is False:
        write_state("failure", method, url)
    else:
        log_error("Script returned failure")
        write_state("Script Failure", method, url)
        ztp_exit("script_failure")
    release_lock(lock)
    if success is False:
        ztp_exit(False)
    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')
    except OSError:
        log_info('Errors while checking ZTP script:' + file)
        if (ofd != -1):
                os.close(ofd)


def processweb(url, scriptcontent, doexec, method, partition=None):
    success = False
    if scriptcontent.find(script_flag) > -1:
        log_info(method + ": Found Marker CUMULUS-AUTOPROVISIONING")
        log_info(method + ": Executing " + url)
        with str_to_tmp_file(scriptcontent) as script_file:
            os.chmod(script_file, 0o700)
            if save_script is True:
                shutil.copyfile(script_file, ztp_script)
            ztp_env = os.environ.copy()
            ztp_env['ZTP_URL'] = url
            ztp_env['PATH'] += (':/usr/cumulus/bin')

            if method == "ZTP USB":
                ztp_env['ZTP_USB_MOUNTPOINT'] = MOUNT_POINT
                mount_status = subprocess.check_call(["/bin/mount", '-o',
                                    'remount,rw', partition, MOUNT_POINT])
                if mount_status > 0:
                    log_error("ZTP USB: Mount failed with exit code "
                                                                + mount_status)
                    return False

            check_script(script_file)

            if method == "ZTP LOCAL" or method == "ZTP USB":
                # 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:
                log_error(method + ": Payload returned code %s"
                                            % cpe.returncode)
            except OSError, ose:
                if ose.errno == 2:
                    log_error(method + ": Could not find referenced "
                            "script/interpreter in downloaded payload.")
                elif ose.errno == 8:
                    log_error(method + ": Could not find interpreter "
                                        "line (#!<path_to_interpreter>) "
                                        "in downloaded payload.")
                else:
                    log_error(method + ': Unexpected OS error: %s' % str(ose))
            except SystemExit, se:
                if se.code == 0 and cleanup_signal == signal.SIGTERM:
                    log_info(method + ': user ztp script initiated system shut down.')
                    success = True
                else:
                    log_error(method + ': Unexpected SystemExit code: %s' % str(se))
                    raise
            except BaseException, be:
                log_error(method + ': Unexpected error: %s' % str(be))
            else:
                log_info(method + ": Script returned success")
                success = True
    else:
        # code ok, but no markers
        log_error(method + ": No marker '%s' found" % script_flag)
    return success


def parse_partitions():
    """
    lsblk gives us physical device, removable bool & current mounted location
    """
    global MOUNT_POINT
    log_debug("ZTP USB: Parsing partitions")
    supported_fstypes = ['vfat', 'ext2', 'ext3', 'ext4']
    partitions = subprocess.check_output(['/bin/lsblk', '--noheadings',
                                    '--pairs', '--ascii', '--output',
                                    'NAME,UUID,MOUNTPOINT,SIZE,RM,FSTYPE'])
    for line in partitions.splitlines():
        line = line.strip()
        line_fields = line.split()

        '''
        Extract values from lsblk, simple string removal - no need for regex's
        '''
        blk_devname = line_fields[0].replace('NAME="',
                                                    '/dev/').replace('"', '')
        blk_mountpoint = str(line_fields[2]).replace('MOUNTPOINT="',
                                                        '').replace('"', '')
        blk_fstype = line_fields[5].replace('FSTYPE="', '').replace('"', '')

        '''
        Only find removable partitions with a UUID (avoids root device name)
        '''
        if not line or line.endswith('RM="0"') or 'UUID=""' in line:
            continue
        if len(blk_mountpoint) > 0:
            #log_warn('ZTP USB: Already mounted as ' + blk_mountpoint)
            continue
        if blk_fstype not in supported_fstypes:
            #log_warn('ZTP USB: Unsupported partition type found ' + blk_fstype)
            continue
        else:
            log_info("ZTP USB: Unmounted device found: " + blk_devname)
            MOUNT_POINT = tempfile.mkdtemp()
            if mount_and_exec(blk_devname) is True:
                unmount_partition()
                return True
            unmount_partition()
    log_debug("ZTP USB: Device not found")
    return False


def mount_and_exec(partition):
    if os.path.isdir(MOUNT_POINT) is False:
        log_error("ZTP USB: Could not create mount point directory: \"%s\""
                                                            % MOUNT_POINT)
        return False

    mount_status = subprocess.check_call(["/bin/mount", '-o', 'ro',
                                            partition, MOUNT_POINT])

    if mount_status > 0:
        log_error("ZTP USB: Mount failed with exit code " + mount_status)
        return False
    else:
        return waterfall_search("ZTP USB", MOUNT_POINT, partition)


def unmount_partition():
    global MOUNT_POINT
    if MOUNT_POINT:
        umount_status = subprocess.check_call(["/bin/umount",
                                                "-l", MOUNT_POINT])
        if umount_status > 0:
            log_warn('ZTP USB: Unmount failed with exit code ' + umount_status)
        else:
            log_info('ZTP USB: Unmount successful.')

        shutil.rmtree(MOUNT_POINT)
        if os.path.exists(MOUNT_POINT) is True:
            log_warn('ZTP USB: Unable to remove temporary mount point '
                                        'directory: \"%s\"' % MOUNT_POINT)
        MOUNT_POINT = None


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

def ztp_exit(success):
    if MOUNT_POINT:
        unmount_partition()
    if success is True:
        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)
    else:
        log_error("ZTP failed. Exiting and disabling ZTP...")
        systemctl_call("disable")
        sys.exit(1)

# mgmt vrf enabled by default in 4.0, put ourselves in mgmt vrf
def ztp_setup_vrf():
    try:
        mypid = str(os.getpid())
        cmd = ['/usr/bin/vrf', 'task', 'set', 'mgmt', mypid]
        subprocess.check_call(cmd, stdout=DEVNULL, stderr=DEVNULL)
    except:
        pass

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

    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")
    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_info("Found ZTP DHCP Request")
        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:
            ztp_exit(True)
        else:
            sys.exit(1)

    if args.enable is True:
        if systemctl_call("enable") is True:
            ztp_exit(True)
        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)
            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:
        # Check if the box has been manually configured
        if check_initialized_config() is True:
            ztp_exit("systemd_exit")

	ztp_setup_vrf()

        # Check for local ztp script file once at boot
        if ztp_local_event() is True:
            ztp_exit(True)

        # Check for USB and DHCP ztp events forever, unless the
        # machine is manually configured.
        while True:
            if check_initialized_config() is True:
                ztp_exit("systemd_exit")
            if ztp_usb_event() is True:
                ztp_exit(True)
            if os.path.isfile(ztp_dhcp):
                if ztp_dhcp_event() is True:
                    ztp_exit(True)
            time.sleep(60)

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

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


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