""" Utilitários para interagir com a AWS, incluindo execução de queries e espera por resultados. Inclui autenticação e configuração de credenciais, sincronização de diretórios com S3 e execução de comandos AWS CLI. Também fornece funções para criar sessões e clientes AWS, assumir roles e obter informações de identidade do usuário. Este módulo é projetado para ser usado em scripts de CI/CD e automação de tarefas relacionadas à AWS. Pode instalar o cliente AWS CLI se não estiver presente na máquina, garantindo que as dependências necessárias estejam disponíveis. """ from __future__ import annotations import shutil from pathlib import Path from typing import Any from zipfile import ZipFile import boto3 from botocore.config import Config from colorama import Fore from common import print_message from src.datalake.commons.ci_utils import RED, RESET, require_env, run_cmd def make_session( env: dict[str, str], *, region_key: str = "AWS_DEFAULT_REGION" ) -> boto3.session.Session: """ Cria uma sessão do boto3 usando as variáveis de ambiente fornecidas. Args: env (dict[str, str]): Dicionário contendo as variáveis de ambiente. region_key (str, opcional): Chave para a região na variável de ambiente. Padrão é "AWS_DEFAULT_REGION". Returns: boto3.session.Session: Sessão do boto3 configurada com as credenciais e região fornecidas. """ region = env.get(region_key) or env.get("REGION") or "us-east-1" access_key = env.get("AWS_ACCESS_KEY_ID") secret_key = env.get("AWS_SECRET_ACCESS_KEY") session_token = env.get("AWS_SESSION_TOKEN") if access_key and secret_key: return boto3.session.Session( aws_access_key_id=access_key, aws_secret_access_key=secret_key, aws_session_token=session_token, region_name=region, ) return boto3.session.Session(region_name=region) def make_client( session: boto3.session.Session, service_name: str, *, region_name: str = "us-east-1", timeout_config: Config | None = None, **kwargs: Any, ) -> boto3.client.BaseClient: """ Cria um cliente do boto3 para um serviço específico usando a sessão fornecida. Args: session (boto3.session.Session): Sessão do boto3. service_name (str): Nome do serviço AWS. region_name (str | None, opcional): Região da AWS. Padrão é None. timeout_config (Config | None, opcional): Configuração de timeout do botocore. Padrão é None. **kwargs: Argumentos adicionais para o cliente do boto3. Returns: boto3.client.BaseClient: Cliente do boto3 para o serviço especificado. """ if timeout_config is None: timeout_config = Config( connect_timeout=120, read_timeout=900, retries={"max_attempts": 10, "mode": "standard"}, ) return session.client( service_name, region_name=region_name, config=timeout_config, **kwargs ) def assume_role_credentials( session: boto3.session.Session, role_arn: str ) -> dict[str, str]: """ Assume um papel (role) na AWS e retorna as credenciais temporárias. Args: session (boto3.session.Session): Sessão do boto3. role_arn (str): ARN do papel a ser assumido. Returns: dict[str, str]: Credenciais temporárias (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_SESSION_TOKEN). """ sts = make_client(session, "sts") response = sts.assume_role(RoleArn=role_arn, RoleSessionName="AWSCLI-Session") credentials = response["Credentials"] return { "AWS_ACCESS_KEY_ID": credentials["AccessKeyId"], "AWS_SECRET_ACCESS_KEY": credentials["SecretAccessKey"], "AWS_SESSION_TOKEN": credentials["SessionToken"], } def get_caller_identity(session: boto3.session.Session) -> dict[str, str]: """ Obtém a identidade do chamador atual usando o serviço STS. Args: session (boto3.session.Session): Sessão do boto3. Returns: dict[str, str]: Informações da identidade do chamador, incluindo Account, Arn e UserId. """ sts = make_client(session, "sts") response = sts.get_caller_identity() return { "Account": str(response.get("Account", "")), "Arn": str(response.get("Arn", "")), "UserId": str(response.get("UserId", "")), } def resolve_stack_output( session: boto3.session.Session, *, stack_name: str, output_key: str, region_name: str, ) -> str: """ Resolve a saída de um stack do CloudFormation. Args: session (boto3.session.Session): Sessão do boto3. stack_name (str): Nome do stack do CloudFormation. output_key (str): Chave da saída a ser resolvida. region_name (str): Região da AWS. Returns: str: Valor da saída correspondente à chave fornecida. Retorna uma string vazia se não encontrado. """ cf = make_client(session, "cloudformation", region_name=region_name) response = cf.describe_stacks(StackName=stack_name) stacks = response.get("Stacks", []) if not stacks: return "" outputs = stacks[0].get("Outputs", []) for output in outputs: if output.get("OutputKey") == output_key: return str(output.get("OutputValue", "")) return "" def install_aws_dependencies() -> None: """ Instala o cliente AWS CLI se não estiver presente na máquina. Verifica se o cliente AWS CLI está disponível. Se não estiver, baixa e instala o cliente AWS CLI na máquina. Isso garante que as dependências necessárias para interagir com a AWS estejam disponíveis. """ aws_binary = shutil.which("aws") if aws_binary: return print_message( "[Dados] sem cliente AWS na maquina, instalando...", fore_color=Fore.BLUE ) aws_install_dir = Path("temp/aws") installer = aws_install_dir / "install" if not installer.exists(): print_message("[Dados] Efetuando download do cliente AWS...") aws_install_dir.mkdir(parents=True, exist_ok=True) temp_dir = Path("temp") temp_dir.mkdir(exist_ok=True) run_cmd( [ "curl", "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip", "-o", "awscliv2.zip", ], cwd=temp_dir, ) zip_path = temp_dir / "awscliv2.zip" with ZipFile(zip_path) as zip_file: for member in zip_file.infolist(): extracted_path = Path(zip_file.extract(member, temp_dir)) permissions = member.external_attr >> 16 if permissions: extracted_path.chmod(permissions) print_message("[Dados] Instalando cliente AWS...") run_cmd([str(installer)]) run_cmd(["/usr/local/bin/aws", "--version"]) print_message(f"[Dados] Dependencias AWS preparadas!{RESET}") access_ok = False def configure_aws_access(env: dict[str, str], force: bool = False) -> dict[str, str]: """ Configura o acesso AWS usando as variáveis de ambiente fornecidas. Args: env (dict[str, str]): Variáveis de ambiente contendo informações de configuração AWS. Returns: dict[str, str]: Variáveis de ambiente atualizadas com credenciais AWS temporárias. Raises: RuntimeError: Se não for possível obter a identidade do usuário após assumir as roles. """ global access_ok print_message(f"{RED}[Dados] cloudformation - configure access") if access_ok and not force: print_message( f"{RED}[Dados] cloudformation - configure access - acesso AWS já configurado, pulando...{RESET}" ) return env project_type = env.get("DATAOFFICE_PROJECT_TYPE", "") print_message(f"[Dados] DATAOFFICE_PROJECT_TYPE = {project_type}") project_name = project_type.upper() if project_type else "DATA" environment = require_env(env, "ENVIRONMENT") aws_account_id_name = f"AWS_ACCOUNT_ID_{project_name}_{environment.upper()}" aws_default_region_name = f"AWS_DEFAULT_REGION_{project_name}_{environment.upper()}" print_message(f"AWS_ACCOUNT_ID_NAME = {aws_account_id_name}") print_message(f"AWS_DEFAULT_REGION_NAME = {aws_default_region_name}") env["AWS_ACCOUNT_ID"] = require_env(env, aws_account_id_name) env["AWS_DEFAULT_REGION"] = env.get(aws_default_region_name, "us-east-1").strip() print_message(f"AWS_ACCOUNT_ID {env['AWS_ACCOUNT_ID']}") print_message(f"AWS_DEFAULT_REGION {env['AWS_DEFAULT_REGION']}") print_message( "[Dados] cloudformation - configure access - removendo qualquer configuração prévia" ) env.pop("AWS_ACCESS_KEY_ID", None) env.pop("AWS_SECRET_ACCESS_KEY", None) env.pop("AWS_SESSION_TOKEN", None) base_session = make_session(env) identity_before = get_caller_identity(base_session) print_message(f"[Dados] caller identity antes do assume-role: {identity_before}") print_message("[Dados] cloudformation - configure access - devops") print_message(f"[Dados] caller identity devops: {identity_before}") account_id = env["AWS_ACCOUNT_ID"] print_message( "[Dados] cloudformation - configure access - colocando nossa ROLE de entrada - GitlabPipeline" ) first_role = f"arn:aws:iam::{account_id}:role/GitlabPipeline" env.update(assume_role_credentials(base_session, first_role)) first_session = make_session(env) print_message( f"[Dados] cloudformation - configure access - colocando nossa ROLE de ambiente - GitlabPipeline-{environment}" ) second_role = f"arn:aws:iam::{account_id}:role/GitlabPipeline-{environment}" env.update(assume_role_credentials(first_session, second_role)) print_message( "[Dados] cloudformation - configure access - dados de conexão gerados" f"\nAWS_ACCESS_KEY_ID = {env.get('AWS_ACCESS_KEY_ID', '')}" f"\nAWS_SECRET_ACCESS_KEY = {env.get('AWS_SECRET_ACCESS_KEY', '')}" f"\nAWS_SESSION_TOKEN = {env.get('AWS_SESSION_TOKEN', '')}", debug_message=True, ) print_message( "[Dados] cloudformation - configure access - referencia final que efetivamente será utilizada" ) final_session = make_session(env) final_identity = get_caller_identity(final_session) print_message(f"[Dados] caller identity final: {final_identity}") print_message(f"[Dados] cloudformation - configure access - efetivado{RESET}") access_ok = True return env def sync_directory_to_s3( dir_name: str, bucket: str, prefix: str, env: dict[str, str], delete: bool = True, exclude: list[str] | None = None, include: list[str] | None = None, ) -> None: """ Sincroniza um diretório local para um bucket S3 usando o AWS CLI. O diretório é especificado por `dir_name`, e os arquivos são enviados para o bucket e prefixo definidos por `bucket` e `prefix`. A função suporta exclusão de arquivos no destino que não estão presentes na origem (`delete`) e permite especificar padrões de exclusão e inclusão de arquivos. Usa o comando `aws s3 sync` para realizar a sincronização, pois o boto3 ainda não tem suporte a essa função, garantindo que apenas os arquivos necessários sejam transferidos. O ambiente AWS é configurado usando as variáveis de ambiente fornecidas em `env`. Args: dir_name (str): O caminho do diretório local a ser sincronizado. bucket (str): O nome do bucket S3 de destino. prefix (str): O prefixo dentro do bucket S3 onde os arquivos serão armazenados. env (dict[str, str]): Variáveis de ambiente a serem usadas durante a execução do comando AWS CLI. delete (bool, optional): Se True, arquivos no destino que não estão presentes na origem serão excluídos. Padrão é True. exclude (list[str] | None, optional): Lista de padrões de arquivos a serem excluídos da sincronização. Padrão é None. include (list[str] | None, optional): Lista de padrões de arquivos a serem incluídos na sincronização. Padrão é None. Raises: RuntimeError: Se o diretório especificado não existir ou não for um diretório. """ dir_path: Path = Path(dir_name) print_message(f"{Fore.YELLOW}[Dados] procurando pasta {dir_path}") if not dir_path.exists() or not dir_path.is_dir(): return print_message(f"[Dados] pasta encontrada: {dir_path}") install_aws_dependencies() exclude_parameters = [] if exclude: for pattern in exclude: exclude_parameters.append("--exclude") exclude_parameters.append(pattern) include_parameters = [] if include: for pattern in include: include_parameters.append("--include") include_parameters.append(pattern) aws_cli_args = [ "aws", "s3", "sync", dir_name, f"s3://{bucket}/{prefix}", *exclude_parameters, *include_parameters, *(["--delete"] if delete else []), ] print_message( f"{Fore.YELLOW}Executando comando AWS CLI para sincronização de diretório com S3: {aws_cli_args}" ) run_cmd_result = run_cmd(aws_cli_args, env=env, capture_output=False) print_message( f"Resultado da execução do comando AWS CLI para sincronização de diretório com S3: {run_cmd_result}" ) def delete_files( bucket: str, prefix: str, env: dict[str, str], target_name: str | None = None, target_extension: str | None = ".whl", ) -> None: """ Deleta arquivos de um bucket S3 com base em um prefixo e nome de arquivo específico. Args: bucket (str): O nome do bucket S3 de onde os arquivos serão deletados. prefix (str): O prefixo dentro do bucket S3 para filtrar os arquivos a serem deletados. env (dict[str, str]): Variáveis de ambiente a serem usadas para configurar a sessão AWS. target_name (str): O nome do arquivo alvo a ser deletado (sem extensão). target_extension (str, optional): A extensão do arquivo alvo a ser deletado. Padrão é ".whl". Returns: None """ session = make_session(env) s3 = make_client(session, "s3", region_name=env.get("REGION", "")) paginator = s3.get_paginator("list_objects_v2") pages = paginator.paginate(Bucket=bucket, Prefix=prefix) for page in pages: if "Contents" in page and page["Contents"]: for obj in page["Contents"]: if (not target_name and not target_extension) or ( target_name and target_name in obj["Key"] and target_extension and obj["Key"].endswith(target_extension) ): s3.delete_object(Bucket=bucket, Key=obj["Key"]) print_message(f"Removido {obj['Key']} de s3://{bucket}/{prefix}")