[GCS]-Parse configurations from Kibana UI

connectors/source.py

@@ -20,6 +20,16 @@
 """
 
 
+class DataSourceConfigurationError(Exception):
+    """Raise when configurations are invalid.
+
+    Args:
+        Exception (DataSourceConfigurationError): Invalid configurations exception.
+    """
+
+    pass
+
+
 class Field:
     def __init__(self, name, label=None, value="", type="str"):
         if label is None:

connectors/sources/google_cloud_storage.py

@@ -10,6 +10,7 @@
 import os
 import urllib.parse
 from functools import partial
+from json.decoder import JSONDecodeError
 
 import aiofiles
 from aiofiles.os import remove
@@ -18,7 +19,7 @@
 from aiogoogle.auth.creds import ServiceAccountCreds
 
 from connectors.logger import logger
-from connectors.source import BaseDataSource
+from connectors.source import BaseDataSource, DataSourceConfigurationError
 from connectors.utils import TIKA_SUPPORTED_FILETYPES, convert_to_b64
 
 CLOUD_STORAGE_READ_ONLY_SCOPE = "https://www.googleapis.com/auth/devstorage.read_only"
@@ -72,27 +73,14 @@ def __init__(self, configuration):
             configuration (DataSourceConfiguration): Object of DataSourceConfiguration class.
         """
         super().__init__(configuration=configuration)
-        if not self.configuration["service_account_credentials"]:
-            raise Exception("service_account_credentials can't be empty.")
-
-        self.credentials = (
-            self.configuration["service_account_credentials"]
-            .strip()
-            .encode("unicode_escape")
-            .decode()
-        )
-        self.service_account_credentials = ServiceAccountCreds(
-            scopes=[CLOUD_STORAGE_READ_ONLY_SCOPE],
-            **json.loads(
-                self.configuration["service_account_credentials"]
-                if RUNNING_FTEST
-                else self.credentials
-            ),
-        )
-        self.user_project_id = self.service_account_credentials.project_id
+
         self.retry_count = self.configuration["retry_count"]
         self.enable_content_extraction = self.configuration["enable_content_extraction"]
 
+        self.service_account_credentials = None
+        self.user_project_id = None
+        self.json_creds = None
+
     @classmethod
     def get_default_configuration(cls):
         """Get the default configuration for Google Cloud Storage.
@@ -113,7 +101,7 @@ def get_default_configuration(cls):
         return {
             "service_account_credentials": {
                 "value": json.dumps(default_credentials),
-                "label": "JSON string for Google Cloud service account",
+                "label": "Google Cloud service account json",
                 "type": "str",
             },
             "retry_count": {
@@ -199,8 +187,61 @@ async def _api_call(
                 )
                 await asyncio.sleep(DEFAULT_WAIT_MULTIPLIER**retry_counter)
 
+    def _validate_configurations(self):
+        """Validates whether user inputs are valid or not for configuration field.
+
+        Raises:
+            DataSourceConfigurationError: The service account json is empty.
+            DataSourceConfigurationError: The service account json is in invalid format.
+        """
+        if self.configuration["service_account_credentials"] == "":
+            raise DataSourceConfigurationError("Service account json can't be empty.")
+        try:
+            self.json_creds = json.loads(
+                self.configuration["service_account_credentials"]
+            )
+        except JSONDecodeError:
+            raise DataSourceConfigurationError(
+                "Service account json is not in valid format."
+            )
+
+    def get_pem_format(self, private_key, max_split=-1):
+        """Convert key into PEM format.
+
+        Args:
+            private_key (str): Private_key in raw format.
+            max_split (int): Specifies how many splits to do. Defaults to -1.
+
+        Returns:
+            string: PEM format
+        """
+        private_key = private_key.replace(" ", "\n")
+        private_key = " ".join(private_key.split("\n", max_split))
+        private_key = " ".join(private_key.rsplit("\n", max_split))
+        return private_key
+
+    def _initialize_configurations(self):
+        """Initialize the ServiceAccountCreds"""
+        if (
+            self.json_creds.get("private_key")
+            and "\n" not in self.json_creds["private_key"]
+        ):
+            self.json_creds["private_key"] = self.get_pem_format(
+                private_key=self.json_creds["private_key"].strip(), max_split=2
+            )
+
+        self.service_account_credentials = ServiceAccountCreds(
+            scopes=[CLOUD_STORAGE_READ_ONLY_SCOPE],
+            **self.json_creds,
+        )
+        self.user_project_id = self.service_account_credentials.project_id
+
     async def ping(self):
-        """Verify the connection with Google Cloud Storage"""
+        """Verify the configurations and connection with Google Cloud Storage"""
+
+        self._validate_configurations()
+        self._initialize_configurations()
+
         if RUNNING_FTEST:
             return
         try:

connectors/sources/tests/test_google_cloud_storage.py

@@ -32,6 +32,8 @@ def get_mocked_source_object():
         {"service_account_credentials": SERVICE_ACCOUNT_CREDENTIALS, "retry_count": 0}
     )
     mocked_gcs_object = GoogleCloudStorageDataSource(configuration=configuration)
+    mocked_gcs_object._validate_configurations()
+    mocked_gcs_object._initialize_configurations()
     return mocked_gcs_object
 
 
@@ -62,10 +64,25 @@ async def test_empty_configuration():
 
     # Setup
     configuration = DataSourceConfiguration({"service_account_credentials": ""})
+    gcs_object = GoogleCloudStorageDataSource(configuration=configuration)
 
     # Execute
-    with pytest.raises(Exception, match="service_account_credentials can't be empty."):
-        _ = GoogleCloudStorageDataSource(configuration=configuration)
+    with pytest.raises(Exception, match="Service account json can't be empty."):
+        gcs_object._validate_configurations()
+
+
+def test_invalid_configuration():
+    """Tests that the service account credential is in invalid json format"""
+
+    # Setup
+    configuration = DataSourceConfiguration(
+        {"service_account_credentials": "{'abc':'bcd','cd'}"}
+    )
+    gcs_object = GoogleCloudStorageDataSource(configuration=configuration)
+
+    # Execute
+    with pytest.raises(Exception, match="Service account json is not in valid format."):
+        gcs_object._validate_configurations()
 
 
 @pytest.mark.asyncio
@@ -499,3 +516,28 @@ async def test_api_call_for_attribute_error(catch_stdout, patch_logger):
             userProject=mocked_gcs_object.user_project_id,
         ):
             print("Method called successfully....")
+
+
+def test_get_pem_format():
+    """This function tests prepare private key with dummy values"""
+    # Setup
+    expected_formated_pem_key = """-----BEGIN PRIVATE KEY-----
+PrivateKey
+-----END PRIVATE KEY-----"""
+    source = get_mocked_source_object()
+    private_key = "-----BEGIN PRIVATE KEY----- PrivateKey -----END PRIVATE KEY-----"
+
+    # Execute
+    formated_privat_key = source.get_pem_format(private_key, max_split=2)
+    assert formated_privat_key == expected_formated_pem_key
+
+    # Setup
+    expected_formated_certificate = """-----BEGIN CERTIFICATE-----
+Certificate1
+Certificate2
+-----END CERTIFICATE-----"""
+    private_key = "-----BEGIN CERTIFICATE----- Certificate1 Certificate2 -----END CERTIFICATE-----"
+
+    # Execute
+    formated_privat_key = source.get_pem_format(private_key, max_split=1)
+    assert formated_privat_key == expected_formated_certificate