Merge pull request #13 from ImMin5/master

Modify additional options schema

src/plugin/connector/base.py

@@ -53,8 +53,9 @@ def set_connect(self, secret_data: dict, tenant_id: str = None) -> None:
             credential=DefaultAzureCredential()
         )
 
-    def _make_request_headers(self, secret_data, client_type=None):
-        access_token = self._get_access_token(secret_data)
+    def _make_request_headers(self, secret_data, access_token=None):
+        if not access_token:
+            access_token = self._get_access_token(secret_data)
         headers = {
             "Authorization": f"Bearer {access_token}",
             "Content-Type": "application/json",

src/plugin/connector/billing_connector.py

@@ -28,6 +28,23 @@ def list_customers(self, billing_account_id: str) -> list:
         )
         return customers
 
+    def list_departments(self, secret_data: dict, billing_account_id: str) -> list:
+        try:
+            api_version = "2019-10-01-preview"
+            self.next_link = f"https://management.azure.com/providers/Microsoft.Billing/billingAccounts/{billing_account_id}/departments?api-version={api_version}"
+
+            while self.next_link:
+                url = self.next_link
+
+                headers = self._make_request_headers(secret_data)
+                response = requests.get(url=url, headers=headers)
+                response_json = response.json()
+
+                self.next_link = response_json.get("properties").get("nextLink", None)
+                yield response_json
+        except Exception as e:
+            raise ERROR_UNKNOWN(message=f"[ERROR] list_departments {e}")
+
     def list_subscription(
         self, secret_data: dict, agreement_type: str, billing_account_id: str
     ) -> list:

src/plugin/main.py

@@ -22,22 +22,21 @@ def account_collector_init(params: dict) -> dict:
         }
     """
 
+    options = params.get("options", {}) or {}
+
     metadata = {
         "additional_options_schema": {
             "type": "object",
             "properties": {
-                "exclude_tenant_root_group": {
+                "exclude_root_management_group": {
                     "title": "Exclude Tenant Root Group",
+                    "oneOf": [True, False],
                     "type": "boolean",
-                    "default": True,
                 },
             },
         }
     }
 
-    if options := params.get("options"):
-        pass
-
     return {"metadata": metadata}
 
 

src/plugin/manager/resource_manager.py

@@ -85,6 +85,10 @@ def sync(
                             tenant_id, subscription_id, name, location
                         )
                     result_subscription_map[subscription_id] = result
+            if agreement_type == "EnterpriseAgreement":
+                departments = billing_connector.list_departments(
+                    secret_data, billing_account_id
+                )
 
             tenants = subscription_connector.list_tenants()
             try:
@@ -179,7 +183,7 @@ def _create_location_from_entity_info(
         parent_display_name_chain = entity_info.get("parent_display_name_chain", [])
         parent_name_chain = entity_info.get("parent_name_chain", [])
         for idx, name in enumerate(parent_display_name_chain):
-            if options.get("exclude_tenant_root_group") and idx == 0:
+            if options.get("exclude_root_management_group") and idx == 0:
                 continue
             location.append(
                 {
@@ -245,6 +249,8 @@ def _get_location(
             location_name = subscription_info.get("customer_display_name", "")
 
         if location_name and resource_id:
+            if resource_id == "23fd0c55-250c-4a62-b124-f7ca35ce2061":
+                print(subscription_info, "okok")
             location = [
                 {
                     "name": location_name,