allow deployment to non-default subscriptions (#774)

src/deployment/deploy.py

@@ -122,7 +122,9 @@ class Client:
         log_service_principal: bool,
         multi_tenant_domain: str,
         upgrade: bool,
+        subscription_id: Optional[str],
     ):
+        self.subscription_id = subscription_id
         self.resource_group = resource_group
         self.arm_template = arm_template
         self.location = location
@@ -171,11 +173,16 @@ class Client:
             self.workbook_data = json.load(f)
 
     def get_subscription_id(self) -> str:
+        if self.subscription_id:
+            return self.subscription_id
         profile = get_cli_profile()
-        return cast(str, profile.get_subscription_id())
+        self.subscription_id = cast(str, profile.get_subscription_id())
+        return self.subscription_id
 
     def get_location_display_name(self) -> str:
-        location_client = get_client_from_cli_profile(SubscriptionClient)
+        location_client = get_client_from_cli_profile(
+            SubscriptionClient, subscription_id=self.get_subscription_id()
+        )
         locations = location_client.subscriptions.list_locations(
             self.get_subscription_id()
         )
@@ -194,7 +201,9 @@ class Client:
         with open(self.arm_template, "r") as handle:
             arm = json.load(handle)
 
-        client = get_client_from_cli_profile(ResourceManagementClient)
+        client = get_client_from_cli_profile(
+            ResourceManagementClient, subscription_id=self.get_subscription_id()
+        )
         providers = {x.namespace: x for x in client.providers.list()}
 
         unsupported = []
@@ -233,7 +242,7 @@ class Client:
             sys.exit(1)
 
     def create_password(self, object_id: UUID) -> Tuple[str, str]:
-        return add_application_password(object_id)
+        return add_application_password(object_id, self.get_subscription_id())
 
     def setup_rbac(self) -> None:
         """
@@ -245,7 +254,9 @@ class Client:
             logger.info("using existing client application")
             return
 
-        client = get_client_from_cli_profile(GraphRbacManagementClient)
+        client = get_client_from_cli_profile(
+            GraphRbacManagementClient, subscription_id=self.get_subscription_id()
+        )
         logger.info("checking if RBAC already exists")
 
         try:
@@ -400,7 +411,10 @@ class Client:
                 "subscription, creating a new one"
             )
             app_info = register_application(
-                "onefuzz-cli", self.application_name, OnefuzzAppRole.CliClient
+                "onefuzz-cli",
+                self.application_name,
+                OnefuzzAppRole.CliClient,
+                self.get_subscription_id(),
             )
             if self.multi_tenant_domain:
                 authority = COMMON_AUTHORITY
@@ -429,7 +443,9 @@ class Client:
         with open(self.arm_template, "r") as template_handle:
             template = json.load(template_handle)
 
-        client = get_client_from_cli_profile(ResourceManagementClient)
+        client = get_client_from_cli_profile(
+            ResourceManagementClient, subscription_id=self.get_subscription_id()
+        )
         client.resource_groups.create_or_update(
             self.resource_group, {"location": self.location}
         )
@@ -512,6 +528,7 @@ class Client:
         assign_scaleset_role(
             self.application_name,
             self.results["deploy"]["scaleset-identity"]["value"],
+            self.get_subscription_id(),
         )
 
     def apply_migrations(self) -> None:
@@ -548,7 +565,9 @@ class Client:
         logger.info("creating eventgrid subscription")
         src_resource_id = self.results["deploy"]["fuzz-storage"]["value"]
         dst_resource_id = self.results["deploy"]["func-storage"]["value"]
-        client = get_client_from_cli_profile(StorageManagementClient)
+        client = get_client_from_cli_profile(
+            StorageManagementClient, subscription_id=self.get_subscription_id()
+        )
         event_subscription_info = EventSubscription(
             destination=StorageQueueEventSubscriptionDestination(
                 resource_id=dst_resource_id, queue_name="file-changes"
@@ -565,7 +584,9 @@ class Client:
             ),
         )
 
-        client = get_client_from_cli_profile(EventGridManagementClient)
+        client = get_client_from_cli_profile(
+            EventGridManagementClient, subscription_id=self.get_subscription_id()
+        )
         result = client.event_subscriptions.create_or_update(
             src_resource_id, "onefuzz1", event_subscription_info
         ).result()
@@ -639,7 +660,8 @@ class Client:
         )
 
         app_insight_client = get_client_from_cli_profile(
-            ApplicationInsightsManagementClient
+            ApplicationInsightsManagementClient,
+            subscription_id=self.get_subscription_id(),
         )
 
         to_delete = []
@@ -829,7 +851,7 @@ class Client:
     def update_registration(self) -> None:
         if not self.create_registration:
             return
-        update_pool_registration(self.application_name)
+        update_pool_registration(self.application_name, self.get_subscription_id())
 
     def done(self) -> None:
         logger.info(TELEMETRY_NOTICE)
@@ -967,6 +989,10 @@ def main() -> None:
         default=None,
         help="enable multi-tenant authentication with this tenant domain",
     )
+    parser.add_argument(
+        "--subscription_id",
+        type=str,
+    )
     args = parser.parse_args()
 
     if shutil.which("func") is None:
@@ -992,6 +1018,7 @@ def main() -> None:
         log_service_principal=args.log_service_principal,
         multi_tenant_domain=args.multi_tenant_domain,
         upgrade=args.upgrade,
+        subscription_id=args.subscription_id,
     )
     if args.verbose:
         level = logging.DEBUG

src/deployment/registration.py

@@ -80,9 +80,10 @@ def query_microsoft_graph(
         )
 
 
-def get_graph_client() -> GraphRbacManagementClient:
+def get_graph_client(subscription_id: str) -> GraphRbacManagementClient:
     client: GraphRbacManagementClient = get_client_from_cli_profile(
-        GraphRbacManagementClient
+        GraphRbacManagementClient,
+        subscription_id=subscription_id,
     )
     return client
 
@@ -99,10 +100,13 @@ class OnefuzzAppRole(Enum):
 
 
 def register_application(
-    registration_name: str, onefuzz_instance_name: str, approle: OnefuzzAppRole
+    registration_name: str,
+    onefuzz_instance_name: str,
+    approle: OnefuzzAppRole,
+    subscription_id: str,
 ) -> ApplicationInfo:
     logger.info("retrieving the application registration %s" % registration_name)
-    client = get_graph_client()
+    client = get_graph_client(subscription_id)
     apps: List[Application] = list(
         client.applications.list(filter="displayName eq '%s'" % registration_name)
     )
@@ -110,7 +114,7 @@ def register_application(
     if len(apps) == 0:
         logger.info("No existing registration found. creating a new one")
         app = create_application_registration(
-            onefuzz_instance_name, registration_name, approle
+            onefuzz_instance_name, registration_name, approle, subscription_id
         )
     else:
         app = apps[0]
@@ -136,7 +140,7 @@ def register_application(
     if app.app_id not in [app.app_id for app in pre_authorized_applications]:
         authorize_application(UUID(app.app_id), UUID(onefuzz_app.app_id))
 
-    password = create_application_credential(registration_name)
+    password = create_application_credential(registration_name, subscription_id)
 
     return ApplicationInfo(
         client_id=app.app_id,
@@ -145,27 +149,27 @@ def register_application(
     )
 
 
-def create_application_credential(application_name: str) -> str:
+def create_application_credential(application_name: str, subscription_id: str) -> str:
     """ Add a new password to the application registration """
 
     logger.info("creating application credential for '%s'" % application_name)
-    client = get_graph_client()
+    client = get_graph_client(subscription_id)
     apps: List[Application] = list(
         client.applications.list(filter="displayName eq '%s'" % application_name)
     )
 
     app: Application = apps[0]
 
-    (_, password) = add_application_password(app.object_id)
+    (_, password) = add_application_password(app.object_id, subscription_id)
     return str(password)
 
 
 def create_application_registration(
-    onefuzz_instance_name: str, name: str, approle: OnefuzzAppRole
+    onefuzz_instance_name: str, name: str, approle: OnefuzzAppRole, subscription_id: str
 ) -> Application:
     """ Create an application registration """
 
-    client = get_graph_client()
+    client = get_graph_client(subscription_id)
     apps: List[Application] = list(
         client.applications.list(filter="displayName eq '%s'" % onefuzz_instance_name)
     )
@@ -207,7 +211,6 @@ def create_application_registration(
         try:
             time.sleep(5)
 
-            client = get_graph_client()
             update_param = ApplicationUpdateParameters(
                 reply_urls=["https://%s.azurewebsites.net" % onefuzz_instance_name]
             )
@@ -221,7 +224,9 @@ def create_application_registration(
     return registered_app
 
 
-def add_application_password(app_object_id: UUID) -> Tuple[str, str]:
+def add_application_password(
+    app_object_id: UUID, subscription_id: str
+) -> Tuple[str, str]:
     # Work-around the race condition where the app is created but passwords cannot
     # be created yet.
 
@@ -234,7 +239,7 @@ def add_application_password(app_object_id: UUID) -> Tuple[str, str]:
         if count > 1:
             logging.info("retrying app password creation")
         try:
-            password = add_application_password_impl(app_object_id)
+            password = add_application_password_impl(app_object_id, subscription_id)
             logging.info("app password created")
             return password
         except GraphQueryError as err:
@@ -255,11 +260,13 @@ def add_application_password(app_object_id: UUID) -> Tuple[str, str]:
         raise Exception("unable to create password")
 
 
-def add_application_password_legacy(app_object_id: UUID) -> Tuple[str, str]:
+def add_application_password_legacy(
+    app_object_id: UUID, subscription_id: str
+) -> Tuple[str, str]:
     key = str(uuid4())
     password = str(uuid4())
 
-    client = get_graph_client()
+    client = get_graph_client(subscription_id)
     password_cred = [
         PasswordCredential(
             start_date="%s" % datetime.now(TZ_UTC).strftime("%Y-%m-%dT%H:%M.%fZ"),
@@ -275,7 +282,9 @@ def add_application_password_legacy(app_object_id: UUID) -> Tuple[str, str]:
     return (key, password)
 
 
-def add_application_password_impl(app_object_id: UUID) -> Tuple[str, str]:
+def add_application_password_impl(
+    app_object_id: UUID, subscription_id: str
+) -> Tuple[str, str]:
     key = uuid4()
     password_request = {
         "passwordCredential": {
@@ -296,7 +305,7 @@ def add_application_password_impl(app_object_id: UUID) -> Tuple[str, str]:
         )
         return (str(key), password["secretText"])
     except adal.AdalError:
-        return add_application_password_legacy(app_object_id)
+        return add_application_password_legacy(app_object_id, subscription_id)
 
 
 def get_application(app_id: UUID) -> Optional[Any]:
@@ -359,13 +368,17 @@ def authorize_application(
 
 
 def create_and_display_registration(
-    onefuzz_instance_name: str, registration_name: str, approle: OnefuzzAppRole
+    onefuzz_instance_name: str,
+    registration_name: str,
+    approle: OnefuzzAppRole,
+    subscription_id: str,
 ) -> None:
     logger.info("Updating application registration")
     application_info = register_application(
         registration_name=registration_name,
         onefuzz_instance_name=onefuzz_instance_name,
         approle=approle,
+        subscription_id=subscription_id,
     )
     logger.info("Registration complete")
     logger.info("These generated credentials are valid for a year")
@@ -373,19 +386,20 @@ def create_and_display_registration(
     logger.info("client_secret: %s" % application_info.client_secret)
 
 
-def update_pool_registration(onefuzz_instance_name: str) -> None:
+def update_pool_registration(onefuzz_instance_name: str, subscription_id: str) -> None:
     create_and_display_registration(
         onefuzz_instance_name,
         "%s_pool" % onefuzz_instance_name,
         OnefuzzAppRole.ManagedNode,
+        subscription_id,
     )
 
 
 def assign_scaleset_role_manually(
-    onefuzz_instance_name: str, scaleset_name: str
+    onefuzz_instance_name: str, scaleset_name: str, subscription_id: str
 ) -> None:
 
-    client = get_graph_client()
+    client = get_graph_client(subscription_id)
     apps: List[Application] = list(
         client.applications.list(filter="displayName eq '%s'" % onefuzz_instance_name)
     )
@@ -441,7 +455,9 @@ def assign_scaleset_role_manually(
     )
 
 
-def assign_scaleset_role(onefuzz_instance_name: str, scaleset_name: str) -> None:
+def assign_scaleset_role(
+    onefuzz_instance_name: str, scaleset_name: str, subscription_id: str
+) -> None:
     """
     Allows the nodes in the scaleset to access the service by assigning
     their managed identity to the ManagedNode Role
@@ -513,7 +529,9 @@ def assign_scaleset_role(onefuzz_instance_name: str, scaleset_name: str) -> None
                 },
             )
     except adal.AdalError:
-        assign_scaleset_role_manually(onefuzz_instance_name, scaleset_name)
+        assign_scaleset_role_manually(
+            onefuzz_instance_name, scaleset_name, subscription_id
+        )
 
 
 def set_app_audience(objectId: str, audience: str) -> None:
@@ -555,6 +573,7 @@ def main() -> None:
     parent_parser.add_argument(
         "onefuzz_instance", help="the name of the onefuzz instance"
     )
+    parent_parser.add_argument("subscription_id")
 
     parser = argparse.ArgumentParser(
         formatter_class=formatter,
@@ -593,14 +612,19 @@ def main() -> None:
 
     onefuzz_instance_name = args.onefuzz_instance
     if args.command == "update_pool_registration":
-        update_pool_registration(onefuzz_instance_name)
+        update_pool_registration(onefuzz_instance_name, args.subscription_id)
     elif args.command == "create_cli_registration":
         registration_name = args.registration_name or ("%s_cli" % onefuzz_instance_name)
         create_and_display_registration(
-            onefuzz_instance_name, registration_name, OnefuzzAppRole.CliClient
+            onefuzz_instance_name,
+            registration_name,
+            OnefuzzAppRole.CliClient,
+            args.subscription_id,
         )
     elif args.command == "assign_scaleset_role":
-        assign_scaleset_role(onefuzz_instance_name, args.scaleset_name)
+        assign_scaleset_role(
+            onefuzz_instance_name, args.scaleset_name, args.subscription_id
+        )
     else:
         raise Exception("invalid arguments")