From c545bfe2a3ccb53fa5ae2eb725a1696677703c0a Mon Sep 17 00:00:00 2001
From: Kaizen Conroy <36202692+kaizencc@users.noreply.github.com>
Date: Thu, 28 Apr 2022 05:01:39 -0400
Subject: [PATCH] feat(iam): add convenience method `inOrganization` to
 ArnPrincipal (#20109)

Add a convenience method to ArnPrincipal. ArnPrincipal is extended by AccountPrincipal and AnyPrincipal, which are the only principals that could reasonably want to add a condition on organization.

```ts
new AccountPrincipal('123456789012').inOrganization('o-xxxxxxxxxx');
```

Related: https://github.com/aws/aws-cdk/pull/19975#discussion_r857385168. With this method, the API in #19975 will look like:

```ts
fn.grantInvoke(new AccountPrincipal('123456789012').inOrganization('o-xxxxxxxxxx');
```

Which is really slick!

----

### All Submissions:

* [x] Have you followed the guidelines in our [Contributing guide?](https://github.com/aws/aws-cdk/blob/master/CONTRIBUTING.md)

### Adding new Unconventional Dependencies:

* [ ] This PR adds new unconventional dependencies following the process described [here](https://github.com/aws/aws-cdk/blob/master/CONTRIBUTING.md/#adding-new-unconventional-dependencies)

### New Features

* [ ] Have you added the new feature to an [integration test](https://github.com/aws/aws-cdk/blob/master/INTEGRATION_TESTS.md)?
	* [ ] Did you use `yarn integ` to deploy the infrastructure and generate the snapshot (i.e. `yarn integ` without `--dry-run`)?

*By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license*
---
 packages/@aws-cdk/aws-iam/lib/principals.ts   | 14 +++++-
 .../@aws-cdk/aws-iam/test/principals.test.ts  | 49 +++++++++++++++++++
 2 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/packages/@aws-cdk/aws-iam/lib/principals.ts b/packages/@aws-cdk/aws-iam/lib/principals.ts
index a52e8d1e0dda2..8cb94c33a0b1d 100644
--- a/packages/@aws-cdk/aws-iam/lib/principals.ts
+++ b/packages/@aws-cdk/aws-iam/lib/principals.ts
@@ -368,6 +368,18 @@ export class ArnPrincipal extends PrincipalBase {
   public toString() {
     return `ArnPrincipal(${this.arn})`;
   }
+
+  /**
+   * A convenience method for adding a condition that the principal is part of the specified
+   * AWS Organization.
+   */
+  public inOrganization(organizationId: string) {
+    return this.withConditions({
+      StringEquals: {
+        'aws:PrincipalOrgID': organizationId,
+      },
+    });
+  }
 }
 
 /**
@@ -397,7 +409,7 @@ export interface ServicePrincipalOpts {
   /**
    * The region in which the service is operating.
    *
-   * @default the current Stack's region.
+   * @default - the current Stack's region.
    * @deprecated You should not need to set this. The stack's region is always correct.
    */
   readonly region?: string;
diff --git a/packages/@aws-cdk/aws-iam/test/principals.test.ts b/packages/@aws-cdk/aws-iam/test/principals.test.ts
index 7a07a50e80fb9..34206540def53 100644
--- a/packages/@aws-cdk/aws-iam/test/principals.test.ts
+++ b/packages/@aws-cdk/aws-iam/test/principals.test.ts
@@ -245,6 +245,55 @@ test('PrincipalWithConditions inherits principalAccount from AccountPrincipal ',
   expect(principalWithConditions.principalAccount).toStrictEqual('123456789012');
 });
 
+test('AccountPrincipal can specify an organization', () => {
+  // GIVEN
+  const stack = new Stack();
+
+  // WHEN
+  const pol = new iam.PolicyDocument({
+    statements: [
+      new iam.PolicyStatement({
+        actions: ['service:action'],
+        resources: ['*'],
+        principals: [
+          new iam.AccountPrincipal('123456789012').inOrganization('o-xxxxxxxxxx'),
+        ],
+      }),
+    ],
+  });
+
+  // THEN
+  expect(stack.resolve(pol)).toEqual({
+    Statement: [
+      {
+        Action: 'service:action',
+        Effect: 'Allow',
+        Principal: {
+          AWS: {
+            'Fn::Join': [
+              '',
+              [
+                'arn:',
+                {
+                  Ref: 'AWS::Partition',
+                },
+                ':iam::123456789012:root',
+              ],
+            ],
+          },
+        },
+        Condition: {
+          StringEquals: {
+            'aws:PrincipalOrgID': 'o-xxxxxxxxxx',
+          },
+        },
+        Resource: '*',
+      },
+    ],
+    Version: '2012-10-17',
+  });
+});
+
 test('ServicePrincipal in agnostic stack generates lookup table', () => {
   // GIVEN
   const stack = new Stack();