diff --git a/sos/policies/distros/__init__.py b/sos/policies/distros/__init__.py index 1ef9c4c7a6..4d820d7214 100644 --- a/sos/policies/distros/__init__.py +++ b/sos/policies/distros/__init__.py @@ -62,12 +62,12 @@ class LinuxPolicy(Policy): _upload_user = None _upload_password = None _upload_method = None - _upload_s3_endpoint = None + _upload_s3_endpoint = 'https://s3.amazonaws.com' _upload_s3_bucket = None _upload_s3_access_key = None _upload_s3_secret_key = None _upload_s3_region = None - _upload_s3_object_prefix = None + _upload_s3_object_prefix = '' default_container_runtime = 'docker' _preferred_hash_name = None upload_url = None @@ -339,9 +339,15 @@ def pre_work(self): cmdline_opts.quiet: try: # Policies will need to handle the prompts for user information - if cmdline_opts.upload and self.get_upload_url(): + if cmdline_opts.upload and self.get_upload_url() and \ + not cmdline_opts.upload_protocol == 's3': self.prompt_for_upload_user() self.prompt_for_upload_password() + elif cmdline_opts.upload_protocol == 's3': + self.prompt_for_upload_s3_bucket() + self.prompt_for_upload_s3_endpoint() + self.prompt_for_upload_s3_access_key() + self.prompt_for_upload_s3_secret_key() self.ui_log.info('') except KeyboardInterrupt: raise @@ -384,32 +390,48 @@ def prompt_for_upload_s3_access_key(self): to be provided for upload or not """ if not self.get_upload_s3_access_key(): - msg = "Please provide the upload access key for %s: " + + msg = ("Please provide the upload access key for bucket %s via " + "endpoint %s: " % (self.get_upload_s3_bucket(), + self.get_upload_s3_endpoint())) self.upload_s3_access_key = input(_(msg)) + def prompt_for_upload_s3_secret_key(self): + """Should be overridden by policies to determine if a secret key needs + to be provided for upload or not + """ + if not self.get_upload_s3_secret_key(): + msg = ("Please provide the upload secret key for bucket %s via " + "endpoint %s: " % (self.get_upload_s3_bucket(), + self.get_upload_s3_endpoint())) + self.upload_s3_secret_key = getpass(msg) + def prompt_for_upload_s3_bucket(self): """Should be overridden by policies to determine if a bucket needs to be provided for upload or not """ - if not self.get_upload_s3_bucket(): - msg = "Please provide the upload bucket for %s: " - self.upload_s3_bucket = input(_(msg)) + if not self.upload_s3_bucket: + if self.upload_url and self.upload_url.startswith('s3://'): + self.upload_s3_bucket = self.upload_url[5:] + else: + user_input = input(_("Please provide the upload bucket: ")) + self.upload_s3_bucket = user_input.strip('/') + return self.upload_s3_bucket def prompt_for_upload_s3_endpoint(self): """Should be overridden by policies to determine if an endpoint needs to be provided for upload or not """ - if not self.get_upload_s3_endpoint(): - msg = "Please provide the upload endpoint for %s: " - self.upload_s3_endpoint = input(_(msg)) - - def prompt_for_upload_s3_secret_key(self): - """Should be overridden by policies to determine if a secret key needs - to be provided for upload or not - """ - if not self.get_upload_s3_secret_key(): - msg = "Please provide the upload secret key for %s: " - self.upload_s3_secret_key = getpass(msg) + default_endpoint = self._upload_s3_endpoint + if not self.upload_s3_endpoint: + msg = "Please provide the upload endpoint for bucket %s " \ + "(default: %s): " % ( + self.get_upload_s3_bucket(), + default_endpoint + ) + user_input = input(_(msg)) + self.upload_s3_endpoint = user_input or default_endpoint + return self.upload_s3_endpoint def prompt_for_upload_user(self): """Should be overridden by policies to determine if a user needs to @@ -544,7 +566,9 @@ def get_upload_s3_endpoint(self): :returns: The S3 Endpoint to use for upload :rtype: ``str`` """ - return self.upload_s3_endpoint or self._upload_s3_endpoint + if not self.upload_s3_endpoint: + self.prompt_for_upload_s3_endpoint() + return self.upload_s3_endpoint def get_upload_s3_region(self): """Helper function to determine if we should use the policy default @@ -562,6 +586,10 @@ def get_upload_s3_bucket(self): :returns: The S3 bucket to use for upload :rtype: ``str`` """ + if self.upload_url and self.upload_url.startswith('s3://'): + return self.upload_url[5:].strip('/') + if not self.upload_s3_bucket: + self.prompt_for_upload_s3_bucket() return self.upload_s3_bucket or self._upload_s3_bucket def get_upload_s3_object_prefix(self): @@ -591,6 +619,14 @@ def get_upload_url(self): :returns: The URL to use for upload :rtype: ``str`` """ + if not self.upload_url and ( + self.upload_s3_bucket and + self.upload_s3_access_key and + self.upload_s3_secret_key + ): + bucket = self.get_upload_s3_bucket() + prefix = self.get_upload_s3_object_prefix() + self._upload_url = f"s3://{bucket}/{prefix}" return self.upload_url or self._upload_url def get_upload_url_string(self): @@ -915,10 +951,14 @@ def upload_s3(self, endpoint=None, region=None, bucket=None, prefix=None, region = self.get_upload_s3_region() if not bucket: - bucket = self.get_upload_s3_bucket() + bucket = self.get_upload_s3_bucket().strip('/') if not prefix: prefix = self.get_upload_s3_object_prefix() + if prefix != '' and prefix.startswith('/'): + prefix = prefix[1:] + if prefix != '' and not prefix.endswith('/'): + prefix = f'{prefix}/' if prefix else '' if not access_key: access_key = self.get_upload_s3_access_key() @@ -932,8 +972,9 @@ def upload_s3(self, endpoint=None, region=None, bucket=None, prefix=None, aws_secret_access_key=secret_key) try: + key = prefix + self.upload_archive_name.split('/')[-1] s3_client.upload_file(self.upload_archive_name, - bucket, prefix + self.upload_archive_name) + bucket, key) return True except Exception as e: raise Exception("Failed to upload to S3: %s" % str(e))