From a8c63f437afd7ef37a124752fddbd6599f28da01 Mon Sep 17 00:00:00 2001
From: Pepe Fagoaga <pepe@verica.io>
Date: Thu, 3 Aug 2023 19:14:31 +0200
Subject: [PATCH] chore(azure): Improve service class with __set_clients__

---
 .../providers/azure/lib/service/service.py    | 30 +++++++++++++++----
 .../services/defender/defender_service.py     | 23 +-------------
 .../azure/services/iam/iam_service.py         | 23 +-------------
 .../services/sqlserver/sqlserver_service.py   | 23 +-------------
 .../azure/services/storage/storage_service.py | 23 +-------------
 5 files changed, 29 insertions(+), 93 deletions(-)

diff --git a/prowler/providers/azure/lib/service/service.py b/prowler/providers/azure/lib/service/service.py
index fe3697ca07..f888bfe41c 100644
--- a/prowler/providers/azure/lib/service/service.py
+++ b/prowler/providers/azure/lib/service/service.py
@@ -1,8 +1,28 @@
+from prowler.lib.logger import logger
+
+
 class AzureService:
-    def __init__(self, service, audit_info):
-        # We receive the service using __class__.__name__ or the service name in lowercase
-        # e.g.: Storage --> we need a lowercase string, so service.lower()
-        self.service = service.lower() if not service.islower() else service
+    def __init__(
+        self,
+        service,
+        audit_info,
+    ):
+        self.clients = self.__set_clients__(
+            audit_info.identity.subscriptions, audit_info.credentials, service
+        )
 
-        self.credentials = audit_info.credentials
         self.subscriptions = audit_info.identity.subscriptions
+
+    def __set_clients__(self, subscriptions, credentials, service):
+        clients = {}
+        try:
+            for display_name, id in subscriptions.items():
+                clients.update(
+                    {display_name: service(credential=credentials, subscription_id=id)}
+                )
+        except Exception as error:
+            logger.error(
+                f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
+            )
+        else:
+            return clients
diff --git a/prowler/providers/azure/services/defender/defender_service.py b/prowler/providers/azure/services/defender/defender_service.py
index 50a3af8d78..d234c8c06a 100644
--- a/prowler/providers/azure/services/defender/defender_service.py
+++ b/prowler/providers/azure/services/defender/defender_service.py
@@ -10,31 +10,10 @@
 ########################## Defender
 class Defender(AzureService):
     def __init__(self, audit_info):
-        super().__init__(__class__.__name__, audit_info)
+        super().__init__(SecurityCenter, audit_info)
 
-        self.clients = self.__set_clients__(
-            audit_info.identity.subscriptions, audit_info.credentials
-        )
         self.pricings = self.__get_pricings__()
 
-    def __set_clients__(self, subscriptions, credentials):
-        clients = {}
-        try:
-            for display_name, id in subscriptions.items():
-                clients.update(
-                    {
-                        display_name: SecurityCenter(
-                            credential=credentials, subscription_id=id
-                        )
-                    }
-                )
-        except Exception as error:
-            logger.error(
-                f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
-            )
-        else:
-            return clients
-
     def __get_pricings__(self):
         logger.info("Defender - Getting pricings...")
         pricings = {}
diff --git a/prowler/providers/azure/services/iam/iam_service.py b/prowler/providers/azure/services/iam/iam_service.py
index d38b68533c..b282222135 100644
--- a/prowler/providers/azure/services/iam/iam_service.py
+++ b/prowler/providers/azure/services/iam/iam_service.py
@@ -10,30 +10,9 @@
 ########################## IAM
 class IAM(AzureService):
     def __init__(self, audit_info):
-        super().__init__(__class__.__name__, audit_info)
-        self.clients = self.__set_clients__(
-            audit_info.identity.subscriptions, audit_info.credentials
-        )
+        super().__init__(AuthorizationManagementClient, audit_info)
         self.roles = self.__get_roles__()
 
-    def __set_clients__(self, subscriptions, credentials):
-        clients = {}
-        try:
-            for display_name, id in subscriptions.items():
-                clients.update(
-                    {
-                        display_name: AuthorizationManagementClient(
-                            credential=credentials, subscription_id=id
-                        )
-                    }
-                )
-        except Exception as error:
-            logger.error(
-                f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
-            )
-        else:
-            return clients
-
     def __get_roles__(self):
         logger.info("IAM - Getting roles...")
         roles = {}
diff --git a/prowler/providers/azure/services/sqlserver/sqlserver_service.py b/prowler/providers/azure/services/sqlserver/sqlserver_service.py
index 07d8388f93..165593cbef 100644
--- a/prowler/providers/azure/services/sqlserver/sqlserver_service.py
+++ b/prowler/providers/azure/services/sqlserver/sqlserver_service.py
@@ -14,30 +14,9 @@
 ########################## SQLServer
 class SQLServer(AzureService):
     def __init__(self, audit_info):
-        super().__init__(__class__.__name__, audit_info)
-        self.clients = self.__set_clients__(
-            audit_info.identity.subscriptions, audit_info.credentials
-        )
+        super().__init__(SqlManagementClient, audit_info)
         self.sql_servers = self.__get_sql_servers__()
 
-    def __set_clients__(self, subscriptions, credentials):
-        clients = {}
-        try:
-            for display_name, id in subscriptions.items():
-                clients.update(
-                    {
-                        display_name: SqlManagementClient(
-                            credential=credentials, subscription_id=id
-                        )
-                    }
-                )
-        except Exception as error:
-            logger.error(
-                f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
-            )
-        else:
-            return clients
-
     def __get_sql_servers__(self):
         logger.info("SQL Server - Getting SQL servers...")
         sql_servers = {}
diff --git a/prowler/providers/azure/services/storage/storage_service.py b/prowler/providers/azure/services/storage/storage_service.py
index 7837c922eb..b74d988de8 100644
--- a/prowler/providers/azure/services/storage/storage_service.py
+++ b/prowler/providers/azure/services/storage/storage_service.py
@@ -10,30 +10,9 @@
 ########################## Storage
 class Storage(AzureService):
     def __init__(self, audit_info):
-        super().__init__(__class__.__name__, audit_info)
-        self.clients = self.__set_clients__(
-            audit_info.identity.subscriptions, audit_info.credentials
-        )
+        super().__init__(StorageManagementClient, audit_info)
         self.storage_accounts = self.__get_storage_accounts__()
 
-    def __set_clients__(self, subscriptions, credentials):
-        clients = {}
-        try:
-            for display_name, id in subscriptions.items():
-                clients.update(
-                    {
-                        display_name: StorageManagementClient(
-                            credential=credentials, subscription_id=id
-                        )
-                    }
-                )
-        except Exception as error:
-            logger.error(
-                f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
-            )
-        else:
-            return clients
-
     def __get_storage_accounts__(self):
         logger.info("Storage - Getting storage accounts...")
         storage_accounts = {}