Merge branch 'staging' into azure-subscriptions

atst/domain/csp/cloud/azure_cloud_provider.py

@@ -26,6 +26,10 @@ from .models import (
     BillingProfileVerificationCSPResult,
     KeyVaultCredentials,
     ManagementGroupCSPResponse,
+    ProductPurchaseCSPPayload,
+    ProductPurchaseCSPResult,
+    ProductPurchaseVerificationCSPPayload,
+    ProductPurchaseVerificationCSPResult,
     PrincipalAdminRoleCSPPayload,
     PrincipalAdminRoleCSPResult,
     TaskOrderBillingCreationCSPPayload,
@@ -94,6 +98,7 @@ class AzureCloudProvider(CloudProviderInterface):
         self.ps_client_id = config["POWERSHELL_CLIENT_ID"]
         self.owner_role_def_id = config["AZURE_OWNER_ROLE_DEF_ID"]
         self.graph_resource = config["AZURE_GRAPH_RESOURCE"]
+        self.default_aadp_qty = config["AZURE_AADP_QTY"]
 
         if azure_sdk_provider is None:
             self.sdk = AzureSDKProvider()
@@ -526,6 +531,70 @@ class AzureCloudProvider(CloudProviderInterface):
         else:
             return self._error(result.json())
 
+    def create_product_purchase(self, payload: ProductPurchaseCSPPayload):
+        sp_token = self._get_root_provisioning_token()
+        if sp_token is None:
+            raise AuthenticationException(
+                "Could not resolve token for aad premium product purchase"
+            )
+
+        create_product_purchase_body = {
+            "type": "AADPremium",
+            "sku": "AADP1",
+            "productProperties": {"beneficiaryTenantId": payload.tenant_id,},
+            "quantity": self.default_aadp_qty,
+        }
+        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}/purchaseProduct?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(ProductPurchaseVerificationCSPResult(**result.json()))
+        else:
+            return self._error(result.json())
+
+    def create_product_purchase_verification(
+        self, payload: ProductPurchaseVerificationCSPPayload
+    ):
+        sp_token = self._get_root_provisioning_token()
+        if sp_token is None:
+            raise AuthenticationException(
+                "Could not resolve token for aad premium product purchase 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:
+            premium_purchase_date = result.json()["properties"]["purchaseDate"]
+            return self._ok(
+                ProductPurchaseVerificationCSPResult(
+                    premium_purchase_date=premium_purchase_date
+                )
+            )
+        else:
+            return self._error(result.json())
+
     def create_tenant_admin_ownership(self, payload: TenantAdminOwnershipCSPPayload):
         mgmt_token = self._get_elevated_management_token(payload.tenant_id)
 
@@ -799,7 +868,7 @@ class AzureCloudProvider(CloudProviderInterface):
     def _get_root_provisioning_token(self):
         creds = self._source_creds()
         return self._get_sp_token(
-            creds.tenant_id, creds.root_sp_client_id, creds.root_sp_key
+            creds.root_tenant_id, creds.root_sp_client_id, creds.root_sp_key
         )
 
     def _get_sp_token(self, tenant_id, client_id, secret_key):

atst/domain/csp/cloud/mock_cloud_provider.py

@@ -25,6 +25,10 @@ from .models import (
     BillingProfileTenantAccessCSPResult,
     BillingProfileVerificationCSPPayload,
     BillingProfileVerificationCSPResult,
+    ProductPurchaseCSPPayload,
+    ProductPurchaseCSPResult,
+    ProductPurchaseVerificationCSPPayload,
+    ProductPurchaseVerificationCSPResult,
     PrincipalAdminRoleCSPPayload,
     PrincipalAdminRoleCSPResult,
     SubscriptionCreationCSPPayload,
@@ -315,11 +319,33 @@ class MockCloudProvider(CloudProviderInterface):
             }
         )
 
-    def create_tenant_admin_ownership(self, payload: TenantAdminOwnershipCSPPayload):
+    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(
+            **dict(
+                product_purchase_verify_url="https://zombo.com",
+                product_purchase_retry_after=10,
+            )
+        )
+
+    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(
+            **dict(premium_purchase_date="2020-01-30T18:57:05.981Z")
+        )
+
+    def create_tenant_admin_ownership(self, payload: TenantAdminOwnershipCSPPayload):
+        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 TenantAdminOwnershipCSPResult(**dict(id="admin_owner_assignment_id"))
 
     def create_tenant_principal_ownership(

atst/domain/csp/cloud/models.py

@@ -450,3 +450,27 @@ class SuscriptionVerificationCSPResult(AliasModel):
 
     class Config:
         fields = {"subscription_id": "subscriptionLink"}
+
+
+class ProductPurchaseCSPPayload(BaseCSPPayload):
+    billing_account_name: str
+    billing_profile_name: str
+
+
+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):
+    premium_purchase_date: str

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"
     TENANT_PRINCIPAL_APP = "tenant principal application"
     TENANT_PRINCIPAL = "tenant principal"
     TENANT_PRINCIPAL_CREDENTIAL = "tenant principal credential"