#!/usr/bin/env python3

# ==============================================================================
# Copyright 2011 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import cfnbootstrap
from cfnbootstrap.sqs_client import SQSClient
from optparse import OptionParser
import os
import sys
from cfnbootstrap import util
import json

parser = OptionParser(usage="usage: %prog [options] [Command Result Data]")

parser.add_option_group(util.get_proxy_options(parser))

parser.add_option("", "--access-key", help="An AWS Access Key",
                  type="string", dest="access_key", default=os.environ.get('RESULT_ACCESS_KEY'))
parser.add_option("", "--secret-key", help="An AWS Secret Key",
                  type="string", dest="secret_key", default=os.environ.get('RESULT_SECRET_KEY'))
parser.add_option("", "--token", help="An AWS Session Token",
                  type="string", dest="security_token", default=os.environ.get('RESULT_SESSION_TOKEN'))
parser.add_option("-q", "--queue-url", help="SQS Queue URL for storing the command result",
                  type="string", dest="queue_url", default=os.environ.get('RESULT_QUEUE'))

parser.add_option("-d", "--dispatcher-id", help="The dispatcher ID",
                  type="string", dest="dispatcher_id", default=os.environ.get('DISPATCHER_ID'))
parser.add_option("-c", "--command-name", help="The command name",
                  type="string", dest="command_name", default=os.environ.get('CMD_NAME'))
parser.add_option("-i", "--invocation-id", help="The invocation ID",
                  type="string", dest="invocation_id", default=os.environ.get('INVOCATION_ID'))
parser.add_option("-l", "--listener-id", help="The listener ID",
                  type="string", dest="listener_id", default=os.environ.get('LISTENER_ID'))

# Optional arguments
parser.add_option("-s", "--success", help="If true, signal success; if false, signal failure. Default: true",
                  dest="success", action="store", type="choice", choices=["true", "false"], default="true")
parser.add_option("-e", "--exit-code",
                  help="Derive success or failure from specified exit code. Note: This takes precedence over the success flag",
                  dest="exit_code", type="int", action="store")

(options, args) = parser.parse_args()


def fail_on_unspecified(value, name):
    if not value:
        print("Error: You must specify %s" % name, file=sys.stderr)
        parser.print_help(sys.stderr)
        sys.exit(1)


fail_on_unspecified(options.dispatcher_id, "DispatcherId")
fail_on_unspecified(options.command_name, "CommandName")
fail_on_unspecified(options.invocation_id, "InvocationId")
fail_on_unspecified(options.listener_id, "ListenerId")
fail_on_unspecified(options.queue_url, "QueueUrl")

if (not options.access_key or not options.secret_key or not options.security_token):
    print("Error: You must specify a token and an access key/secret key pair", file=sys.stderr)
    parser.print_help(sys.stderr)
    exit(1)

if options.exit_code is None:
    success = options.success == 'true'
else:
    success = options.exit_code == 0

result_msg = {
    'DispatcherId': options.dispatcher_id,
    'CommandName': options.command_name,
    'ListenerId': options.listener_id,
    'InvocationId': options.invocation_id,
    'Data': ' '.join(args) if success else '',
    'Message': ' '.join(args) if not success else '',
    'Status': 'SUCCESS' if success else 'FAILURE',
}

cfnbootstrap.configureLogging("DEBUG")

try:
    from cfnbootstrap._ca_install import (
        ensure_ca_override_installed,
        CaOverrideError,
        format_install_failure,
        format_install_tolerated,
        is_benign_install_failure,
    )
    ensure_ca_override_installed()
except ImportError as _e:
    # Override CA subsystem not present in this build -- skip (commercial
    # Linux + commercial Windows are unaffected). MUST be first so the
    # CaOverrideError name is never evaluated after a failed import.
    import logging
    logging.getLogger(__name__).warning(
        "[CaOverrideInstall] override CA subsystem unavailable; "
        "skipping override CA install: %s", _e)
except CaOverrideError as _e:
    import logging
    if is_benign_install_failure(_e):
        # Transient/benign (concurrent_install / budget_exceeded / marker_write):
        # a peer process is installing, or the certs installed but only the
        # bookkeeping marker failed. Trust is (or will be) intact -- log and
        # continue rather than fail the boot (e.g. a first-boot cfn-init/cfn-hup
        # install-lock race must NOT roll back the stack).
        logging.getLogger(__name__).warning(format_install_tolerated(_e))
    else:
        # Real install failure on a host that HAS an override PEM: refuse to
        # fall back to system-store-only trust.
        logging.getLogger(__name__).error(format_install_failure(_e))
        sys.exit(1)
except Exception as _e:
    # Any other unexpected error: keep the commercial boot un-brickable.
    import logging
    logging.getLogger(__name__).warning("override CA install: %s", _e)

client = SQSClient(util.Credentials(options.access_key,
                                    options.secret_key,
                                    options.security_token),
                   proxyinfo=util.get_proxyinfo(options))
try:
    client.send_message(options.queue_url, json.dumps(result_msg))
except Exception as e:
    print("Error: Could not send command result: {}".format(
        str(e)), file=sys.stderr)
    sys.exit(1)
