SERVER-34996 Save console_log & console_screenshot in aws_ec2.py

1 files changed, 63 insertions, 7 deletions

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