[ENG-2790] Fix broken p0 kubeconfig command (#129)

This PR unbreaks the `p0 kubeconfig` command. Some values and structures
got changed around in the k8s integration config and responses from the
backend:
- Required AWS IDC metadata is no longer found in the `eksGenerated`
object; it's now in `awsResourcePermission`
- AWS account ID is no longer stored separately in the k8s integration
configuration, and now is extracted from the EKS cluster ARN
- `provider` was renamed `hosting` in the k8s integration configuration

It also adds some minor quality-of-live improvements to the `p0
kubeconfig` command, such as utilizing `spinUntil()` for some of the
longer `await`s, and adds a more robust and reusable simple AWS ARN
parser.

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@p0security/cli",
-  "version": "0.11.2",
+  "version": "0.11.3",
   "description": "Execute infra CLI commands with P0 grants",
   "main": "index.ts",
   "repository": {

src/commands/kubeconfig.ts

@@ -12,7 +12,8 @@ import { retryWithSleep } from "../common/retry";
 import { AnsiSgr } from "../drivers/ansi";
 import { authenticate } from "../drivers/auth";
 import { guard } from "../drivers/firestore";
-import { print2 } from "../drivers/stdio";
+import { print2, spinUntil } from "../drivers/stdio";
+import { parseArn } from "../plugins/aws/utils";
 import {
   awsCloudAuth,
   profileName,
@@ -88,12 +89,13 @@ const kubeconfigAction = async (
     throw "Required dependencies are missing; please try again after installing them, or check that they are available on the PATH.";
   }
 
+  // No spinUntil(); there is one inside requestAccessToCluster() if needed
   const request = await requestAccessToCluster(authn, args, clusterId, role);
 
   const awsAuth = await awsCloudAuth(
     authn,
     awsAccountId,
-    request.generated,
+    request,
     awsLoginType
   );
 
@@ -121,11 +123,14 @@ const kubeconfigAction = async (
   try {
     // Federated access especially sometimes takes some time to propagate, so
     // retry for up to 20 seconds just in case it takes a while.
-    const awsResult = await retryWithSleep(
-      async () => await exec("aws", updateKubeconfigArgs, { check: true }),
-      () => true,
-      8,
-      2500
+    const awsResult = await spinUntil(
+      "Waiting for AWS resources to be provisioned and updating kubeconfig for EKS",
+      retryWithSleep(
+        async () => await exec("aws", updateKubeconfigArgs, { check: true }),
+        () => true,
+        8,
+        2500
+      )
     );
     print2(awsResult.stdout);
   } catch (error: any) {
@@ -233,22 +238,11 @@ const validateResourceArg = (resource: string): void => {
 const extractClusterNameAndRegion = (clusterArn: string) => {
   const INVALID_ARN_MSG = `Invalid EKS cluster ARN: ${clusterArn}`;
   // Example EKS cluster ARN: arn:aws:eks:us-west-2:123456789012:cluster/my-testing-cluster
-  const parts = clusterArn.split(":");
+  const arn = parseArn(clusterArn);
+  const { region: clusterRegion, resource: resourceStr } = arn;
+  const [resourceType, clusterName] = resourceStr.split("/");
 
-  if (parts.length < 6 || !parts[3] || !parts[5]) {
-    throw INVALID_ARN_MSG;
-  }
-
-  const clusterRegion = parts[3];
-  const resource = parts[5].split("/");
-
-  if (resource[0] !== "cluster") {
-    throw INVALID_ARN_MSG;
-  }
-
-  const clusterName = resource[1];
-
-  if (!clusterName) {
+  if (resourceType !== "cluster" || !clusterName || !clusterRegion) {
     throw INVALID_ARN_MSG;
   }
 

src/plugins/aws/__tests__/utils.test.ts

@@ -0,0 +1,95 @@
+/** Copyright © 2024-present P0 Security
+
+This file is part of @p0security/cli
+
+@p0security/cli is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, version 3 of the License.
+
+@p0security/cli is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License along with @p0security/cli. If not, see <https://www.gnu.org/licenses/>.
+**/
+import { parseArn } from "../utils";
+
+describe("parseArn() function", () => {
+  it.each([
+    "badarn:aws:ec2:us-east-1:123456789012:vpc/vpc-0e9801d129EXAMPLE", // Bad prefix
+    ":aws:ec2:us-east-1:123456789012:vpc/vpc-0e9801d129EXAMPLE", // Missing prefix
+    "arn:aws:ec2:us-east-1:123456789012", // Too few elements
+  ])('Raises an "Invalid ARN" error', (arn) => {
+    expect(() => parseArn(arn)).toThrow("Invalid AWS ARN");
+  });
+
+  it("Parses a valid ARN with all fields correctly", () => {
+    const arn = "arn:aws:ec2:us-east-1:123456789012:vpc/vpc-0e9801d129EXAMPLE";
+
+    const parsed = parseArn(arn);
+
+    expect(parsed).toEqual({
+      partition: "aws",
+      service: "ec2",
+      region: "us-east-1",
+      accountId: "123456789012",
+      resource: "vpc/vpc-0e9801d129EXAMPLE",
+    });
+  });
+
+  it("Parses a valid ARN with colons in the resource correctly", () => {
+    // Note: This is not the format we would expect an EKS ARN to actually be in (it should
+    // use a / instead of a : in the resource); this is just for unit testing purposes.
+    const arn = "arn:aws:eks:us-west-2:123456789012:cluster:my-testing-cluster";
+
+    const parsed = parseArn(arn);
+
+    expect(parsed).toEqual({
+      partition: "aws",
+      service: "eks",
+      region: "us-west-2",
+      accountId: "123456789012",
+      resource: "cluster:my-testing-cluster",
+    });
+  });
+
+  it("Parses a valid ARN with no region correctly", () => {
+    const arn = "arn:aws:iam::123456789012:user/johndoe";
+
+    const parsed = parseArn(arn);
+
+    expect(parsed).toEqual({
+      partition: "aws",
+      service: "iam",
+      region: "",
+      accountId: "123456789012",
+      resource: "user/johndoe",
+    });
+  });
+
+  it("Parses a valid ARN with no account ID correctly", () => {
+    // Note: This is not a valid SNS ARN; they would ordinarily have an account ID. This is
+    // just for unit testing purposes.
+    const arn = "arn:aws-us-gov:sns:us-east-1::example-sns-topic-name";
+
+    const parsed = parseArn(arn);
+
+    expect(parsed).toEqual({
+      partition: "aws-us-gov",
+      service: "sns",
+      region: "us-east-1",
+      accountId: "",
+      resource: "example-sns-topic-name",
+    });
+  });
+
+  it("Parses a valid ARN with no region or account ID correctly", () => {
+    const arn = "arn:aws-cn:s3:::my-corporate-bucket";
+
+    const parsed = parseArn(arn);
+
+    expect(parsed).toEqual({
+      partition: "aws-cn",
+      service: "s3",
+      region: "",
+      accountId: "",
+      resource: "my-corporate-bucket",
+    });
+  });
+});

src/plugins/aws/utils.ts

@@ -0,0 +1,65 @@
+/** Copyright © 2024-present P0 Security
+
+This file is part of @p0security/cli
+
+@p0security/cli is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, version 3 of the License.
+
+@p0security/cli is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License along with @p0security/cli. If not, see <https://www.gnu.org/licenses/>.
+**/
+const ARN_PATTERN =
+  /^arn:(aws|aws-us-gov|aws-cn|aws-iso|aws-iso-b):([^:]*):([^:]*):([^:]*):(.*)$/;
+
+/**
+ * Parses out Amazon Resource Names (ARNs) from AWS into their components. Note
+ * that not all components are present in all ARNs (depending on the service;
+ * for example, S3 ARNs don't have a region or account ID), and the final
+ * component of the ARN (`resource`) may contain its own internal structure that
+ * is also service-dependent and which may also include colons. In particular,
+ * quoting the Amazon docs: "Be aware that the ARNs for some resources omit the
+ * Region, the account ID, or both the Region and the account ID."
+ *
+ * @param arn The ARN to parse as a string.
+ * @return A structure representing the components of the ARN. All fields will
+ * be defined, but some may be empty strings if they are not present in the ARN.
+ */
+export const parseArn = (
+  arn: string
+): {
+  partition: string;
+  service: string;
+  region: string;
+  accountId: string;
+  resource: string;
+} => {
+  // Reference: https://docs.aws.amazon.com/IAM/latest/UserGuide/reference-arns.html
+  const invalidArnMsg = `Invalid AWS ARN: ${arn}`;
+  const match = arn.match(ARN_PATTERN);
+
+  if (!match) {
+    throw invalidArnMsg;
+  }
+
+  const [_, partition, service, region, accountId, resource] = match;
+
+  // We know these are all defined based on the regex, but TypeScript doesn't.
+  // Empty string is okay, so explicitly check for undefined.
+  if (
+    partition === undefined ||
+    service === undefined ||
+    accountId === undefined ||
+    region === undefined ||
+    resource === undefined
+  ) {
+    throw invalidArnMsg;
+  }
+
+  return {
+    partition,
+    service,
+    region,
+    accountId,
+    resource,
+  };
+};

src/plugins/kubeconfig/index.ts

@@ -12,15 +12,16 @@ import { KubeconfigCommandArgs } from "../../commands/kubeconfig";
 import { waitForProvisioning } from "../../commands/shared";
 import { request } from "../../commands/shared/request";
 import { doc } from "../../drivers/firestore";
-import { print2 } from "../../drivers/stdio";
+import { spinUntil } from "../../drivers/stdio";
 import { Authn } from "../../types/identity";
 import { Request } from "../../types/request";
 import { assertNever } from "../../util";
 import { getAwsConfig } from "../aws/config";
 import { assumeRoleWithIdc } from "../aws/idc";
 import { AwsCredentials } from "../aws/types";
+import { parseArn } from "../aws/utils";
 import { assumeRoleWithOktaSaml } from "../okta/aws";
-import { K8sConfig, K8sGenerated, K8sPermissionSpec } from "./types";
+import { K8sConfig, K8sPermissionSpec } from "./types";
 import { getDoc } from "firebase/firestore";
 import { pick } from "lodash";
 import yargs from "yargs";
@@ -46,20 +47,21 @@ export const getAndValidateK8sIntegration = async (
     throw `Cluster with ID ${clusterId} not found`;
   }
 
-  if (config.state !== "installed" || config.provider.type !== "aws") {
+  if (config.state !== "installed") {
     throw `Cluster with ID ${clusterId} is not installed`;
   }
 
-  const { provider } = config;
-  const { accountId: awsAccountId, clusterArn: awsClusterArn } = provider;
+  const { hosting } = config;
 
-  if (!awsAccountId || !awsClusterArn) {
+  if (hosting.type !== "aws") {
     throw (
       `This command currently only supports AWS EKS clusters, and ${clusterId} is not configured as one.\n` +
       "You can request access to the cluster using the `p0 request k8s` command."
     );
   }
 
+  const { arn: awsClusterArn } = hosting;
+  const { accountId: awsAccountId } = parseArn(awsClusterArn);
   const { config: awsConfig } = await getAwsConfig(authn, awsAccountId);
   const { login: awsLogin } = awsConfig;
 
@@ -109,14 +111,12 @@ export const requestAccessToCluster = async (
   if (!response) {
     throw "Did not receive access ID from server";
   }
-  const { id, isPreexisting } = response;
-  if (!isPreexisting) {
-    print2(
-      "Waiting for access to be provisioned. This may take up to a minute."
-    );
-  }
+  const { id } = response;
 
-  return await waitForProvisioning<K8sPermissionSpec>(authn, id);
+  return await spinUntil(
+    "Waiting for access to be provisioned. This may take up to a minute.",
+    waitForProvisioning<K8sPermissionSpec>(authn, id)
+  );
 };
 
 export const profileName = (eksCluterName: string): string =>
@@ -125,23 +125,27 @@ export const profileName = (eksCluterName: string): string =>
 export const awsCloudAuth = async (
   authn: Authn,
   awsAccountId: string,
-  generated: K8sGenerated,
+  request: Request<K8sPermissionSpec>,
   loginType: "federated" | "idc"
 ): Promise<AwsCredentials> => {
+  const { permission, generated } = request;
   const { eksGenerated } = generated;
-  const { name, idc } = eksGenerated;
+  const { name } = eksGenerated;
 
   switch (loginType) {
-    case "idc":
-      if (!idc) {
+    case "idc": {
+      const { idcId, idcRegion } = permission.awsResourcePermission ?? {};
+
+      if (!idcId || !idcRegion) {
         throw "AWS is configured to use Identity Center, but IDC information wasn't received in the request.";
       }
 
       return await assumeRoleWithIdc({
         accountId: awsAccountId,
         permissionSet: name,
-        idc,
+        idc: { id: idcId, region: idcRegion },
       });
+    }
     case "federated":
       return await assumeRoleWithOktaSaml(authn, {
         accountId: awsAccountId,

src/plugins/kubeconfig/types.ts

@@ -17,9 +17,7 @@ export type K8sClusterConfig = {
   isProxy: boolean;
   token: string;
   publicJwk?: string; // only present for proxy installs
-  provider:
-    | { type: "aws"; clusterArn: string; accountId: string }
-    | { type: "email" };
+  hosting: { type: "aws"; arn: string } | { type: "azure" } | { type: "gcp" };
   state: string;
 };
 
@@ -42,15 +40,16 @@ export type K8sResourcePermission = {
   role: string;
   clusterId: string;
   type: "resource";
+  awsResourcePermission?: {
+    idcRegion?: string;
+    idcId?: string;
+  };
 };
 
 export type K8sGenerated = {
   eksGenerated: {
     // For IDC, the name of the permission set. For Federated, the name of the assumed role
     name: string;
-
-    // Only present for IDC; the ID and region of the IDC installation
-    idc?: { id: string; region: string };
   };
   role: string; // The name of the generated role in k8s itself
 };