#!/usr/bin/python
#
# (c) Copyright 2015 Hewlett Packard Enterprise Development LP.
#
"""
em_installer.py - EM installer for EM Smart Component

Run with -h or --help for details on cmd line args.
"""

import io
import pycurl
import os
import json
import sys
import traceback
import optparse
import re
import time
import subprocess


# exit codes, per fusion requirements
class ExitCode:
    em_flash_success         =   0 # Successful update
    no_em_found              = 101 # No EM found at given IP address
    user_terminated          = 120 # User terminated the script
    incorrect_parameters     = 140 # Invalid cmd line args
    login_failed             = 160 # Invalid credentials
    em_not_active            = 162 # Standby EM was targeted. Must be active.
    em_update_timeout        = 163 # Update did not complete
    em_not_resp_after_reboot = 164 # EM not responding after reboot
    em_internal_error        = 165 # EM failed to update
    internal_error           = 166 # This script had an internal error
    no_local_privileges      = 169 # This script lacked needed privileges
    update_in_progress       = 170 # Another update is already being run
    invalid_package          = 171 # Firmware package is invalid
    em_update_interrupted    = 172 # Update was interrupted on EM
    standby_em_failed        = 180 # Failed to update standby EM

    @staticmethod
    def contains(val):
        """
        Check if a given number is one of our exit codes.
        Return True or False.
        """

        for attr_name in ExitCode.__dict__:
            attr = getattr(ExitCode, attr_name)
            if (type(attr) is int and attr == val):
                return True
        return False


# RIS response codes
ris_rc_success = 200
ris_rc_created = 201
ris_rc_accepted = 202
ris_rc_bad_request = 400
ris_rc_unauthorized = 401
ris_rc_forbidden = 403
ris_rc_not_found = 404
ris_rc_resource_conflict = 409
ris_rc_server_error = 500

# EM error result messages
em_err_invalid_fw_pkg = 'emRegistry.1.0.InvalidFirmwarePackage'
em_err_update_interrupted = 'emRegistry.1.0.UpdateInterrupted'
em_err_standby_failed = 'emRegistry.1.0.StandbyUpdateFailed'

# Other global variables
EM_FWPKGUTIL = "/usr/bin/em-fwpkgutil"


# Making a pass-through io.BytesIO subclass only to enable some unit test mocks
class MyBytesIO(io.BytesIO):
    pass


class EmCurlFactory():
    """ Factory class for making curl objects for a given EM. """

    def __init__(self, ip_addr, user, password, updating_peer_em):
        """
        Construct the factory that will be used to make curl objects for the EM.
        Also gets list of links from RIS root, because we'll need to use those
        later to translate resources to URIs.
        If we end up needing to access resources not linked from the RIS root,
        we'll probably need to add general support for walking the RIS tree.
        """

        self.ip_addr = ip_addr
        self.user = user
        self.password = password
        self.logged_in = False
        self.use_peer_login = updating_peer_em

        #TODO: Later we will want to only use HTTPS. When we decide it's time
        # to retire http access for good, remove this and any associated code
        # where we use this flag to not use https.
        self.use_https = (user == "Administrator" or updating_peer_em)

        print("Getting RIS root..")
        curl = self.make_curl_obj("/rest/v1")
        (rc, data) = ris_perform(curl)
        if rc != ris_rc_success:
            if ip_addr[0] == "[" and ip_addr[-1] == "]":
                print("will retry without square brackets..")
                self.ip_addr = ip_addr[1:-1]
            else:
                print("will retry with square brackets..")
                self.ip_addr = "[" + ip_addr + "]"
            curl = self.make_curl_obj("/rest/v1")
            (rc, data) = ris_perform(curl)
            if rc != ris_rc_success:
                print("ERROR: Target is not an EM or is not responsive")
                exit(ExitCode.no_em_found)

        self.resources = {}
        hp_links = data["Oem"]["Hp"]["links"]
        for key in hp_links:
            self.resources[key] = hp_links[key]["href"].encode('ascii')
        # data["links"] is not yet populated on the standby EM, so don't
        # just assume it is there yet. It will be soon though!
        if "links" in data:
            links = data["links"]
            for key in links:
                self.resources[key] = links[key]["href"].encode('ascii')

        self.login()
        if not updating_peer_em:
            self._replace_floating_ip_addr()


    def _replace_floating_ip_addr(self):
        """
        Make sure we use the bay-static IP address, not the active floating IP
        address.
        """

        print("")
        print("Checking the EM static IP address..")

        curl = self.make_curl_obj("UpdateService")
        (rc, data) = ris_perform(curl)
        if rc != ris_rc_success:
            print("ERROR: Failed to get the UpdateService RIS resource.")
            exit(ExitCode.em_internal_error)

        em_ip_addr = str(data["EmIpAddress"])
        em_ip_addr = re.sub(r'[\da-fA-F:]+', em_ip_addr, self.ip_addr, 1)

        if self.ip_addr.lower() != em_ip_addr.lower():
            # We must have been given the floating IP address and will now use
            # the static IP address instead. There's an intermittent bug in the
            # CI manager where it can change the source IP address it claims to
            # send RIS requests from when we change which IP address we use
            # here. RIS requests on a session must use the same source IP
            # address as when it created the session, so that CI manager bug
            # was causing some updates to fail. We can work around that by
            # creating a new session here.
            print("")
            print("Will switch from floating to static IP address.")
            print("The IP address we were given:", self.ip_addr)
            print("The IP address we will use instead:", em_ip_addr)
            print("Logging out and creating a new session..")
            self.logout()
            self.ip_addr = em_ip_addr
            self.login()

    def login(self):
        """ Log in to the EM to start a RIS session. """

        self.logged_in = False
        if not self.use_https:
            return

        if self.use_peer_login:

            print("")
            print("Creating session with peer EM..")

            # There are some cases where this is used after a reboot, during which there may be a short time window
            # where Apache has been launched but the underlying cluster parameters have not yet been been
            # synchronized.  To handle those cases, retry peer RIS session creation for up to 2 minutes.

            # The session will live for 30 minutes, which should be plenty of time to do what we need to do (including retries).
            attempts = 0
            tryAgain = True
            while tryAgain:
                try:
                    output = subprocess.check_output(["create_peer_em_ris_session", "1800"])
                    session_id, token = re.findall('\S+', output)
                    # Success!
                    tryAgain = False
                except:
                    attempts += 1
                    if attempts < 24:
                        print("Session creation attempt #%d failed... waiting for retry" % attempts)
                        time.sleep(5)
                    else:
                        print("ERROR: unable to create session with the standby EM")
                        exit(ExitCode.login_failed)

            self.auth_token = token
            self.session_uri = self.resources["Sessions"] + "/" + session_id

        else:

            print("")
            print("Creating session..")

            curl = self.make_curl_obj("Sessions", http_header=['Content-Type: application/json'])
            curl.setopt(curl.POST, 1)
            data = json.dumps({"UserName":self.user, "Password":self.password})
            curl.setopt(curl.POSTFIELDS, data)
            curl.setopt(curl.HEADER, 1)

            # Some transient errors have been seen after reboots, such as "empty response".  We expect to
            # either get a "200 OK" when creating a session.  If we get anything else then retry (for up
            # to 2 minutes).

            attempts = 0
            while attempts < 24:
                (rc, resp) = ris_perform(curl, jsonify_resp=False)
                if rc == ris_rc_created:
                    break
                # Wait 5 seconds before another attempt.
                attempts += 1
                time.sleep(5)

            if rc != ris_rc_created:
                print("ERROR: Failed to create session with the given credentials")
                exit(ExitCode.login_failed)

            m = re.search(r"X-Auth-Token: (\S+)", resp)
            if m is None:
                print("ERROR: Failed to parse auth token from login response")
                exit(ExitCode.login_failed)
            self.auth_token = m.group(1)

            m = re.search(r"Location: (\S+)", resp)
            if m is None:
                print("ERROR: Failed to parse session location from login response")
                exit(ExitCode.login_failed)
            self.session_uri = m.group(1)

        self.logged_in = True


    def logout(self):
        """ Log out of the RIS session we've created. """

        if not self.logged_in:
            return

        print("")
        print("Logging out..")

        try:
            curl = self.make_curl_obj(self.session_uri)
            curl.setopt(curl.CUSTOMREQUEST, 'DELETE')

            (rc, resp) = ris_perform(curl, jsonify_resp=False)
            if rc != ris_rc_success:
                print("ERROR: Failed to delete session")
            else:
                self.logged_in = False
        except:
            print("ERROR: Failed to delete session due to exception..")
            traceback.print_exc(file=sys.stdout)


    def __del__(self):
        """ Destructor to clean up RIS session. """
        self.logout()


    def make_curl_obj(self, resource, timeout=60, http_header=None):
        """
        Return a curl object pointing to the given RIS resource.
        The RIS resource must be a full path OR the name of one of the links
        shown by the RIS root.
        """

        # Why not just default this param to the empty list? Because if I do
        # that, appending to it will change the list that the next call to
        # this function defaults to. WHY WOULD PYTHON DO THAT!?
        if http_header is None:
            http_header = []

        if resource[0] == "/":
            # We are using full path rather than a name to look up
            pass
        elif resource not in self.resources:
            print("ERROR: Unknown resource: ", resource)
            exit(ExitCode.internal_error)
        else:
            resource = self.resources[resource]

        curl = pycurl.Curl()
        curl.setopt(curl.TIMEOUT, timeout)

        #explictly enable TLS 1.x -- Works around bug in curl-7.19.7-46.el6 on CIM
        try:
            curl.setopt(curl.SSLVERSION, curl.SSLVERSION_TLSv1)
        except AttributeError as e:
            print("Error while setting curl.SSLVERSION_TLSv1:", e)
        except pycurl.error as (errno, errstr):
            print("curl error:", errno, errstr)

        # The following two lines are necessary due to using self-signed
        # certificates. If we change that such that this isn't necessary, then
        # we should remove this. It is the equivalent of using "--insecure"
        # from the curl command line
        curl.setopt(curl.SSL_VERIFYPEER, False)
        curl.setopt(curl.SSL_VERIFYHOST, False)
        if self.logged_in:
            http_header.append('X-Auth-Token: ' + self.auth_token)
        if len(http_header) > 0:
            curl.setopt(curl.HTTPHEADER, http_header)
        if self.use_https:
            curl.setopt(curl.URL, "https://" + self.ip_addr + resource)
        else:
            curl.setopt(curl.URL, "http://" + self.ip_addr + resource)

        return curl


def get_cmd_line_args(argv):
    """Parse, verify, & return cmd line args."""

    # Make OptionParser subclass to use the exit code we want
    class MyOptionParser(optparse.OptionParser):
        def error(self, msg):
            self.print_usage(sys.stderr)
            self.exit(ExitCode.incorrect_parameters, "%s: error: %s\n" % (self.get_prog_name(), msg))

    # construct usage and option definitions
    usage = "usage: %prog -e <ip_addr> -f <fw_package> -u <username> [-p <password>]"
    parser = MyOptionParser(usage=usage)
    parser.add_option("-f", "--file", dest="fw_pkg", help="EM firmware package filename", metavar="FILE")
    parser.add_option("-e", "--em", dest="em_ip_addr", help="Targeted EM IP address", metavar="IP_ADDR")
    parser.add_option("-u", "--user", dest="user", help="Username for EM", metavar="STRING")
    parser.add_option("-p", "--pass", dest="password", help="Password for EM", metavar="STRING")
    parser.add_option("-s", "--allow-standby", dest="allow_standby", help="Allow targeting the standby EM only", action="store_true")
    parser.add_option("--updating-peer-em", dest="updating_peer_em", help="I am an EM updating my peer. Don't bother creating session.", action="store_true")

    # parse the cmd line and sanity check it
    (opts, args) = parser.parse_args(argv)
    if not opts.em_ip_addr:
        parser.error('EM IP address not given')
    if not opts.fw_pkg:
        parser.error('Firmware package not given')
    if not os.path.exists(opts.fw_pkg):
        parser.error("Firmware package doesn't exist")
    if not opts.user:
        parser.error('Username not given')
    if not opts.password:
        import getpass
        opts.password = getpass.getpass()
    if opts.updating_peer_em:
        opts.allow_standby = True

    return opts


def verify_em_active(curl_factory):
    """
    Make sure we are targeting an active EM. Exit if standby.
    """

    print("")
    print("Verifying EM role..")

    curl = curl_factory.make_curl_obj("UpdateService")
    (rc, data) = ris_perform(curl)
    if rc != ris_rc_success:
        print("ERROR: Failed to get UpdateService RIS resource")
        exit(ExitCode.internal_error)
    if not data["EmIsActive"]:
        print("ERROR: Target EM is not active. Use '--updating-standby' to override.")
        exit(ExitCode.em_not_active)


def pkg_size(fw_pkg, updating_peer_em):
    """ Return package size. """
    if not updating_peer_em:
        return os.path.getsize(fw_pkg)
    else:
        cmd = [EM_FWPKGUTIL, "rawcatsize", fw_pkg]
        return int(subprocess.check_output(cmd))


def open_pkg(fw_pkg, updating_peer_em):
    """ Return file object to read firmware package contents. """
    if updating_peer_em:
        cmd = [EM_FWPKGUTIL, "rawcat", fw_pkg, "/dev/stdout"]
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE)
        return proc.stdout
    else:
        try:
            return open(fw_pkg, 'rb')
        except IOError:
            print("ERROR: Unable to open", fw_pkg)
            exit(ExitCode.no_local_privileges)


def post_update(curl_factory, fw_pkg, updating_peer_em):
    """
    POST an update request to an EM via RIS.
    Exits on error with appropriate ExitCode.
    """

    attempts = 1
    while True:

        print("")
        print("Posting update request..")

        # construct RIS message
        curl = curl_factory.make_curl_obj("UpdateService", timeout=300, http_header=['Content-Type: application/octet-stream'])
        curl.setopt(curl.POST, 1)
        curl.setopt(curl.POSTFIELDSIZE, pkg_size(fw_pkg, updating_peer_em))
        fin = open_pkg(fw_pkg, updating_peer_em)
        curl.setopt(curl.READFUNCTION, fin.read)

        # perform the request and interpret result
        (rc, data) = ris_perform(curl)
        if (rc == ris_rc_accepted):
            return
        elif (rc == ris_rc_resource_conflict):
            print("ERROR: Another update is already in progress")
            exit(ExitCode.update_in_progress)
        elif (rc == ris_rc_bad_request and data and data["ResultMessage"] == em_err_invalid_fw_pkg):
            print("ERROR: Invalid firmware package")
            exit(ExitCode.invalid_package)
        elif (rc == ris_rc_bad_request):
            print("ERROR: Update request failed")
            exit(ExitCode.em_internal_error)
        elif (rc == ris_rc_server_error):
            print("ERROR: Update request failed")
            exitCode = ExitCode.em_internal_error
        else:
            print("ERROR: Unexpected response to update request")
            exitCode = ExitCode.internal_error

        # For unexpected errors, try up to three times with a 5 second delay between retries.
        if attempts < 3:
            attempts += 1
            print("")
            print("Waiting to retry...")
            time.sleep(5)
        else:
            exit(exitCode)

def ris_perform(curl, jsonify_resp=True):
    """
    Perform RIS operation with prepared curl object.
    Returns (RIS response code, RIS JSON response data).
    """

    # use BytesIO to capture curl response data
    fout = MyBytesIO()
    curl.setopt(curl.WRITEFUNCTION, fout.write)

    # perform the curl request
    try:
        curl.perform()
    except pycurl.error as (errno, errstr):
        print("curl error:", errno, errstr)
        return (None, None)

    # check response code
    rc = curl.getinfo(curl.RESPONSE_CODE)
    print("response code:", rc)
    if (rc == ris_rc_accepted):
        # No data is returned with this response code
        return (rc, None)

    # check response
    resp_buf = fout.getvalue()

    if not jsonify_resp:
        print(resp_buf)
        return (rc, resp_buf)

    # jsonify the response
    json_data = None
    try:
        json_data = json.loads(resp_buf)
        print("data:", json_data)
    except ValueError as error:
        print(resp_buf)
        print(error)

    return (rc, json_data)


def poll_em(curl_factory, url, max_wait_time, sleep_time, break_condition):
    """
    Poll EM, by calling GET on update URI, for a break condition to be met.
    Returns (Boolean if break condition was met, RIS response code, JSON data).
    """

    start_time = time.time()
    while (time.time() - start_time < max_wait_time):
        time.sleep(sleep_time)
        curl = curl_factory.make_curl_obj(url)
        (rc, data) = ris_perform(curl)
        if break_condition(rc, data):
            return (True, rc, data)
    return (False, None, None)


def wait_for_update_complete(curl_factory):
    """
    Wait for update to complete.
    Exits with appropriate ExitCode on error.
    """

    print("")
    print("Waiting for update to complete..")
    def break_condition(rc, data): return ((data and data["Successful"] is not None) or rc == ris_rc_unauthorized)
    (completed, rc, data) = poll_em(curl_factory, "UpdateService", 900, 15, break_condition)
    if not completed:
        print("ERROR: Timed out waiting for update to complete.")
        exit(ExitCode.em_update_timeout)
    elif data and data["ResultMessage"] == em_err_update_interrupted:
        print("ERROR: Update interrupted")
        exit(ExitCode.em_update_interrupted)
    elif data and data["ResultMessage"] == em_err_standby_failed:
        print("ERROR: Failed to update the standby EM, update aborted.")
        exit(ExitCode.standby_em_failed)
    elif data and not data["Successful"]:
        print("ERROR: Update failed")
        exit(ExitCode.em_internal_error)

    # It's possible, but very unlikely, that the EM has already booted back up
    # when we get to this point.
    # If the EM still needs to shutdown, wait for it to do so.
    if data and data["RebootPending"]:
        wait_for_em_shutdown(curl_factory)

    verify_update_status_after_reboot(curl_factory)


def wait_for_em_shutdown(curl_factory):
    """
    Wait for EM to shut down for reboot.
    Exits with appropriate ExitCode on error.
    """

    # wait for EM to stop responding
    print("")
    print("Waiting for EM to shutdown for reboot..")
    def break_condition(rc, data): return (rc is None)
    (completed, rc, data) = poll_em(curl_factory, "UpdateService", 600, 5, break_condition)
    if not completed:
        print("ERROR: Did not see EM go down for reboot")
        exit(ExitCode.em_internal_error)


def verify_update_status_after_reboot(curl_factory):
    """
    After EM has rebooted, log back in and verify update status reports success.
    """
    wait_for_ris_root(curl_factory)
    wait_for_login(curl_factory)
    wait_for_final_update_status(curl_factory)


def wait_for_ris_root(curl_factory):
    print("")
    print("Waiting for RIS root to be available..")

    curl_factory.logged_in = False
    def break_condition(rc, data): return (rc == ris_rc_success)
    (completed, rc, data) = poll_em(curl_factory, "/rest/v1", 600, 5, break_condition)
    if not completed:
        print("ERROR: EM or RIS not available after reboot")
        exit(ExitCode.em_not_resp_after_reboot)


def wait_for_login(curl_factory):
    start_time = time.time()
    while True:
        try:
            curl_factory.login()
            break
        except:
            if time.time() - start_time < 150:
                print("ERROR: Unable to login after reboot")
                exit(ExitCode.em_internal_error)
            else:
                print("Unable to login yet, will retry...")
        time.sleep(5)


def wait_for_final_update_status(curl_factory):
    print("")
    print("Check final update status..")

    def break_condition(rc, data): return (rc is not None)
    (completed, rc, data) = poll_em(curl_factory, "UpdateService", 120, 5, break_condition)
    if not completed or data is None:
        print("ERROR: Failed to get update status after EM reboot")
        exit(ExitCode.em_internal_error)
    if data["ResultMessage"] == em_err_update_interrupted:
        print("ERROR: Update interrupted")
        exit(ExitCode.em_update_interrupted)
    elif not data["Successful"]:
        print("ERROR: Update failed")
        exit(ExitCode.em_internal_error)


def main():
    try:
        opts = get_cmd_line_args(sys.argv)
        curl_factory = EmCurlFactory(opts.em_ip_addr, opts.user, opts.password, opts.updating_peer_em)
        if not opts.allow_standby:
            verify_em_active(curl_factory)
        post_update(curl_factory, opts.fw_pkg, opts.updating_peer_em)
        wait_for_update_complete(curl_factory)
        print("")
        print("Success!")
        exit(ExitCode.em_flash_success)
    except KeyboardInterrupt:
        print("Terminated by user")
        exit(ExitCode.user_terminated)
    except SystemExit as e:
        if (not ExitCode.contains(e.code)):
            print("ERROR: Unknown exit code attempted:", e.code)
            e.code = ExitCode.internal_error
        raise e
    except:
        print("ERROR: Unhandled exception")
        traceback.print_exc(file=sys.stdout)
        exit(ExitCode.internal_error)


if __name__ == "__main__":
    main()
