Adds IAM examples for Python. (GoogleCloudPlatform#1318)

gguuss · andrewsg · commit 49802f11b5d6 · 2018-02-05T13:35:49.000-08:00
diff --git a/iot/api-client/manager/manager.py b/iot/api-client/manager/manager.py
@@ -422,6 +422,42 @@ def get_config_versions(
     return configs
 
 
+def get_iam_permissions(
+        service_account_json, project_id, cloud_region, registry_id):
+    """Retrieves IAM permissions for the given registry."""
+    client = get_client(service_account_json)
+    registry_path = 'projects/{}/locations/{}/registries/{}'.format(
+        project_id, cloud_region, registry_id)
+
+    policy = client.projects().locations().registries().getIamPolicy(
+            resource=registry_path, body={}).execute()
+
+    return policy
+
+
+def set_iam_permissions(
+        service_account_json, project_id, cloud_region, registry_id, role,
+        member):
+    """Retrieves IAM permissions for the given registry."""
+    client = get_client(service_account_json)
+
+    registry_path = 'projects/{}/locations/{}/registries/{}'.format(
+        project_id, cloud_region, registry_id)
+    body = {
+        "policy":
+        {
+            "bindings":
+            [{
+                "members": [member],
+                "role": role
+            }]
+        }
+    }
+
+    return client.projects().locations().registries().setIamPolicy(
+            resource=registry_path, body=body).execute()
+
+
 def parse_command_line_args():
     """Parse command line arguments."""
     default_registry = 'cloudiot_device_manager_example_registry_{}'.format(
@@ -473,6 +509,14 @@ def parse_command_line_args():
             '--version',
             default=None,
             help='Version number for setting device configuration.')
+    parser.add_argument(
+            '--member',
+            default=None,
+            help='Member used for IAM commands.')
+    parser.add_argument(
+            '--role',
+            default=None,
+            help='Role used for IAM commands.')
 
     # Command subparser
     command = parser.add_subparsers(dest='command')
@@ -485,14 +529,16 @@ def parse_command_line_args():
     command.add_parser('delete-device', help=delete_device.__doc__)
     command.add_parser('delete-registry', help=delete_registry.__doc__)
     command.add_parser('get', help=get_device.__doc__)
+    command.add_parser('get-config-versions', help=get_config_versions.__doc__)
+    command.add_parser('get-iam-permissions', help=get_iam_permissions.__doc__)
     command.add_parser('get-registry', help=get_registry.__doc__)
     command.add_parser('get-state', help=get_state.__doc__)
     command.add_parser('list', help=list_devices.__doc__)
     command.add_parser('list-registries', help=list_registries.__doc__)
     command.add_parser('patch-es256', help=patch_es256_auth.__doc__)
     command.add_parser('patch-rs256', help=patch_rsa256_auth.__doc__)
     command.add_parser('set-config', help=patch_rsa256_auth.__doc__)
-    command.add_parser('get-config-versions', help=get_config_versions.__doc__)
+    command.add_parser('set-iam-permissions', help=set_iam_permissions.__doc__)
 
     return parser.parse_args()
 
@@ -525,15 +571,45 @@ def run_create(args):
         create_iot_topic(args.project_id, args.pubsub_topic)
 
 
+def run_get(args):
+    if args.command == 'get':
+        get_device(
+                args.service_account_json, args.project_id,
+                args.cloud_region, args.registry_id, args.device_id)
+
+    elif args.command == 'get-config-versions':
+        get_device(
+                args.service_account_json, args.project_id,
+                args.cloud_region, args.registry_id, args.device_id)
+
+    elif args.command == 'get-state':
+        get_state(
+                args.service_account_json, args.project_id,
+                args.cloud_region, args.registry_id, args.device_id)
+
+    elif args.command == 'get-iam-permissions':
+        print(get_iam_permissions(
+                args.service_account_json, args.project_id,
+                args.cloud_region, args.registry_id))
+
+    elif args.command == 'get-registry':
+        print(get_registry(
+                args.service_account_json, args.project_id,
+                args.cloud_region, args.registry_id))
+
+
 def run_command(args):
     """Calls the program using the specified command."""
     if args.project_id is None:
         print('You must specify a project ID or set the environment variable.')
         return
 
-    if args.command.startswith('create'):
+    elif args.command.startswith('create'):
         run_create(args)
 
+    elif args.command.startswith('get'):
+        run_get(args)
+
     elif args.command == 'delete-device':
         delete_device(
                 args.service_account_json, args.project_id,
@@ -544,21 +620,6 @@ def run_command(args):
                 args.service_account_json, args.project_id,
                 args.cloud_region, args.registry_id)
 
-    elif args.command == 'get':
-        get_device(
-                args.service_account_json, args.project_id,
-                args.cloud_region, args.registry_id, args.device_id)
-
-    elif args.command == 'get-state':
-        get_state(
-                args.service_account_json, args.project_id,
-                args.cloud_region, args.registry_id, args.device_id)
-
-    elif args.command == 'get-registry':
-        print(get_registry(
-                args.service_account_json, args.project_id,
-                args.cloud_region, args.registry_id))
-
     elif args.command == 'list':
         list_devices(
                 args.service_account_json, args.project_id,
@@ -585,6 +646,15 @@ def run_command(args):
                 args.cloud_region, args.registry_id, args.device_id,
                 args.rsa_certificate_file)
 
+    elif args.command == 'set-iam-permissions':
+        if (args.member is None):
+            sys.exit('Error: specify --member')
+        if (args.role is None):
+            sys.exit('Error: specify --role')
+        set_iam_permissions(
+                args.service_account_json, args.project_id,
+                args.cloud_region, args.registry_id, args.role, args.member)
+
     elif args.command == 'set-config':
         if (args.config is None):
             sys.exit('Error: specify --config')
@@ -595,11 +665,6 @@ def run_command(args):
                 args.cloud_region, args.registry_id, args.device_id,
                 args.version, args.config)
 
-    elif args.command == 'get-config-versions':
-        get_device(
-                args.service_account_json, args.project_id,
-                args.cloud_region, args.registry_id, args.device_id)
-
 
 def main():
     args = parse_command_line_args()
diff --git a/iot/api-client/manager/manager_test.py b/iot/api-client/manager/manager_test.py
@@ -63,6 +63,37 @@ def test_create_delete_registry(test_topic, capsys):
             service_account_json, project_id, cloud_region, registry_id)
 
 
+def test_get_iam_permissions(test_topic, capsys):
+    manager.open_registry(
+            service_account_json, project_id, cloud_region, pubsub_topic,
+            registry_id)
+
+    manager.list_devices(
+            service_account_json, project_id, cloud_region, registry_id)
+
+    # Test getting IAM permissions
+    print(manager.get_iam_permissions(
+            service_account_json, project_id, cloud_region, registry_id))
+
+    # Test setting IAM permissions
+    MEMBER = "group:dpebot@google.com"
+    ROLE = "roles/viewer"
+    print(manager.set_iam_permissions(
+            service_account_json, project_id, cloud_region, registry_id,
+            ROLE, MEMBER))
+
+    out, _ = capsys.readouterr()
+
+    # Check that create / list worked
+    assert 'Created registry' in out
+    assert 'eventNotificationConfig' in out
+    assert 'etag' in out
+
+    # Clean up
+    manager.delete_registry(
+            service_account_json, project_id, cloud_region, registry_id)
+
+
 def test_add_delete_unauth_device(test_topic, capsys):
     device_id = device_id_template.format('UNAUTH')
     manager.open_registry(
