Skip to content

Commit

Permalink
Refactor _file_transport_command
Browse files Browse the repository at this point in the history
  • Loading branch information
@tremble
tremble committed Jan 19, 2023
1 parent 9733031 commit 851d41f
Showing 1 changed file with 60 additions and 33 deletions.
93 changes: 60 additions & 33 deletions plugins/connection/aws_ssm.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -647,7 +647,7 @@ def _flush_stderr(self, session_process):

return stderr

def _get_url(self, client_method, bucket_name, out_path, http_method, profile_name, extra_args=None):
def _get_url(self, client_method, bucket_name, out_path, http_method, extra_args=None):
''' Generate URL for get_object / put_object '''

region_name = self.get_option('region') or 'us-east-1'
Expand Down Expand Up @@ -695,54 +695,81 @@ def _get_boto_client(self, service, region_name=None, profile_name=None):
)
return client

@_ssm_retry
def _file_transport_command(self, in_path, out_path, ssm_action):
''' transfer a file from using an intermediate S3 bucket '''
def _escape_path(self, path):
return path.replace("\\", "/")

def _generate_encryption_settings(self):
put_args = {}
put_headers = {}
if not self.get_option('bucket_sse_mode'):
return put_args, put_headers

path_unescaped = u"{0}/{1}".format(self.instance_id, out_path)
s3_path = path_unescaped.replace('\\', '/')
bucket_url = 's3://%s/%s' % (self.get_option('bucket_name'), s3_path)
put_args['ServerSideEncryption'] = self.get_option('bucket_sse_mode')
put_headers['x-amz-server-side-encryption'] = self.get_option('bucket_sse_mode')
if self.get_option('bucket_sse_mode') == 'aws:kms' and self.get_option('bucket_sse_kms_key_id'):
put_args['SSEKMSKeyId'] = self.get_option('bucket_sse_kms_key_id')
put_headers['x-amz-server-side-encryption-aws-kms-key-id'] = self.get_option('bucket_sse_kms_key_id')
return put_args, put_headers

profile_name = self.get_option('profile')
def _generate_commands(self, bucket_name, s3_path, in_path, out_path):
put_args, put_headers = self._generate_encryption_settings()

put_args = dict()
put_headers = dict()
if self.get_option('bucket_sse_mode'):
put_args['ServerSideEncryption'] = self.get_option('bucket_sse_mode')
put_headers['x-amz-server-side-encryption'] = self.get_option('bucket_sse_mode')
if self.get_option('bucket_sse_mode') == 'aws:kms' and self.get_option('bucket_sse_kms_key_id'):
put_args['SSEKMSKeyId'] = self.get_option('bucket_sse_kms_key_id')
put_headers['x-amz-server-side-encryption-aws-kms-key-id'] = self.get_option('bucket_sse_kms_key_id')
put_url = self._get_url('put_object', bucket_name, s3_path, 'PUT', extra_args=put_args)
get_url = self._get_url('get_object', bucket_name, s3_path, 'GET')

if self.is_windows:
put_command_headers = "; ".join(["'%s' = '%s'" % (h, v) for h, v in put_headers.items()])
put_command = "Invoke-WebRequest -Method PUT -Headers @{%s} -InFile '%s' -Uri '%s' -UseBasicParsing" % (
put_command_headers, in_path,
self._get_url('put_object', self.get_option('bucket_name'), s3_path, 'PUT', profile_name,
extra_args=put_args))
get_command = "Invoke-WebRequest '%s' -OutFile '%s'" % (
self._get_url('get_object', self.get_option('bucket_name'), s3_path, 'GET', profile_name), out_path)
put_command_headers = "; ".join([f"'{h}' = '{v}'" for h, v in put_headers.items()])
put_command = (
"Invoke-WebRequest -Method PUT "
f"-Headers @{{{put_command_headers}}} " # @{'key' = 'value'; 'key2' = 'value2'}
f"-InFile '{in_path}' "
f"-Uri '{put_url}' "
f"-UseBasicParsing"
)
get_command = (
"Invoke-WebRequest "
f"'{get_url}' "
f"-OutFile '{out_path}'"
)
else:
put_command_headers = "".join(["-H '%s: %s' " % (h, v) for h, v in put_headers.items()])
put_command = "curl --request PUT %s--upload-file '%s' '%s'" % (
put_command_headers, in_path,
self._get_url('put_object', self.get_option('bucket_name'), s3_path, 'PUT', profile_name,
extra_args=put_args))
get_command = "curl '%s' -o '%s'" % (
self._get_url('get_object', self.get_option('bucket_name'), s3_path, 'GET', profile_name), out_path)
put_command_headers = " ".join([f"-H '{h}: {v}'" for h, v in put_headers.items()])
put_command = (
"curl --request PUT "
f"{put_command_headers} "
f"--upload-file '{in_path}' "
f"'{put_url}'"
)
get_command = (
"curl "
f"-o '{out_path}' "
f"'{get_url}'"
)

return get_command, put_command, put_args

@_ssm_retry
def _file_transport_command(self, in_path, out_path, ssm_action):
''' transfer a file from using an intermediate S3 bucket '''

bucket_name = self.get_option("bucket_name")
s3_path = self._escape_path(f"{self.instance_id}/{out_path}")

get_command, put_command, put_args = self._generate_commands(
bucket_name, s3_path, in_path, out_path,
)

client = self._get_boto_client('s3', profile_name=profile_name)
if ssm_action == 'get':
(returncode, stdout, stderr) = self.exec_command(put_command, in_data=None, sudoable=False)
with open(to_bytes(out_path, errors='surrogate_or_strict'), 'wb') as data:
client.download_fileobj(self.get_option('bucket_name'), s3_path, data)
client.download_fileobj(bucket_name, s3_path, data)
else:
with open(to_bytes(in_path, errors='surrogate_or_strict'), 'rb') as data:
client.upload_fileobj(data, self.get_option('bucket_name'), s3_path, ExtraArgs=put_args)
client.upload_fileobj(data, bucket_name, s3_path, ExtraArgs=put_args)
(returncode, stdout, stderr) = self.exec_command(get_command, in_data=None, sudoable=False)

# Remove the files from the bucket after they've been transferred
client.delete_object(Bucket=self.get_option('bucket_name'), Key=s3_path)
client.delete_object(Bucket=bucket_name, Key=s3_path)

# Check the return code
if returncode == 0:
Expand Down

0 comments on commit 851d41f

Please sign in to comment.