diff --git a/app.py b/app.py index 23cd4e2..5dcc449 100644 --- a/app.py +++ b/app.py @@ -97,7 +97,7 @@ async def run_groups_authn(): updated_creds = get_credentials(creds, SCOPES) # create UserService object for API calls - user_service = UserService(creds) + user_service = UserService(updated_creds) # keep track of IAM group and database instance tasks group_tasks = {} diff --git a/iam_groups_authn/mysql.py b/iam_groups_authn/mysql.py index a7f692f..a3dea7b 100644 --- a/iam_groups_authn/mysql.py +++ b/iam_groups_authn/mysql.py @@ -19,6 +19,7 @@ from google.cloud.sql.connector import connector from google.cloud.sql.connector.instance_connection_manager import IPTypes from iam_groups_authn.utils import async_wrap +from google.auth.transport.requests import Request def mysql_username(iam_email): @@ -137,6 +138,11 @@ def init_connection_engine(instance_connection_name, creds, ip_type=IPTypes.PUBL "pool_timeout": 30, # 30 seconds "pool_recycle": 1800, # 30 minutes } + # refresh credentials if not valid + if not creds.valid: + request = Request() + creds.refresh(request) + # service account email to access DB, mysql truncates usernames to before '@' sign service_account_email = mysql_username(creds.service_account_email) diff --git a/iam_groups_authn/sql_admin.py b/iam_groups_authn/sql_admin.py index 30c5e7e..02a02fa 100644 --- a/iam_groups_authn/sql_admin.py +++ b/iam_groups_authn/sql_admin.py @@ -73,7 +73,7 @@ async def add_missing_db_users( missing_db_users = get_users_to_add(iam_users, db_users) # add missing users to database instance for user in missing_db_users: - user_service.insert_db_user( + await user_service.insert_db_user( user, InstanceConnectionName(*instance_connection_name.split(":")) ) return missing_db_users diff --git a/iam_groups_authn/sync.py b/iam_groups_authn/sync.py index 553ea29..7df4d15 100644 --- a/iam_groups_authn/sync.py +++ b/iam_groups_authn/sync.py @@ -14,13 +14,14 @@ # sync.py contains functions for syncing IAM groups with Cloud SQL instances +import asyncio from google.auth import iam from google.auth.transport.requests import Request from google.oauth2 import service_account -from googleapiclient.discovery import build -from googleapiclient.errors import HttpError from iam_groups_authn.mysql import mysql_username -from iam_groups_authn.utils import async_wrap +import json +from aiohttp import ClientSession +from enum import Enum # URI for OAuth2 credentials TOKEN_URI = "https://accounts.google.com/o/oauth2/token" @@ -36,9 +37,11 @@ def __init__(self, creds): creds: OAuth2 credentials to call admin APIs. """ self.creds = creds + self.client_session = ClientSession( + headers={"Content-Type": "application/json"} + ) - @async_wrap - def get_group_members(self, group): + async def get_group_members(self, group): """Get all members of an IAM group. Given an IAM group, get all members (groups or users) that belong to the @@ -51,21 +54,23 @@ def get_group_members(self, group): members: List of all members (groups or users) that belong to the IAM group. """ # build service to call Admin SDK Directory API - service = build("admin", "directory_v1", credentials=self.creds) + url = f"https://admin.googleapis.com/admin/directory/v1/groups/{group}/members" try: # call the Admin SDK Directory API - results = service.members().list(groupKey=group).execute() + resp = await authenticated_request( + self.creds, url, self.client_session, RequestType.get + ) + results = json.loads(await resp.text()) members = results.get("members", []) return members # handle errors if IAM group does not exist etc. - except HttpError as e: - raise HttpError( + except Exception as e: + raise Exception( f"Error: Failed to get IAM members of IAM group `{group}`. Verify group exists and is configured correctly." ) from e - @async_wrap - def get_db_users(self, instance_connection_name): + async def get_db_users(self, instance_connection_name): """Get all database users of a Cloud SQL instance. Given a database instance and a Google Cloud project, get all the database @@ -79,17 +84,17 @@ def get_db_users(self, instance_connection_name): Returns: users: List of all database users that belong to the Cloud SQL instance. """ - # build service to call SQL Admin API - service = build("sqladmin", "v1beta4", credentials=self.creds) + # build request to SQL Admin API + project = instance_connection_name.project + instance = instance_connection_name.instance + url = f"https://sqladmin.googleapis.com/sql/v1beta4/projects/{project}/instances/{instance}/users" + try: - results = ( - service.users() - .list( - project=instance_connection_name.project, - instance=instance_connection_name.instance, - ) - .execute() + # call the SQL Admin API + resp = await authenticated_request( + self.creds, url, self.client_session, RequestType.get ) + results = json.loads(await resp.text()) users = results.get("items", []) return users except Exception as e: @@ -97,7 +102,7 @@ def get_db_users(self, instance_connection_name): f"Error: Failed to get the database users for instance `{instance_connection_name}`. Verify instance connection name and instance details." ) from e - def insert_db_user(self, user_email, instance_connection_name): + async def insert_db_user(self, user_email, instance_connection_name): """Create DB user from IAM user. Given an IAM user's email, insert the IAM user as a DB user for Cloud SQL instance. @@ -108,18 +113,16 @@ def insert_db_user(self, user_email, instance_connection_name): (e.g. InstanceConnectionName(project='my-project', region='my-region', instance='my-instance')) """ - # build service to call SQL Admin API - service = build("sqladmin", "v1beta4", credentials=self.creds) + # build request to SQL Admin API + project = instance_connection_name.project + instance = instance_connection_name.instance + url = f"https://sqladmin.googleapis.com/sql/v1beta4/projects/{project}/instances/{instance}/users" user = {"name": user_email, "type": "CLOUD_IAM_USER"} + try: - results = ( - service.users() - .insert( - project=instance_connection_name.project, - instance=instance_connection_name.instance, - body=user, - ) - .execute() + # call the SQL Admin API + resp = await authenticated_request( + self.creds, url, self.client_session, RequestType.post, body=user ) return except Exception as e: @@ -127,6 +130,57 @@ def insert_db_user(self, user_email, instance_connection_name): f"Error: Failed to add IAM user `{user_email}` to Cloud SQL database instance `{instance_connection_name.instance}`." ) from e + def __del__(self): + """Deconstructor for UserService to close ClientSession and have + graceful exit. + """ + + async def deconstruct(): + if not self.client_session.closed: + await self.client_session.close() + + asyncio.run_coroutine_threadsafe(deconstruct(), loop=asyncio.get_event_loop()) + + +class RequestType(Enum): + """Helper class for supported aiohttp request types.""" + + get = 1 + post = 2 + + +async def authenticated_request(creds, url, client_session, request_type, body=None): + """Helper function to build authenticated aiohttp requests. + + Args: + creds: OAuth2 credentials for authorizing requests. + url: URL for aiohttp request. + client_session: aiohttp ClientSession object. + request_type: RequestType enum determining request type. + body: (optional) JSON body for request. + + Return: + Result from aiohttp request. + """ + if not creds.valid: + request = Request() + creds.refresh(request) + + headers = { + "Authorization": f"Bearer {creds.token}", + } + + if request_type == RequestType.get: + return await client_session.get(url, headers=headers, raise_for_status=True) + elif request_type == RequestType.post: + return await client_session.post( + url, headers=headers, json=body, raise_for_status=True + ) + else: + raise ValueError( + "Request type not recognized! " "Please verify RequestType is valid." + ) + async def get_users_with_roles(role_service, role): """Get mapping of group role grants on DB users. diff --git a/requirements.txt b/requirements.txt index 0fc659c..2bddfba 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,8 +1,8 @@ quart==0.15.1 hypercorn==0.11.2 SQLAlchemy==1.4.22 -google-api-python-client==2.19.1 google-auth==2.0.0 PyMySQL==1.0.2 cloud-sql-python-connector[pymysql]==0.4.1 google-cloud-logging==2.6.0 +aiohttp==3.8.0 diff --git a/tests/requirements-test.txt b/tests/requirements-test.txt index 2f0f677..f457f19 100644 --- a/tests/requirements-test.txt +++ b/tests/requirements-test.txt @@ -1,7 +1,6 @@ pytest==6.2.5 quart==0.15.1 SQLAlchemy==1.4.22 -google-api-python-client==2.19.1 google-auth==2.0.0 PyMySQL==1.0.2 cloud-sql-python-connector[pymysql]==0.4.1