From 024274e133d8acb48fa68e21687b3eb5c796309c Mon Sep 17 00:00:00 2001
From: ImMin5 <mino@megazone.com>
Date: Fri, 19 Apr 2024 11:18:00 +0900
Subject: [PATCH] feat: modify additional options schema

---
 src/plugin/connector/base.py              |  5 +++--
 src/plugin/connector/billing_connector.py | 17 +++++++++++++++++
 src/plugin/main.py                        |  9 ++++-----
 src/plugin/manager/resource_manager.py    |  8 +++++++-
 4 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/src/plugin/connector/base.py b/src/plugin/connector/base.py
index 57d9dee..d6aac8c 100644
--- a/src/plugin/connector/base.py
+++ b/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",
diff --git a/src/plugin/connector/billing_connector.py b/src/plugin/connector/billing_connector.py
index 32f5c19..6e74ecc 100644
--- a/src/plugin/connector/billing_connector.py
+++ b/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:
diff --git a/src/plugin/main.py b/src/plugin/main.py
index 1081b67..0ceb07f 100644
--- a/src/plugin/main.py
+++ b/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}
 
 
diff --git a/src/plugin/manager/resource_manager.py b/src/plugin/manager/resource_manager.py
index 8310ceb..1e6289f 100644
--- a/src/plugin/manager/resource_manager.py
+++ b/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,