# Copyright 2015 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. A copy of

# the License is located at

#

# http://aws.amazon.com/apache2.0/

#

# or in the "license" file accompanying this file. This file 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 zipfile

import copy

from contextlib import closing

from botocore.vendored import six

from awscli.arguments import CustomArgument, CLIArgument

from awscli.customizations import utils

ERROR_MSG = (

"--zip-file must be a zip file with the fileb:// prefix.\n"

"Example usage: --zip-file fileb://path/to/file.zip")

ZIP_DOCSTRING = ('<p>The path to the zip file of the code you are uploading. '

'Example: fileb://code.zip</p>')

def register_lambda_create_function(cli):

cli.register('building-argument-table.lambda.create-function',

_extract_code_and_zip_file_arguments)

cli.register('building-argument-table.lambda.update-function-code',

_modify_zipfile_docstring)

cli.register('process-cli-arg.lambda.update-function-code',

validate_is_zip_file)

def validate_is_zip_file(cli_argument, value, **kwargs):

if cli_argument.name == 'zip-file':

_should_contain_zip_content(value)

def _extract_code_and_zip_file_arguments(session, argument_table, **kwargs):

argument_table['zip-file'] = ZipFileArgument(

'zip-file', help_text=ZIP_DOCSTRING, cli_type_name='blob')

code_argument = argument_table['code']

code_model = copy.deepcopy(code_argument.argument_model)

del code_model.members['ZipFile']

argument_table['code'] = CodeArgument(

name='code',

argument_model=code_model,

operation_model=code_argument._operation_model,

is_required=False,

event_emitter=session.get_component('event_emitter'),

serialized_name='Code'

)

def _modify_zipfile_docstring(session, argument_table, **kwargs):

if 'zip-file' in argument_table:

argument_table['zip-file'].documentation = ZIP_DOCSTRING

def _should_contain_zip_content(value):

if not isinstance(value, bytes):

# If it's not bytes it's basically impossible for

# this to be valid zip content, but we'll at least

# still try to load the contents as a zip file

# to be absolutely sure.

value = value.encode('utf-8')

fileobj = six.BytesIO(value)

try:

with closing(zipfile.ZipFile(fileobj)) as f:

f.infolist()

except zipfile.BadZipfile:

raise ValueError(ERROR_MSG)

class ZipFileArgument(CustomArgument):

def add_to_params(self, parameters, value):

if value is None:

return

_should_contain_zip_content(value)

zip_file_param = {'ZipFile': value}

if parameters.get('Code'):

parameters['Code'].update(zip_file_param)

else:

parameters['Code'] = zip_file_param

class CodeArgument(CLIArgument):

def add_to_params(self, parameters, value):

if value is None:

return

unpacked = self._unpack_argument(value)

if 'ZipFile' in unpacked:

raise ValueError("ZipFile cannot be provided "

"as part of the --code argument. "

"Please use the '--zip-file' "

"option instead to specify a zip file.")

if parameters.get('Code'):

parameters['Code'].update(unpacked)

else:

parameters['Code'] = unpacked