"""Commands to manage DDB.""" import asyncio import csv import importlib import io import os.path from typing import Any, Dict, Tuple from uuid import UUID import aiofiles import typer from pdp import config from pdp.cli import get_data from pdp.connectors.dynamo import DynamoDbConnector from pdp.connectors.redis_client import RedisConnector from pdp.fastapi.schemas import identity as identity_schema from pdp.logic.cache import bust_identity_caches from pdp.logic.identity import attach_and_detach_roles cli: typer.Typer = typer.Typer( short_help="Commands to manage dynamodb", no_args_is_help=True ) DEV_M2M_IDENTITY = "pdpcli" QA_REFRESH_IDENTITY_UUID = "QA_REFRESH_IDENTITY_UUID" def get_m2m_identity() -> str: """Get the identity used to populate the updated_at in the pp_identity table.""" if config.ENVIRONMENT in (config.DEV_ENVIRONMENT, config.TEST_ENVIRONMENT): identity = DEV_M2M_IDENTITY else: identity = os.getenv(QA_REFRESH_IDENTITY_UUID, "") if not identity: typer.secho( "Missing 'QA_REFRESH_IDENTITY_UUID' environment variable." " Required for QA env. Exiting...", fg="red", ) raise typer.Exit(1) return identity def set_env(allowed_environments: Tuple[str, ...]) -> None: """Set environment vars for DEV. Otherwise abort to avoid accidentally changing QA or Prod envs. """ if config.ENVIRONMENT not in allowed_environments: typer.secho( f"This command only supports {allowed_environments}. " f"Found '{config.ENVIRONMENT}'", fg="red", ) raise typer.Exit(1) if config.ENVIRONMENT in (config.DEV_ENVIRONMENT, config.TEST_ENVIRONMENT): if "DYNAMODB_ENDPOINT_URL" not in os.environ: os.environ["DYNAMODB_ENDPOINT_URL"] = "http://localhost:50028" os.environ["AWS_DEFAULT_REGION"] = "us-east-1" os.environ["AWS_ACCESS_KEY_ID"] = "fakeMyKeyId" os.environ["AWS_SECRET_ACCESS_KEY"] = "fakeSecretAccessKey" # pdp.utils.get_opts references config for the DYNAMODB_ENDPOINT_URL. # reload the module to pick up environ values defined above. importlib.reload(config) elif config.ENVIRONMENT == config.QA_ENVIRONMENT: # sanity check the expected DYNAMODB_TABLE_IDENTITY config in QA if config.DYNAMODB_TABLE_IDENTITY not in ( "qa_pp_identity", "pp_identity_refresh", ): typer.secho( f"Unexpected DynamoDB table name for QA: " f"{config.DYNAMODB_TABLE_IDENTITY}.", fg="red", ) raise typer.Exit(1) else: typer.secho( f"This module only supports DEV environment. Found '{config.ENVIRONMENT}'", fg="red", ) raise typer.Exit(1) def assert_development_env() -> None: """Aborts the script unless Environment is 'dev'.""" if config.ENVIRONMENT in (config.QA_ENVIRONMENT, config.PROD_ENVIRONMENT): typer.secho( f"This module only supports DEV environment. Found '{config.ENVIRONMENT}'", fg="red", ) raise typer.Exit(1) def drop_pp_identity_table(connector: DynamoDbConnector) -> None: """Drop the development pp_identity table.""" assert_development_env() tables = connector.client.list_tables() if config.DYNAMODB_TABLE_IDENTITY in tables["TableNames"]: typer.secho(f"Dropping table: {config.DYNAMODB_TABLE_IDENTITY}", fg="green") connector.client.delete_table(TableName=config.DYNAMODB_TABLE_IDENTITY) def create_pp_identity_table(connector: DynamoDbConnector) -> None: """Create the development pp_identity table.""" assert_development_env() typer.secho(f"Creating table: {config.DYNAMODB_TABLE_IDENTITY}", fg="green") connector.client.create_table( TableName=config.DYNAMODB_TABLE_IDENTITY, AttributeDefinitions=[ {"AttributeName": config.IDENTITY_HASH_KEY, "AttributeType": "S"}, {"AttributeName": config.IDENTITY_RANGE_KEY, "AttributeType": "S"}, ], KeySchema=[ {"AttributeName": config.IDENTITY_HASH_KEY, "KeyType": "HASH"}, {"AttributeName": config.IDENTITY_RANGE_KEY, "KeyType": "RANGE"}, ], ProvisionedThroughput={"ReadCapacityUnits": 10, "WriteCapacityUnits": 5}, ) # Enable TTL on the table connector.client.update_time_to_live( TableName=config.DYNAMODB_TABLE_IDENTITY, TimeToLiveSpecification={"Enabled": True, "AttributeName": "expires_at"}, ) @cli.command("create", short_help="Drop and Create the pp_identity table") def create() -> None: """Command to drop and create the DEV table.""" set_env(allowed_environments=(config.DEV_ENVIRONMENT, config.TEST_ENVIRONMENT)) connector = DynamoDbConnector( config.DYNAMODB_TABLE_IDENTITY, config.IDENTITY_HASH_KEY, config.IDENTITY_RANGE_KEY, ) drop_pp_identity_table(connector) create_pp_identity_table(connector) @cli.command("seed", short_help="Insert seed data into pp_identity table") def seed_dynamodb() -> None: """Load roles and identities from seed.csv using Identity logic.""" set_env(allowed_environments=(config.DEV_ENVIRONMENT, config.TEST_ENVIRONMENT)) redis_connector = RedisConnector( config.REDIS_URL, use_redis_cache=config.CACHE_USE_REDIS, ) asyncio.run(_load_csv(get_data("seed.csv"), redis_connector)) @cli.command("load_csv", short_help="Insert csv file data into pp_identity table") def load_csv(input_csv: str) -> None: """CLI method to load a PP roles CSV file to DynamoDB.""" set_env( allowed_environments=( config.DEV_ENVIRONMENT, config.TEST_ENVIRONMENT, config.QA_ENVIRONMENT, ) ) redis_connector = RedisConnector( config.REDIS_URL, use_redis_cache=config.CACHE_USE_REDIS, ) asyncio.run(_load_csv(input_csv, redis_connector)) async def _load_csv( input_csv: str, redis_connector: RedisConnector, ) -> None: """Load a CSV file to DynamoDB using pdp.logic.""" if not os.path.isfile(input_csv): typer.secho(f"seed csv file not found at: {input_csv}", fg="red") raise typer.Exit(1) # Dictionary to load CSV file into memory identities: Dict[str, Any] = {} async with aiofiles.open(input_csv, newline="") as f: content = await f.read() # Treat the string content like a CSV file csvfile = io.StringIO(content) reader = csv.DictReader(csvfile) for row in reader: # remove whitespaces row = {str(k).strip(): str(v).strip() for k, v in dict(row).items()} # Get the Identity if it has already been read from the CSV file identity = identities.get(row["identity_uuid"], {}) # Then transform the new row to store into memory as a Tenant+Role tenant = identity.get(row["tenant_uuid"], {}) tenant.update( {"tenant_uuid": row["tenant_uuid"], "tenant_type": row["tenant_type"]} ) roles_to_attach = tenant.get("roles_to_attach", []) roles_to_detach = tenant.get("roles_to_detach", []) operation = row.get("operation", "detach") # Default to detach if operation == "attach": roles_to_attach.append(row["role"]) else: roles_to_detach.append(row["role"]) tenant["roles_to_attach"] = roles_to_attach tenant["roles_to_detach"] = roles_to_detach # Update the Identity dictionary for the Tenant identity[row["tenant_uuid"]] = tenant # Update the Identities dictionary for the Identity identities[row["identity_uuid"]] = identity typer.secho( f"LOADED: identity:{row['identity_uuid']} " f"tenant:{row['tenant_uuid']} role:{row['role']} operation:{operation}", # noqa: E501 fg="green", ) # Assumes we can load an entire CSV file into memory # If we have issues in the future, we can batch these updates, # but that might be a premature improvement now. await _update_identities(identities) identity_uuids = [UUID(identity_uuid) for identity_uuid in identities.keys()] await bust_identity_caches( identity_uuids, redis_connector=redis_connector, ) async def _update_identities( identities: Dict[str, Any], ) -> None: identity_ddb_connector = DynamoDbConnector( config.DYNAMODB_TABLE_IDENTITY, config.IDENTITY_HASH_KEY, config.IDENTITY_RANGE_KEY, ) authenticated_identity_uuid = get_m2m_identity() for identity_uuid, identity in identities.items(): assert UUID(identity_uuid) for tenant in identity.values(): await attach_and_detach_roles( identity_uuid, UUID(tenant["tenant_uuid"]), tenant["tenant_type"], roles_to_attach=[ identity_schema.Role(role=role) for role in tenant["roles_to_attach"] ], roles_to_detach=[ identity_schema.Role(role=role) for role in tenant["roles_to_detach"] ], authenticated_identity_uuid=authenticated_identity_uuid, identity_ddb_connector=identity_ddb_connector, ) typer.secho( f"UPDATE: identity:{identity_uuid}", fg="green", )