Initial implementation of Azure SSH (#143)

This PR implements the initial version of Azure SSH for the P0 CLI. It
is fully functional and works with P0 instances configured with an Azure
iam-write integration and an Azure SSH integration.

For more details, please see the PR (#143).

src/commands/scp.ts

@@ -84,6 +84,14 @@ const scpAction = async (args: yargs.ArgumentsCamelCase<ScpCommandArgs>) => {
     : [];
   args.sshOptions = sshOptions;
 
+  // TODO(ENG-3142): Azure SSH currently doesn't support specifying a port; throw an error if one is set.
+  if (
+    args.provider === "azure" &&
+    sshOptions.some((opt) => opt.startsWith("-P"))
+  ) {
+    throw "Azure SSH does not currently support specifying a port. SSH on the target VM must be listening on the default port 22.";
+  }
+
   const host = getHostIdentifier(args.source, args.destination);
 
   if (!host) {

src/commands/shared/ssh.ts

@@ -54,6 +54,17 @@ export type SshCommandArgs = BaseSshCommandArgs & {
 
 export type CommandArgs = ScpCommandArgs | SshCommandArgs;
 
+export type SshAdditionalSetup = {
+  /** A list of SSH configuration options, as would be used after '-o' in an SSH command */
+  sshOptions: string[];
+
+  /** The port to connect to, overriding the default */
+  port: string;
+
+  /** Perform any teardown required after the SSH command exits but before terminating the P0 CLI */
+  teardown: () => Promise<void>;
+};
+
 export const SSH_PROVIDERS: Record<
   SupportedSshProvider,
   SshProvider<any, any, any, any>
@@ -81,6 +92,7 @@ const validateSshInstall = async (
       value.state == "installed" &&
       providersToCheck.some((prefix) => key.startsWith(prefix))
   );
+
   if (items.length === 0) {
     throw "This organization is not configured for SSH access via the P0 CLI";
   }
@@ -139,9 +151,6 @@ export const provisionRequest = async (
     authn,
     id
   );
-  if (provisionedRequest.permission.publicKey !== publicKey) {
-    throw "Public key mismatch. Please revoke the request and try again.";
-  }
 
   return { provisionedRequest, publicKey, privateKey };
 };
@@ -156,9 +165,17 @@ export const prepareRequest = async (
     throw "Server did not return a request id. Please contact support@p0.dev for assistance.";
   }
 
-  const { provisionedRequest } = result;
+  const { provisionedRequest, publicKey } = result;
 
   const sshProvider = SSH_PROVIDERS[provisionedRequest.permission.provider];
+
+  if (
+    sshProvider.validateSshKey &&
+    !sshProvider.validateSshKey(provisionedRequest, publicKey)
+  ) {
+    throw "Public key mismatch. Please revoke the request and try again.";
+  }
+
   await sshProvider.ensureInstall();
 
   const cliRequest = await pluginToCliRequest(provisionedRequest, {

src/commands/ssh.ts

@@ -51,7 +51,7 @@ export const sshCommand = (yargs: yargs.Argv) =>
         .option("provider", {
           type: "string",
           describe: "The cloud provider where the instance is hosted",
-          choices: ["aws", "gcloud"],
+          choices: ["aws", "azure", "gcloud"],
         })
         .option("debug", {
           type: "boolean",
@@ -89,6 +89,14 @@ const sshAction = async (args: yargs.ArgumentsCamelCase<SshCommandArgs>) => {
     : [];
   args.sshOptions = sshOptions;
 
+  // TODO(ENG-3142): Azure SSH currently doesn't support specifying a port; throw an error if one is set.
+  if (
+    args.provider === "azure" &&
+    sshOptions.some((opt) => opt.startsWith("-p"))
+  ) {
+    throw "Azure SSH does not currently support specifying a port. SSH on the target VM must be listening on the default port 22.";
+  }
+
   const { request, privateKey, sshProvider } = await prepareRequest(
     authn,
     args,

src/plugins/aws/ssh.ts

@@ -72,6 +72,9 @@ export const awsSshProvider: SshProvider<
         : throwAssertNever(config.login);
   },
 
+  validateSshKey: (request, publicKey) =>
+    request.permission.publicKey === publicKey,
+
   ensureInstall: async () => {
     if (!(await ensureSsmInstall())) {
       throw "Please try again after installing the required AWS utilities";

src/plugins/azure/auth.ts

@@ -0,0 +1,18 @@
+/** 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 { exec } from "../../util";
+
+export const azLogin = async (subscriptionId: string) => {
+  await exec("az", ["login"], { check: true });
+  await exec("az", ["account", "set", "--subscription", subscriptionId], {
+    check: true,
+  });
+};

src/plugins/azure/keygen.ts

@@ -0,0 +1,49 @@
+/** 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 { print2 } from "../../drivers/stdio";
+import { exec } from "../../util";
+import path from "node:path";
+import tmp from "tmp-promise";
+
+// We pass in the name of the certificate file to generate
+export const AD_CERT_FILENAME = "p0cli-azure-ad-ssh-cert.pub";
+
+// The `az ssh cert` command manages key generation, and generates SSH RSA keys with the standard names
+export const AD_SSH_KEY_PRIVATE = "id_rsa";
+
+export const createTempDirectoryForKeys = async (): Promise<{
+  path: string;
+  cleanup: () => Promise<void>;
+}> => {
+  // unsafeCleanup lets us delete the directory even if there are still files in it, which is fine since the
+  // files are no longer needed once we've authenticated to the remote system.
+  const { path, cleanup } = await tmp.dir({
+    mode: 0o700,
+    prefix: "p0cli-",
+    unsafeCleanup: true,
+  });
+
+  return { path, cleanup };
+};
+
+export const generateSshKeyAndAzureAdCert = async (keyPath: string) => {
+  try {
+    await exec(
+      "az",
+      ["ssh", "cert", "--file", path.join(keyPath, AD_CERT_FILENAME)],
+      { check: true }
+    );
+  } catch (error: any) {
+    print2(error.stdout);
+    print2(error.stderr);
+    throw `Failed to generate Azure AD SSH certificate: ${error}`;
+  }
+};

src/plugins/azure/ssh.ts

@@ -9,24 +9,33 @@ 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 { SshProvider } from "../../types/ssh";
-import { exec } from "../../util";
-import { importSshKey } from "../google/ssh-key";
+import { azLogin } from "./auth";
 import { ensureAzInstall } from "./install";
-import { AzureSshPermissionSpec, AzureSshRequest } from "./types";
+import {
+  AD_CERT_FILENAME,
+  AD_SSH_KEY_PRIVATE,
+  createTempDirectoryForKeys,
+  generateSshKeyAndAzureAdCert,
+} from "./keygen";
+import { trySpawnBastionTunnel } from "./tunnel";
+import {
+  AzureLocalData,
+  AzureSshPermissionSpec,
+  AzureSshRequest,
+} from "./types";
+import path from "node:path";
 
 // TODO: Determine what this value should be for Azure
 const PROPAGATION_TIMEOUT_LIMIT_MS = 2 * 60 * 1000;
 
 export const azureSshProvider: SshProvider<
   AzureSshPermissionSpec,
-  { linuxUserName: string },
+  AzureLocalData,
   AzureSshRequest
 > = {
   // TODO: Natively support Azure login in P0 CLI
   cloudProviderLogin: async () => {
-    // Always invoke `az login` before each SSH access. This is needed because
-    // Azure permissions are only updated upon login.
-    await exec("az", ["login"]);
+    // Login is handled as part of setup() below
     return undefined;
   },
 
@@ -45,31 +54,81 @@ export const azureSshProvider: SshProvider<
 
   propagationTimeoutMs: PROPAGATION_TIMEOUT_LIMIT_MS,
 
-  // TODO: Implement
+  // TODO(ENG-3149): Implement sudo access checks here
   preTestAccessPropagationArgs: () => undefined,
 
-  // TODO: Determine if necessary
+  // Azure doesn't support ProxyCommand, as nice as that would be. Yet.
   proxyCommand: () => [],
 
   // TODO: Determine if necessary
   reproCommands: () => undefined,
 
-  // TODO: Placeholder
+  setup: async (request) => {
+    // TODO(ENG-3129): Does this specifically need to be the subscription ID for the Bastion?
+    await azLogin(request.subscriptionId); // Always re-login to Azure CLI
+
+    const { path: keyPath, cleanup: sshKeyPathCleanup } =
+      await createTempDirectoryForKeys();
+
+    const wrappedCreateCertAndTunnel = async () => {
+      try {
+        await generateSshKeyAndAzureAdCert(keyPath);
+        return await trySpawnBastionTunnel(request);
+      } catch (error: any) {
+        await sshKeyPathCleanup();
+        throw error;
+      }
+    };
+
+    const { killTunnel, tunnelLocalPort } = await wrappedCreateCertAndTunnel();
+
+    const sshPrivateKeyPath = path.join(keyPath, AD_SSH_KEY_PRIVATE);
+    const sshCertificateKeyPath = path.join(keyPath, AD_CERT_FILENAME);
+
+    const teardown = async () => {
+      await killTunnel();
+      await sshKeyPathCleanup();
+    };
+
+    return {
+      sshOptions: [
+        `IdentityFile ${sshPrivateKeyPath}`,
+        `CertificateFile ${sshCertificateKeyPath}`,
+        "IdentitiesOnly yes",
+
+        // Because we connect to the Azure Network Bastion tunnel via a local port instead of a ProxyCommand, every
+        // instance connected to will appear to `ssh` to be the same host but presenting a different host key (i.e.,
+        // `ssh` always connects to localhost but each VM will present its own host key), which will trigger MITM attack
+        // warnings. We disable host key checking to avoid this. This is ordinarily very dangerous, but in this case,
+        // security of the connection is ensured by the Azure Bastion Network tunnel, which utilizes HTTPS and thus has
+        // its own MITM protection.
+        "StrictHostKeyChecking no",
+        "UserKnownHostsFile /dev/null",
+      ],
+      port: tunnelLocalPort,
+      teardown,
+    };
+  },
+
   requestToSsh: (request) => ({
     type: "azure",
-    id: request.permission.resource.instanceId,
+    id: "localhost",
+    ...request.cliLocalData,
     instanceId: request.permission.resource.instanceId,
-    linuxUserName: request.cliLocalData.linuxUserName,
+    subscriptionId: request.permission.resource.subscriptionId,
+    instanceResourceGroup: request.permission.resource.resourceGroupId,
+    bastionId: request.permission.bastionHostId,
   }),
 
   // TODO: Implement
   unprovisionedAccessPatterns: [],
 
-  // TODO: Placeholder
-  toCliRequest: async (request, options) => ({
-    ...request,
-    cliLocalData: {
-      linuxUserName: await importSshKey(request.permission.publicKey, options),
-    },
-  }),
+  toCliRequest: async (request) => {
+    return {
+      ...request,
+      cliLocalData: {
+        linuxUserName: request.principal,
+      },
+    };
+  },
 };