Upload s3 protocol get/prompt logic.

- parsing bucket from --upload-url or --upload-s3-bucket
- default s3-endpoint offered during prompt

Signed-off-by: Trevor Benson <[email protected]>

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