1# Copyright 2015 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License"). You
4# may not use this file except in compliance with the License. A copy of
5# the License is located at
6#
7# http://aws.amazon.com/apache2.0/
8#
9# or in the "license" file accompanying this file. This file is
10# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
11# ANY KIND, either express or implied. See the License for the specific
12# language governing permissions and limitations under the License.
13import threading
14import contextlib
15import os
16import tempfile
17import sys
18import zipfile
19
20from s3transfer import S3Transfer
21
22from awscli.customizations.commands import BasicCommand
23from awscli.customizations.s3.utils import human_readable_size
24
25
26class UploadBuildCommand(BasicCommand):
27    NAME = 'upload-build'
28    DESCRIPTION = 'Upload a new build to AWS GameLift.'
29    ARG_TABLE = [
30        {'name': 'name', 'required': True,
31         'help_text': 'The name of the build'},
32        {'name': 'build-version', 'required': True,
33         'help_text': 'The version of the build'},
34        {'name': 'build-root', 'required': True,
35         'help_text':
36         'The path to the directory containing the build to upload'},
37        {'name': 'operating-system', 'required': False,
38         'help_text': 'The operating system the build runs on'}
39    ]
40
41    def _run_main(self, args, parsed_globals):
42        gamelift_client = self._session.create_client(
43            'gamelift', region_name=parsed_globals.region,
44            endpoint_url=parsed_globals.endpoint_url,
45            verify=parsed_globals.verify_ssl
46        )
47        # Validate a build directory
48        if not validate_directory(args.build_root):
49            sys.stderr.write(
50                'Fail to upload %s. '
51                'The build root directory is empty or does not exist.\n'
52                % (args.build_root)
53            )
54
55            return 255
56        # Create a build based on the operating system given.
57        create_build_kwargs = {
58            'Name': args.name,
59            'Version': args.build_version
60        }
61        if args.operating_system:
62            create_build_kwargs['OperatingSystem'] = args.operating_system
63
64        response = gamelift_client.create_build(**create_build_kwargs)
65        build_id = response['Build']['BuildId']
66
67        # Retrieve a set of credentials and the s3 bucket and key.
68        response = gamelift_client.request_upload_credentials(
69            BuildId=build_id)
70        upload_credentials = response['UploadCredentials']
71        bucket = response['StorageLocation']['Bucket']
72        key = response['StorageLocation']['Key']
73
74        # Create the S3 Client for uploading the build based on the
75        # credentials returned from creating the build.
76        access_key = upload_credentials['AccessKeyId']
77        secret_key = upload_credentials['SecretAccessKey']
78        session_token = upload_credentials['SessionToken']
79        s3_client = self._session.create_client(
80            's3', aws_access_key_id=access_key,
81            aws_secret_access_key=secret_key,
82            aws_session_token=session_token,
83            region_name=parsed_globals.region,
84            verify=parsed_globals.verify_ssl
85        )
86
87        s3_transfer_mgr = S3Transfer(s3_client)
88
89        try:
90            fd, temporary_zipfile = tempfile.mkstemp('%s.zip' % build_id)
91            zip_directory(temporary_zipfile, args.build_root)
92            s3_transfer_mgr.upload_file(
93                temporary_zipfile, bucket, key,
94                callback=ProgressPercentage(
95                    temporary_zipfile,
96                    label='Uploading ' + args.build_root + ':'
97                )
98            )
99        finally:
100            os.close(fd)
101            os.remove(temporary_zipfile)
102
103        sys.stdout.write(
104            'Successfully uploaded %s to AWS GameLift\n'
105            'Build ID: %s\n' % (args.build_root, build_id))
106
107        return 0
108
109
110def zip_directory(zipfile_name, source_root):
111    source_root = os.path.abspath(source_root)
112    with open(zipfile_name, 'wb') as f:
113        zip_file = zipfile.ZipFile(f, 'w', zipfile.ZIP_DEFLATED, True)
114        with contextlib.closing(zip_file) as zf:
115            for root, dirs, files in os.walk(source_root):
116                for filename in files:
117                    full_path = os.path.join(root, filename)
118                    relative_path = os.path.relpath(
119                        full_path, source_root)
120                    zf.write(full_path, relative_path)
121
122
123def validate_directory(source_root):
124    # For Python26 on Windows, passing an empty string equates to the
125    # current directory, which is not intended behavior.
126    if not source_root:
127        return False
128    # We walk the root because we want to validate there's at least one file
129    # that exists recursively from the root directory
130    for path, dirs, files in os.walk(source_root):
131        if files:
132            return True
133    return False
134
135
136# TODO: Remove this class once available to CLI from s3transfer
137# docstring.
138class ProgressPercentage(object):
139    def __init__(self, filename, label=None):
140        self._filename = filename
141        self._label = label
142        if self._label is None:
143            self._label = self._filename
144        self._size = float(os.path.getsize(filename))
145        self._seen_so_far = 0
146        self._lock = threading.Lock()
147
148    def __call__(self, bytes_amount):
149        with self._lock:
150            self._seen_so_far += bytes_amount
151            if self._size > 0:
152                percentage = (self._seen_so_far / self._size) * 100
153                sys.stdout.write(
154                    "\r%s  %s / %s  (%.2f%%)" % (
155                        self._label, human_readable_size(self._seen_so_far),
156                        human_readable_size(self._size), percentage
157                    )
158                )
159                sys.stdout.flush()
160