ENG-2884 Add multi-host support

package.json

@@ -70,6 +70,7 @@
   },
   "scripts": {
     "build": "tsc && cp -r public dist/",
+    "test:unit": "NODE_ENV=unit jest --color",
     "clean": "rm -f tsconfig.tsbuildinfo && rm -rf dist/",
     "lint": "yarn prettier --check . &&  yarn run eslint --max-warnings 0 .",
     "p0": "node --no-deprecation ./p0",

src/commands/__tests__/__snapshots__/login.test.ts.snap

@@ -1,9 +1,9 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
-exports[`login organization exists should write the user's identity to the file system 1`] = `
+exports[`login organization exists should write the user's identity & config to the file system 1`] = `
 [
   [
-    "/path/to/home/.p0",
+    "/dummy/identity/file/path",
     "{
   "credential": {
     "access_token": "test-access-token",

src/commands/__tests__/login.test.ts

@@ -18,7 +18,11 @@ jest.spyOn(Date, "now").mockReturnValue(1.6e12);
 jest.mock("fs/promises");
 jest.mock("../../drivers/auth", () => ({
   ...jest.requireActual("../../drivers/auth"),
-  IDENTITY_FILE_PATH: "/path/to/home/.p0",
+  IDENTITY_FILE_PATH: "/dummy/identity/file/path",
+}));
+jest.mock("../../drivers/config", () => ({
+  ...jest.requireActual("../../drivers/config"),
+  saveConfig: jest.fn(),
 }));
 jest.mock("../../drivers/stdio");
 jest.mock("../../plugins/login");
@@ -34,6 +38,7 @@ describe("login", () => {
       `"Could not find organization"`
     );
   });
+
   it("should print a friendly error if unsupported login", async () => {
     mockGetDoc({
       slug: "test-org",
@@ -44,6 +49,7 @@ describe("login", () => {
       `"Unsupported login for your organization"`
     );
   });
+
   describe("organization exists", () => {
     let credentialData: string = "";
     mockReadFile.mockImplementation(async () =>
@@ -60,27 +66,33 @@ describe("login", () => {
           },
         })
     );
+
     beforeEach(() => {
       credentialData = "";
       jest.clearAllMocks();
+
       mockGetDoc({
         slug: "test-org",
         tenantId: "test-tenant",
         ssoProvider: "google",
       });
     });
+
     it("should call the provider's login function", async () => {
       await login({ org: "test-org" });
       expect(pluginLoginMap.google).toHaveBeenCalled();
     });
-    it("should write the user's identity to the file system", async () => {
+
+    it("should write the user's identity & config to the file system", async () => {
       await login({ org: "test-org" });
       expect(mockWriteFile.mock.calls).toMatchSnapshot();
     });
+
     it("validates authentication", async () => {
       await login({ org: "test-org" });
       expect((signInWithCredential as jest.Mock).mock.calls).toMatchSnapshot();
     });
+
     it("returns an error message if firebase cannot determine the user's email", async () => {
       mockSignInWithCredential.mockResolvedValueOnce({
         user: {},

src/commands/allow.ts

@@ -10,7 +10,7 @@ You should have received a copy of the GNU General Public License along with @p0
 **/
 import { fetchCommand } from "../drivers/api";
 import { authenticate } from "../drivers/auth";
-import { guard } from "../drivers/firestore";
+import { fsShutdownGuard } from "../drivers/firestore";
 import { print2 } from "../drivers/stdio";
 import { AllowResponse } from "../types/allow";
 import { Authn } from "../types/identity";
@@ -37,7 +37,7 @@ export const allowCommand = (yargs: yargs.Argv) =>
     "allow [arguments..]",
     "Create standing access for a resource",
     allowArgs,
-    guard(allow)
+    fsShutdownGuard(allow)
   );
 
 export const allow = async (

src/commands/aws/role.ts

@@ -10,7 +10,7 @@ You should have received a copy of the GNU General Public License along with @p0
 **/
 import { parseXml } from "../../common/xml";
 import { authenticate } from "../../drivers/auth";
-import { guard } from "../../drivers/firestore";
+import { fsShutdownGuard } from "../../drivers/firestore";
 import { print1, print2 } from "../../drivers/stdio";
 import { getAwsConfig } from "../../plugins/aws/config";
 import { AwsFederatedLogin, AwsItem } from "../../plugins/aws/types";
@@ -29,7 +29,7 @@ export const role = (yargs: yargs.Argv<{ account: string | undefined }>) =>
         "List available AWS roles",
         identity,
         // TODO: select based on uidLocation
-        guard(oktaAwsListRoles)
+        fsShutdownGuard(oktaAwsListRoles)
       )
       .command(
         "assume <role>",
@@ -41,7 +41,7 @@ export const role = (yargs: yargs.Argv<{ account: string | undefined }>) =>
             describe: "An AWS role name",
           }),
         // TODO: select based on uidLocation
-        guard(oktaAwsAssumeRole)
+        fsShutdownGuard(oktaAwsAssumeRole)
       )
       .demandCommand(1)
   );

src/commands/grant.ts

@@ -8,7 +8,7 @@ This file is part of @p0security/cli
 
 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 { guard } from "../drivers/firestore";
+import { fsShutdownGuard } from "../drivers/firestore";
 import { request, requestArgs } from "./shared/request";
 import yargs from "yargs";
 
@@ -17,5 +17,5 @@ export const grantCommand = (yargs: yargs.Argv) =>
     "grant [arguments..]",
     "Grant access to another identity",
     requestArgs,
-    guard(request("grant"))
+    fsShutdownGuard(request("grant"))
   );

src/commands/kubeconfig.ts

@@ -11,7 +11,7 @@ You should have received a copy of the GNU General Public License along with @p0
 import { retryWithSleep } from "../common/retry";
 import { AnsiSgr } from "../drivers/ansi";
 import { authenticate } from "../drivers/auth";
-import { guard } from "../drivers/firestore";
+import { fsShutdownGuard } from "../drivers/firestore";
 import { print2, spinUntil } from "../drivers/stdio";
 import { parseArn } from "../plugins/aws/utils";
 import {
@@ -65,7 +65,7 @@ export const kubeconfigCommand = (yargs: yargs.Argv) =>
           describe:
             "Requested duration for access (format like '10 minutes', '2 hours', '5 days', or '1 week')",
         }),
-    guard(kubeconfigAction)
+    fsShutdownGuard(kubeconfigAction)
   );
 
 const kubeconfigAction = async (

src/commands/login.ts

@@ -11,9 +11,15 @@ You should have received a copy of the GNU General Public License along with @p0
 import {
   IDENTITY_CACHE_PATH,
   IDENTITY_FILE_PATH,
-  authenticate,
+  loadCredentials,
 } from "../drivers/auth";
-import { doc, guard } from "../drivers/firestore";
+import { saveConfig } from "../drivers/config";
+import { bootstrapConfig } from "../drivers/env";
+import {
+  authenticateToFirebase,
+  fsShutdownGuard,
+  publicDoc,
+} from "../drivers/firestore";
 import { print2 } from "../drivers/stdio";
 import { pluginLoginMap } from "../plugins/login";
 import { TokenResponse } from "../types/oidc";
@@ -32,10 +38,14 @@ export const login = async (
   args: { org: string },
   options?: { skipAuthenticate?: boolean }
 ) => {
-  const orgDoc = await getDoc<RawOrgData, object>(doc(`orgs/${args.org}`));
+  const orgDoc = await getDoc<RawOrgData, object>(
+    publicDoc(`orgs/${args.org}`)
+  );
   const orgData = orgDoc.data();
   if (!orgData) throw "Could not find organization";
 
+  await saveConfig(orgData.config ?? bootstrapConfig);
+
   const orgWithSlug: OrgData = { ...orgData, slug: args.org };
 
   const plugin = orgWithSlug?.ssoProvider;
@@ -49,7 +59,10 @@ export const login = async (
   await writeIdentity(orgWithSlug, tokenResponse);
 
   // validate auth
-  if (!options?.skipAuthenticate) await authenticate({ noRefresh: true });
+  if (!options?.skipAuthenticate) {
+    const identity = await loadCredentials({ noRefresh: true });
+    await authenticateToFirebase(identity);
+  }
 
   print2(`You are now logged in, and can use the p0 CLI.`);
 };
@@ -95,5 +108,5 @@ export const loginCommand = (yargs: yargs.Argv) =>
         type: "string",
         describe: "Your P0 organization ID",
       }),
-    guard(login)
+    fsShutdownGuard(login)
   );

src/commands/ls.ts

@@ -11,7 +11,7 @@ You should have received a copy of the GNU General Public License along with @p0
 import { AnsiSgr } from "../drivers/ansi";
 import { fetchCommand } from "../drivers/api";
 import { authenticate } from "../drivers/auth";
-import { guard } from "../drivers/firestore";
+import { fsShutdownGuard } from "../drivers/firestore";
 import { print2, print1, spinUntil } from "../drivers/stdio";
 import { max, orderBy } from "lodash";
 import pluralize from "pluralize";
@@ -45,7 +45,7 @@ export const lsCommand = (yargs: yargs.Argv) =>
     "ls [arguments..]",
     "List request-command arguments",
     lsArgs,
-    guard(ls)
+    fsShutdownGuard(ls)
   );
 
 const ls = async (

src/commands/request.ts

@@ -8,7 +8,7 @@ This file is part of @p0security/cli
 
 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 { guard } from "../drivers/firestore";
+import { fsShutdownGuard } from "../drivers/firestore";
 import { request, requestArgs } from "./shared/request";
 import yargs from "yargs";
 
@@ -17,5 +17,5 @@ export const requestCommand = (yargs: yargs.Argv) =>
     "request [arguments..]",
     "Manually request permissions on a resource",
     requestArgs,
-    guard(request("request"))
+    fsShutdownGuard(request("request"))
   );

src/commands/scp.ts

@@ -9,7 +9,7 @@ This file is part of @p0security/cli
 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 { authenticate } from "../drivers/auth";
-import { guard } from "../drivers/firestore";
+import { fsShutdownGuard } from "../drivers/firestore";
 import { sshOrScp } from "../plugins/ssh";
 import { SshRequest, SupportedSshProviders } from "../types/ssh";
 import { prepareRequest, ScpCommandArgs } from "./shared/ssh";
@@ -69,7 +69,7 @@ export const scpCommand = (yargs: yargs.Argv) =>
   The '--' argument must be specified between P0-specific args on the left and SCP_ARGS on the right.`
         ),
 
-    guard(scpAction)
+    fsShutdownGuard(scpAction)
   );
 
 /** Transfers files between a local and remote hosts using SSH.

src/commands/ssh.ts

@@ -9,7 +9,7 @@ This file is part of @p0security/cli
 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 { authenticate } from "../drivers/auth";
-import { guard } from "../drivers/firestore";
+import { fsShutdownGuard } from "../drivers/firestore";
 import { sshOrScp } from "../plugins/ssh";
 import { SshCommandArgs, prepareRequest } from "./shared/ssh";
 import yargs from "yargs";
@@ -69,7 +69,7 @@ export const sshCommand = (yargs: yargs.Argv) =>
   $ p0 ssh example-instance --provider gcloud -- -NR '*:8080:localhost:8088' -o 'GatewayPorts yes'`
         ),
 
-    guard(sshAction)
+    fsShutdownGuard(sshAction)
   );
 
 /** Connect to an SSH backend

src/drivers/api.ts

@@ -8,13 +8,12 @@ This file is part of @p0security/cli
 
 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 { config } from "../drivers/env";
 import { Authn } from "../types/identity";
+import { getTenantConfig } from "./config";
 import * as path from "node:path";
 import yargs from "yargs";
 
-const tenantUrl = (tenant: string) => `${config.appUrl}/o/${tenant}`;
-
+const tenantUrl = (tenant: string) => `${getTenantConfig().appUrl}/o/${tenant}`;
 const commandUrl = (tenant: string) => `${tenantUrl(tenant)}/command/`;
 
 export const fetchCommand = async <T>(

src/drivers/auth.ts

@@ -11,13 +11,9 @@ You should have received a copy of the GNU General Public License along with @p0
 import { login } from "../commands/login";
 import { Authn, Identity } from "../types/identity";
 import { P0_PATH } from "../util";
-import { auth } from "./firestore";
+import { loadConfig } from "./config";
+import { authenticateToFirebase, initializeFirebase } from "./firestore";
 import { print2 } from "./stdio";
-import {
-  OAuthProvider,
-  SignInMethod,
-  signInWithCredential,
-} from "firebase/auth";
 import * as fs from "fs/promises";
 import * as path from "path";
 
@@ -96,24 +92,11 @@ export const loadCredentials = async (options?: {
 export const authenticate = async (options?: {
   noRefresh?: boolean;
 }): Promise<Authn> => {
-  const identity = await loadCredentials(options);
-  const { credential } = identity;
+  await loadConfig();
+  initializeFirebase();
 
-  // TODO: Move to map lookup
-  const provider = new OAuthProvider(
-    identity.org.ssoProvider === "google"
-      ? SignInMethod.GOOGLE
-      : identity.org.providerId
-  );
-  const firebaseCredential = provider.credential({
-    accessToken: credential.access_token,
-    idToken: credential.id_token,
-  });
-  auth.tenantId = identity.org.tenantId;
-  const userCredential = await signInWithCredential(auth, firebaseCredential);
-  if (!userCredential?.user?.email) {
-    throw "Can not sign in: this user has previously signed in with a different identity provider.\nPlease contact support@p0.dev to enable this user.";
-  }
+  const identity = await loadCredentials(options);
+  const userCredential = await authenticateToFirebase(identity);
 
   return { userCredential, identity };
 };

src/drivers/config.ts

@@ -0,0 +1,35 @@
+/** 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 { Config } from "../types/org";
+import { P0_PATH } from "../util";
+import { print2 } from "./stdio";
+import fs from "fs/promises";
+import path from "path";
+
+export const CONFIG_FILE_PATH = path.join(P0_PATH, "config.json");
+
+let tenantConfig: Config;
+
+export function getTenantConfig(): Config {
+  return tenantConfig;
+}
+
+export async function saveConfig(config: Config) {
+  print2(`Saving config to ${CONFIG_FILE_PATH}.`);
+  const dir = path.dirname(CONFIG_FILE_PATH);
+  await fs.mkdir(dir, { recursive: true });
+  await fs.writeFile(CONFIG_FILE_PATH, JSON.stringify(config), { mode: "600" });
+}
+
+export async function loadConfig() {
+  const buffer = await fs.readFile(CONFIG_FILE_PATH);
+  tenantConfig = JSON.parse(buffer.toString());
+}

src/drivers/env.ts

@@ -14,7 +14,7 @@ dotenv.config();
 
 const { env } = process;
 
-export const config = {
+export const bootstrapConfig = {
   fs: {
     // Falls back to public production Firestore credentials
     apiKey: env.P0_FS_API_KEY ?? "AIzaSyCaL-Ik_l_5tdmgNUNZ4Nv6NuR4o5_PPfs",