GoogleCloudPlatform · gcf-merge-on-green · Sep 22, 2020 · Sep 18, 2020 · Sep 22, 2020
diff --git a/iam/api-client/grantable_roles_test.py b/iam/api-client/grantable_roles_test.py
@@ -17,7 +17,7 @@
 import grantable_roles
 
 
-def test_service_accounts(capsys):
+def test_grantable_roles(capsys):
     project = os.environ['GOOGLE_CLOUD_PROJECT']
     resource = '//cloudresourcemanager.googleapis.com/projects/' + project
     grantable_roles.view_grantable_roles(resource)

diff --git a/iam/api-client/quickstart.py b/iam/api-client/quickstart.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python
-
-# Copyright 2018 Google Inc. All Rights Reserved.
+#
+# Copyright 2020 Google Inc. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -14,38 +14,111 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# [START iam_quickstart]
+import os
+
+from google.oauth2 import service_account
+import googleapiclient.discovery
+
+
+def quickstart(project_id, member):
+    """Gets a policy, adds a member, prints their permissions, and removes the member."""
+
+    # Role to be granted.
+    role = "roles/logging.logWriter"
 
-def quickstart():
-    # [START iam_quickstart]
-    import os
+    # Initializes service.
+    crm_service = initialize_service()
 
-    from google.oauth2 import service_account
-    import googleapiclient.discovery
+    # Grants your member the 'Log Writer' role for the project.
+    modify_policy_add_role(crm_service, project_id, role, member)
+
+    # Gets the project's policy and prints all members with the 'Log Writer' role.
+    policy = get_policy(crm_service, project_id)
+    binding = next(b for b in policy["bindings"] if b["role"] == role)
+    print(f'Role: {(binding["role"])}')
+    print("Members: ")
+    for m in binding["members"]:
+        print(f'[{m}]')
+
+    # Removes the member from the 'Log Writer' role.
+    modify_policy_remove_member(crm_service, project_id, role, member)
+
+
+def initialize_service():
+    """Initializes a Cloud Resource Manager service."""
 
-    # Get credentials
     credentials = service_account.Credentials.from_service_account_file(
-        filename=os.environ['GOOGLE_APPLICATION_CREDENTIALS'],
-        scopes=['https://www.googleapis.com/auth/cloud-platform'])
+        filename=os.environ["GOOGLE_APPLICATION_CREDENTIALS"],
+        scopes=["https://www.googleapis.com/auth/cloud-platform"],
+    )
+    crm_service = googleapiclient.discovery.build(
+        "cloudresourcemanager", "v1", credentials=credentials
+    )
+    return crm_service
+
+
+def modify_policy_add_role(crm_service, project_id, role, member):
+    """Adds a new role binding to a policy."""
 
-    # Create the Cloud IAM service object
-    service = googleapiclient.discovery.build(
-        'iam', 'v1', credentials=credentials)
+    policy = get_policy(crm_service, project_id)
 
-    # Call the Cloud IAM Roles API
-    # If using pylint, disable weak-typing warnings
-    # pylint: disable=no-member
-    response = service.roles().list().execute()
-    roles = response['roles']
+    binding = None
+    for b in policy["bindings"]:
+        if b["role"] == role:
+            binding = b
+            break
+    if binding is not None:
+        binding["members"].append(member)
+    else:
+        binding = {"role": role, "members": [member]}
+        policy["bindings"].append(binding)
 
-    # Process the response
-    for role in roles:
-        print('Title: ' + role['title'])
-        print('Name: ' + role['name'])
-        if 'description' in role:
-            print('Description: ' + role['description'])
-        print('')
-    # [END iam_quickstart]
+    set_policy(crm_service, project_id, policy)
+
+
+def modify_policy_remove_member(crm_service, project_id, role, member):
+    """Removes a  member from a role binding."""
+
+    policy = get_policy(crm_service, project_id)
+
+    binding = next(b for b in policy["bindings"] if b["role"] == role)
+    if "members" in binding and member in binding["members"]:
+        binding["members"].remove(member)
+
+    set_policy(crm_service, project_id, policy)
+
+
+def get_policy(crm_service, project_id, version=3):
+    """Gets IAM policy for a project."""
+
+    policy = (
+        crm_service.projects()
+        .getIamPolicy(
+            resource=project_id,
+            body={"options": {"requestedPolicyVersion": version}},
+        )
+        .execute()
+    )
+    return policy
+
+
+def set_policy(crm_service, project_id, policy):
+    """Sets IAM policy for a project."""
+
+    policy = (
+        crm_service.projects()
+        .setIamPolicy(resource=project_id, body={"policy": policy})
+        .execute()
+    )
+    return policy
 
 
 if __name__ == '__main__':
-    quickstart()
+
+    # TODO: replace with your project ID
+    project_id = "your_project_id"
+    # TODO: Replace with the ID of your member in the form 'user:member@example.com.
+    member = "your_member"
+    quickstart(project_id, member)
+# [END iam_quickstart]
diff --git a/iam/api-client/quickstart_test.py b/iam/api-client/quickstart_test.py
@@ -1,4 +1,4 @@
-# Copyright 2018 Google Inc. All Rights Reserved.
+# Copyright 2020 Google Inc. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,10 +12,89 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""Tests for quickstart."""
+
+import os
+import uuid
+
+from google.oauth2 import service_account
+from googleapiclient import errors
+import googleapiclient.discovery
+import pytest
+from retrying import retry
+
 import quickstart
 
+# Setting up variables for testing
+GCLOUD_PROJECT = os.environ["GCLOUD_PROJECT"]
+
+
+def retry_if_conflict(exception):
+    return (isinstance(exception, errors.HttpError)
+            and 'There were concurrent policy changes' in str(exception))
+
+
+@pytest.fixture(scope="module")
+def test_member():
+    # section to create service account to test policy updates.
+    # we use the first portion of uuid4 because full version is too long.
+    name = f'test-{uuid.uuid4().hex[:25]}'
+    email = name + "@" + GCLOUD_PROJECT + ".iam.gserviceaccount.com"
+    member = "serviceAccount:" + email
+    create_service_account(
+        GCLOUD_PROJECT, name, "Py Test Account"
+    )
+
+    yield member
+
+    # deleting the service account created above
+    delete_service_account(email)
+
+
+def create_service_account(project_id, name, display_name):
+    """Creates a service account."""
+
+    credentials = service_account.Credentials.from_service_account_file(
+        filename=os.environ['GOOGLE_APPLICATION_CREDENTIALS'],
+        scopes=['https://www.googleapis.com/auth/cloud-platform'])
+
+    service = googleapiclient.discovery.build(
+        'iam', 'v1', credentials=credentials)
+
+    my_service_account = service.projects().serviceAccounts().create(
+        name='projects/' + project_id,
+        body={
+            'accountId': name,
+            'serviceAccount': {
+                'displayName': display_name
+            }
+        }).execute()
+
+    print('Created service account: ' + my_service_account['email'])
+    return my_service_account
+
+
+def delete_service_account(email):
+    """Deletes a service account."""
+
+    credentials = service_account.Credentials.from_service_account_file(
+        filename=os.environ['GOOGLE_APPLICATION_CREDENTIALS'],
+        scopes=['https://www.googleapis.com/auth/cloud-platform'])
+
+    service = googleapiclient.discovery.build(
+        'iam', 'v1', credentials=credentials)
+
+    service.projects().serviceAccounts().delete(
+        name='projects/-/serviceAccounts/' + email).execute()
+
+    print('Deleted service account: ' + email)
+
 
-def test_quickstart(capsys):
-    quickstart.quickstart()
-    out, _ = capsys.readouterr()
-    assert 'Title' in out
+def test_quickstart(test_member, capsys):
+    @retry(wait_exponential_multiplier=1000, wait_exponential_max=10000,
+           stop_max_attempt_number=5, retry_on_exception=retry_if_conflict)
+    def test_call():
+        quickstart.quickstart(GCLOUD_PROJECT, test_member)
+        out, _ = capsys.readouterr()
+        assert test_member in out
+    test_call()
diff --git a/iam/api-client/quickstartv2.py b/iam/api-client/quickstartv2.py