diff --git a/atst/domain/csp/cloud/azure_cloud_provider.py b/atst/domain/csp/cloud/azure_cloud_provider.py index 84a9238c..2b50028f 100644 --- a/atst/domain/csp/cloud/azure_cloud_provider.py +++ b/atst/domain/csp/cloud/azure_cloud_provider.py @@ -19,6 +19,10 @@ from .models import ( BillingProfileVerificationCSPResult, KeyVaultCredentials, ManagementGroupCSPResponse, + ProductPurchaseCSPPayload, + ProductPurchaseCSPResult, + ProductPurchaseVerificationCSPPayload, + ProductPurchaseVerificationCSPResult, TaskOrderBillingCreationCSPPayload, TaskOrderBillingCreationCSPResult, TaskOrderBillingVerificationCSPPayload, @@ -493,6 +497,64 @@ class AzureCloudProvider(CloudProviderInterface): else: return self._error(result.json()) + def create_product_purchase( + self, payload: ProductPurchaseCSPPayload + ): + sp_token = self._get_sp_token(payload.creds) + if sp_token is None: + raise AuthenticationException( + "Could not resolve token for aad premium product purchase" + ) + + create_product_purchase_body = payload.dict(by_alias=True) + create_product_purchase_headers = { + "Authorization": f"Bearer {sp_token}", + } + + product_purchase_url = f"https://management.azure.com//providers/Microsoft.Billing/billingAccounts/{payload.billing_account_name}/billingProfiles/{payload.billing_profile_name}/products?api-version=2019-10-01-preview" + + result = self.sdk.requests.post( + product_purchase_url, + json=create_product_purchase_body, + headers=create_product_purchase_headers, + ) + + if result.status_code == 202: + # 202 has location/retry after headers + return self._ok(ProductPurchaseCSPResult(**result.headers)) + elif result.status_code == 200: + # NB: Swagger docs imply call can sometimes resolve immediately + return self._ok(ProductPurchaseCSPResult(**result.json())) + else: + return self._error(result.json()) + + + def create_product_purchase_verification( + self, payload: ProductPurchaseVerificationCSPPayload + ): + sp_token = self._get_sp_token(payload.creds) + if sp_token is None: + raise AuthenticationException( + "Could not resolve token for task order billing validation" + ) + + auth_header = { + "Authorization": f"Bearer {sp_token}", + } + + result = self.sdk.requests.get( + payload.product_purchase_verify_url, headers=auth_header + ) + + if result.status_code == 202: + # 202 has location/retry after headers + return self._ok(ProductPurchaseCSPResult(**result.headers)) + elif result.status_code == 200: + return self._ok(ProductPurchaseVerificationCSPResult(**result.json())) + else: + return self._error(result.json()) + + def create_remote_admin(self, creds, tenant_details): # create app/service principal within tenant, with name constructed from tenant details # assign principal global admin diff --git a/atst/domain/csp/cloud/mock_cloud_provider.py b/atst/domain/csp/cloud/mock_cloud_provider.py index 10d62e15..bd29f796 100644 --- a/atst/domain/csp/cloud/mock_cloud_provider.py +++ b/atst/domain/csp/cloud/mock_cloud_provider.py @@ -26,6 +26,10 @@ from .models import ( BillingProfileCreationCSPResult, BillingProfileVerificationCSPPayload, BillingProfileVerificationCSPResult, + ProductPurchaseCSPPayload, + ProductPurchaseCSPResult, + ProductPurchaseVerificationCSPPayload, + ProductPurchaseVerificationCSPResult, TaskOrderBillingCreationCSPPayload, TaskOrderBillingCreationCSPResult, TaskOrderBillingVerificationCSPPayload, @@ -277,6 +281,37 @@ class MockCloudProvider(CloudProviderInterface): } ) + + def create_product_purchase(self, payload: ProductPurchaseCSPPayload): + self._maybe_raise(self.NETWORK_FAILURE_PCT, self.NETWORK_EXCEPTION) + self._maybe_raise(self.SERVER_FAILURE_PCT, self.SERVER_EXCEPTION) + self._maybe_raise(self.UNAUTHORIZED_RATE, self.AUTHORIZATION_EXCEPTION) + + return ProductPurchaseCSPResult( + **{ + "name": "TO1:CLIN001", + "properties": { + "amount": 1000.0, + "endDate": "2020-03-01T00:00:00+00:00", + "startDate": "2020-01-01T00:00:00+00:00", + }, + "type": "Microsoft.Billing/billingAccounts/billingProfiles/billingInstructions", + } + ) + + def create_product_purchase_verification( + self, payload: ProductPurchaseVerificationCSPPayload + ): + self._maybe_raise(self.NETWORK_FAILURE_PCT, self.NETWORK_EXCEPTION) + self._maybe_raise(self.SERVER_FAILURE_PCT, self.SERVER_EXCEPTION) + self._maybe_raise(self.UNAUTHORIZED_RATE, self.AUTHORIZATION_EXCEPTION) + + return ProductPurchaseVerificationCSPResult( + **{ + } + ) + + def create_or_update_user(self, auth_credentials, user_info, csp_role_id): self._authorize(auth_credentials) diff --git a/atst/domain/csp/cloud/models.py b/atst/domain/csp/cloud/models.py index b4ff9232..ad9d171e 100644 --- a/atst/domain/csp/cloud/models.py +++ b/atst/domain/csp/cloud/models.py @@ -341,3 +341,50 @@ class KeyVaultCredentials(BaseModel): ) return values + +class AadPremiumProductParameter(AliasModel): + type: str + sku: str + quantity: int + productProperties: Dict + + #{ + # "type": "string", + # "sku": "string", + # "quantity": 0, + # "productProperties": { + # "beneficiaryTenantId": "string" + # } + #} + +class ProductPurchaseCSPPayload(BaseCSPPayload): + billing_account_name: str + billing_profile_name: str + parameters: List[AadPremiumProductParameter] + +class ProductPurchaseCSPResult(AliasModel): + product_purchase_verify_url: str + product_purchase_retry_after: int + + class Config: + fields = { + "product_purchase_verify_url": "Location", + "product_purchase_retry_after": "Retry-After", + } + + +class ProductPurchaseVerificationCSPPayload(BaseCSPPayload): + product_purchase_verify_url: str + + +class ProductPurchaseVerificationCSPResult(AliasModel): + billing_profile_id: str + billing_profile_name: str + billing_profile_enabled_plan_details: BillingProfileEnabledPlanDetails + + class Config: + fields = { + "billing_profile_id": "id", + "billing_profile_name": "name", + "billing_profile_enabled_plan_details": "properties", + } diff --git a/atst/models/mixins/state_machines.py b/atst/models/mixins/state_machines.py index 37682e5b..1db257df 100644 --- a/atst/models/mixins/state_machines.py +++ b/atst/models/mixins/state_machines.py @@ -17,6 +17,8 @@ class AzureStages(Enum): TASK_ORDER_BILLING_CREATION = "task order billing creation" TASK_ORDER_BILLING_VERIFICATION = "task order billing verification" BILLING_INSTRUCTION = "billing instruction" + PRODUCT_PURCHASE = "purchase aad premium product" + PRODUCT_PURCHASE_VERIFICATION = "purchase aad premium product verification" def _build_csp_states(csp_stages):