NAG Utilities¶

def apply_compliance_aspects(stack: Stack) -> None:
    """Attach every cdk-nag rule pack this project runs to ``stack``."""
    Aspects.of(stack).add(AwsSolutionsChecks(verbose=True))
    Aspects.of(stack).add(ServerlessChecks(verbose=True))
    Aspects.of(stack).add(NIST80053R5Checks(verbose=True))
    Aspects.of(stack).add(HIPAASecurityChecks(verbose=True))
    Aspects.of(stack).add(PCIDSS321Checks(verbose=True))

def attach_async_failure_destination(
    scope: IConstruct,
    singleton_id: str,
    *,
    encryption_key: kms.Key,
    queue_id: str,
) -> sqs.Queue | None:
    """Wire an SQS DLQ to a CDK-managed async singleton Lambda.

    AwsCustomResource provider Lambdas are invoked asynchronously by
    CloudFormation during stack lifecycle events. Without an on_failure
    destination, a provider crash that exhausts Lambda's two automatic
    async retries is silently dropped — the stack rollback still surfaces
    a CFN error, but the *cause* (Python traceback, AWS API error response)
    is gone unless someone catches it in CloudWatch within the retention
    window. SQS as the on_failure destination preserves the failed-event
    envelope (full request payload + responseContext) for post-mortem.

    The queue uses the same CMK as the surrounding stack, with 14-day
    retention (Lambda's max meaningful window — events older than that
    have already aged past most rollback investigations).

    Returns the created queue so callers can attach alarms or outputs;
    returns None if the singleton isn't present under ``scope`` (which
    happens when no AwsCustomResource has been instantiated in this stack).
    """
    singleton = scope.node.try_find_child(singleton_id)
    # IFunction is a JSII protocol that isn't runtime-checkable, so we check
    # the concrete Function class. SingletonFunction is a subclass, so the
    # isinstance check covers both.
    if not isinstance(singleton, _lambda.Function):
        return None

    dlq = sqs.Queue(
        cast(Construct, scope),
        queue_id,
        encryption=sqs.QueueEncryption.KMS,
        encryption_master_key=encryption_key,
        retention_period=Duration.days(14),
        enforce_ssl=True,
        removal_policy=RemovalPolicy.DESTROY,
    )

    # This queue IS the dead-letter destination. cdk-nag flags any SQS queue
    # without a DLQ or a redrive policy, but recursing DLQs into more DLQs
    # makes no sense — when this terminal queue's consumer fails, manual
    # inspection of the queue content is the recovery path, not another DLQ.
    dlq_terminal_reason = (
        "Terminal DLQ: this queue IS the dead-letter destination — recursing into another DLQ has no recovery value"
    )
    NagSuppressions.add_resource_suppressions(
        dlq,
        [
            {"id": "AwsSolutions-SQS3", "reason": dlq_terminal_reason},
            {"id": "Serverless-SQSRedrivePolicy", "reason": dlq_terminal_reason},
        ],
    )

    singleton.configure_async_invoke(on_failure=destinations.SqsDestination(dlq))

    # configure_async_invoke + SqsDestination + KMS-encrypted queue adds
    # kms:GenerateDataKey* and kms:ReEncrypt* wildcards to the singleton's
    # auto-generated default policy so it can encrypt messages to the DLQ.
    # These are granular IAM5 findings that need applies_to scoping rather
    # than the blanket suppression in CDK_LAMBDA_SUPPRESSIONS. Also
    # re-applies the inline-policy suppressions because the DefaultPolicy
    # resource only materialized when configure_async_invoke modified the
    # role above, after the initial suppress_cdk_singletons run.
    kms_wildcard_reason = (
        "KMS wildcards required by configure_async_invoke to encrypt messages to the CMK-encrypted DLQ"
    )
    NagSuppressions.add_resource_suppressions(
        cast(Construct, singleton),
        [
            {
                "id": "AwsSolutions-IAM5",
                "applies_to": ["Action::kms:GenerateDataKey*", "Action::kms:ReEncrypt*"],
                "reason": kms_wildcard_reason,
            },
            {
                "id": "NIST.800.53.R5-IAMNoInlinePolicy",
                "reason": "CDK-generated inline policy on singleton service role",
            },
            {
                "id": "HIPAA.Security-IAMNoInlinePolicy",
                "reason": "CDK-generated inline policy on singleton service role",
            },
            {"id": "PCI.DSS.321-IAMNoInlinePolicy", "reason": "CDK-generated inline policy on singleton service role"},
        ],
        apply_to_children=True,
    )

    return dlq

def grant_guardduty_service_to_key(key: kms.Key, *, region: str, account: str, partition: str) -> None:
    """Grant GuardDuty ``kms:Decrypt`` on a CMK so Lambda Protection can introspect.

    GuardDuty Lambda Protection (and similar foundational-detection features)
    needs to read Lambda function configuration — including env vars encrypted
    with a customer-managed key. Without this grant the assumed
    ``AWSServiceRoleForAmazonGuardDuty`` role is denied ``kms:Decrypt`` against
    the CMK, leaving GuardDuty's coverage of CMK-encrypted resources incomplete
    (the original CloudTrail finding that motivated this grant).

    Scoped to GuardDuty detectors in this account+region only via
    ``aws:SourceAccount`` and ``aws:SourceArn`` — the cross-account
    confused-deputy guard AWS documents for service-principal grants.
    """
    key.add_to_resource_policy(
        iam.PolicyStatement(
            sid="AllowGuardDutyDecrypt",
            actions=["kms:Decrypt"],
            principals=[iam.ServicePrincipal("guardduty.amazonaws.com")],
            resources=["*"],
            conditions={
                "StringEquals": {"aws:SourceAccount": account},
                "ArnLike": {"aws:SourceArn": f"arn:{partition}:guardduty:{region}:{account}:detector/*"},
            },
        )
    )

def grant_logs_service_to_key(key: kms.Key, *, region: str, account: str, partition: str) -> None:
    """Add the standard CloudWatch Logs service-principal grant to a CMK.

    Three CMKs in this project (backend, frontend, WAF) need the same statement:
    a grant to ``logs.{region}.amazonaws.com`` for symmetric encrypt/decrypt
    operations, conditioned via ``kms:EncryptionContext:aws:logs:arn`` so only
    log groups in this account+region can request key operations. Defining it
    in one place keeps the three call sites in lockstep — pylint's R0801
    duplicate-code check correctly flags any drift between them, and the
    confused-deputy condition is exactly the kind of thing that's harmful to
    forget on one of the three CMKs.
    """
    key.add_to_resource_policy(
        iam.PolicyStatement(
            actions=["kms:Encrypt*", "kms:Decrypt*", "kms:ReEncrypt*", "kms:GenerateDataKey*", "kms:Describe*"],
            principals=[iam.ServicePrincipal(f"logs.{region}.amazonaws.com")],
            resources=["*"],
            conditions={
                "ArnLike": {
                    "kms:EncryptionContext:aws:logs:arn": f"arn:{partition}:logs:{region}:{account}:log-group:*",
                },
            },
        )
    )

def suppress_cdk_singletons(scope: IConstruct, singleton_ids: Iterable[str]) -> None:
    """Apply ``CDK_LAMBDA_SUPPRESSIONS`` to any CDK-managed singletons present under ``scope``.

    Resolves each ID via ``node.try_find_child`` rather than an absolute path
    string so suppressions survive being nested in a ``cdk.Stage``. Missing IDs
    are tolerated — some singletons only appear when the construct that needs
    them is instantiated.
    """
    for singleton_id in singleton_ids:
        singleton = scope.node.try_find_child(singleton_id)
        if singleton is not None:
            NagSuppressions.add_resource_suppressions(
                cast(Construct, singleton),
                CDK_LAMBDA_SUPPRESSIONS,
                apply_to_children=True,
            )