From 5abdd989fbf8f21f9fc01addaf45e661ec793c81 Mon Sep 17 00:00:00 2001 From: Jonathan Abrahams Date: Sat, 23 Jun 2018 22:21:23 -0400 Subject: SERVER-34996 Save console_log & console_screenshot in aws_ec2.py --- buildscripts/aws_ec2.py | 70 ++++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 63 insertions(+), 7 deletions(-) (limited to 'buildscripts/aws_ec2.py') diff --git a/buildscripts/aws_ec2.py b/buildscripts/aws_ec2.py index 8a92749d531..259948d85fe 100755 --- a/buildscripts/aws_ec2.py +++ b/buildscripts/aws_ec2.py @@ -3,6 +3,7 @@ from __future__ import print_function +import base64 import collections import datetime import optparse @@ -16,6 +17,24 @@ import yaml _MODES = ("status", "create", "start", "stop", "force-stop", "reboot", "terminate") +def write_binary_file(path, buffer): + """Write buffer to path in binary format.""" + with open(path, "wb") as fh: + fh.write(buffer) + + +def write_utf8_file(path, buffer): + """Write buffer to path in utf-8 format.""" + with open(path, "w") as fh: + fh.write(buffer.encode("utf-8")) + + +def write_yaml_file(path, dict): + """Write dict to path in YML format.""" + with open(path, "w") as ystream: + yaml.safe_dump(dict, ystream) + + class AwsEc2(object): """Class to support controlling AWS EC2 istances.""" @@ -74,7 +93,8 @@ class AwsEc2(object): sys.stdout.flush() return 0 if reached_state else 1 - def control_instance(self, mode, image_id, wait_time_secs=0, show_progress=False): + def control_instance(self, mode, image_id, wait_time_secs=0, show_progress=False, + console_output_file=None, console_screenshot_file=None): """Control an AMI instance. Returns 0 & status information, if successful.""" if mode not in _MODES: raise ValueError("Invalid mode '{}' specified, choose from {}.".format(mode, _MODES)) @@ -116,6 +136,29 @@ class AwsEc2(object): getattr(instance, "public_ip_address", None), getattr(instance, "private_dns_name", None), getattr(instance, "public_dns_name", None), getattr(instance, "tags", None)) + + if console_output_file: + try: + console_ouput = instance.console_output() + if console_ouput and "Output" in console_ouput: + write_utf8_file(console_output_file, console_ouput["Output"]) + else: + print("Unable to generate console_ouptut file, data not available") + except botocore.exceptions.ClientError as err: + print("Unable to generate console_ouptut file: {}".format(err.message)) + + if console_screenshot_file: + client = boto3.client("ec2") + try: + console_screenshot = client.get_console_screenshot(InstanceId=image_id) + if console_screenshot and "ImageData" in console_screenshot: + write_binary_file(console_screenshot_file, + base64.decodestring(console_screenshot["ImageData"])) + else: + print("Unable to generate console_screenshot file, data not available") + except botocore.exceptions.ClientError as err: + print("Unable to generate console_screenshot file: {}".format(err.message)) + except botocore.exceptions.ClientError as err: return 1, err.message @@ -141,7 +184,7 @@ class AwsEc2(object): def launch_instance( # pylint: disable=too-many-arguments,too-many-locals self, ami, instance_type, block_devices=None, key_name=None, security_group_ids=None, security_groups=None, subnet_id=None, tags=None, wait_time_secs=0, show_progress=False, - **kwargs): + console_output_file=None, console_screenshot_file=None, **kwargs): """Launch and tag an AMI instance. Return the tuple (0, status_information), if successful. @@ -179,7 +222,9 @@ class AwsEc2(object): self.tag_instance(instance.instance_id, tags) - return self.control_instance("status", instance.instance_id) + return self.control_instance("status", instance.instance_id, + console_output_file=console_output_file, + console_screenshot_file=console_screenshot_file) def main(): # pylint: disable=too-many-locals,too-many-statements @@ -250,6 +295,15 @@ def main(): # pylint: disable=too-many-locals,too-many-statements status_options.add_option("--yamlFile", dest="yaml_file", default=None, help="Save the status into the specified YAML file.") + status_options.add_option("--consoleOutputFile", dest="console_output_file", default=None, + help="Save the console output into the specified file, if" + " available.") + + status_options.add_option("--consoleScreenshotFile", dest="console_screenshot_file", + default=None, + help="Save the console screenshot (JPG format) into the specified" + " file, if available.") + parser.add_option_group(control_options) parser.add_option_group(create_options) parser.add_option_group(status_options) @@ -287,7 +341,9 @@ def main(): # pylint: disable=too-many-locals,too-many-statements ami=options.ami, instance_type=options.instance_type, block_devices=block_devices, key_name=options.key_name, security_group_ids=options.security_group_ids, security_groups=options.security_groups, subnet_id=options.subnet_id, tags=tags, - wait_time_secs=options.wait_time_secs, show_progress=True, **my_kwargs) + wait_time_secs=options.wait_time_secs, show_progress=True, + console_output_file=options.console_output_file, + console_screenshot_file=options.console_screenshot_file, **my_kwargs) else: if not getattr(options, "image_id", None): parser.print_help() @@ -295,7 +351,8 @@ def main(): # pylint: disable=too-many-locals,too-many-statements (ret_code, instance_status) = aws_ec2.control_instance( mode=options.mode, image_id=options.image_id, wait_time_secs=options.wait_time_secs, - show_progress=True) + show_progress=True, console_output_file=options.console_output_file, + console_screenshot_file=options.console_screenshot_file) if ret_code: print("Return code: {}, {}".format(ret_code, instance_status)) @@ -307,8 +364,7 @@ def main(): # pylint: disable=too-many-locals,too-many-statements if options.yaml_file: print("Saving status to {}".format(options.yaml_file)) - with open(options.yaml_file, "w") as ystream: - yaml.safe_dump(status_dict, ystream) + write_yaml_file(options.yaml_file, status_dict) print(yaml.safe_dump(status_dict)) -- cgit v1.2.1