Implement reproCommands for Azure SSH (#144)

This PR improves the Azure SSH debugging experience by adding
reproduction commands for Azure SSH, which are displayed when invoking
`p0 ssh` with the `--debug` option. It outputs all the `az` commands, as
used by the P0 CLI, that are needed to manually reproduce the SSH
connection process, as well as the final `ssh`/`scp` command too. I also
included some minor refactoring which both made the implementation
easier, as well as removed some redundancy from the `ssh` options we end
up passing for Azure SSH (e.g., previously, we technically specified two
separate private keys to `ssh`, the usual P0 CLI-managed key as well as
the ephemeral SSH certificate private key generated by `az ssh cert`).

src/commands/shared/ssh.ts

@@ -58,6 +58,9 @@ export type SshAdditionalSetup = {
   /** A list of SSH configuration options, as would be used after '-o' in an SSH command */
   sshOptions: string[];
 
+  /** The path to the private key file to use for the SSH connection, instead of the default P0 CLI managed key */
+  identityFile: string;
+
   /** The port to connect to, overriding the default */
   port: string;
 

src/plugins/azure/auth.ts

@@ -10,9 +10,21 @@ You should have received a copy of the GNU General Public License along with @p0
 **/
 import { exec } from "../../util";
 
+export const azLoginCommand = () => ({
+  command: "az",
+  args: ["login"],
+});
+
+export const azAccountSetCommand = (subscriptionId: string) => ({
+  command: "az",
+  args: ["account", "set", "--subscription", subscriptionId],
+});
+
 export const azLogin = async (subscriptionId: string) => {
-  await exec("az", ["login"], { check: true });
-  await exec("az", ["account", "set", "--subscription", subscriptionId], {
-    check: true,
-  });
+  const { command: azLoginExe, args: azLoginArgs } = azLoginCommand();
+  await exec(azLoginExe, azLoginArgs, { check: true });
+
+  const { command: azAccountSetExe, args: azAccountSetArgs } =
+    azAccountSetCommand(subscriptionId);
+  await exec(azAccountSetExe, azAccountSetArgs, { check: true });
 };

src/plugins/azure/keygen.ts

@@ -19,6 +19,11 @@ 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 azSshCertCommand = (keyPath: string) => ({
+  command: "az",
+  args: ["ssh", "cert", "--file", path.join(keyPath, AD_CERT_FILENAME)],
+});
+
 export const createTempDirectoryForKeys = async (): Promise<{
   path: string;
   cleanup: () => Promise<void>;
@@ -36,11 +41,8 @@ export const createTempDirectoryForKeys = async (): Promise<{
 
 export const generateSshKeyAndAzureAdCert = async (keyPath: string) => {
   try {
-    await exec(
-      "az",
-      ["ssh", "cert", "--file", path.join(keyPath, AD_CERT_FILENAME)],
-      { check: true }
-    );
+    const { command, args } = azSshCertCommand(keyPath);
+    await exec(command, args, { check: true });
   } catch (error: any) {
     print2(error.stdout);
     print2(error.stderr);

src/plugins/azure/ssh.ts

@@ -9,15 +9,16 @@ 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 { azLogin } from "./auth";
+import { azAccountSetCommand, azLogin, azLoginCommand } from "./auth";
 import { ensureAzInstall } from "./install";
 import {
   AD_CERT_FILENAME,
   AD_SSH_KEY_PRIVATE,
+  azSshCertCommand,
   createTempDirectoryForKeys,
   generateSshKeyAndAzureAdCert,
 } from "./keygen";
-import { trySpawnBastionTunnel } from "./tunnel";
+import { azBastionTunnelCommand, trySpawnBastionTunnel } from "./tunnel";
 import {
   AzureLocalData,
   AzureSshPermissionSpec,
@@ -60,11 +61,49 @@ export const azureSshProvider: SshProvider<
   // Azure doesn't support ProxyCommand, as nice as that would be. Yet.
   proxyCommand: () => [],
 
-  // TODO: Determine if necessary
-  reproCommands: () => undefined,
+  reproCommands: (request, additionalData) => {
+    const { command: azLoginExe, args: azLoginArgs } = azLoginCommand();
+    const { command: azAccountSetExe, args: azAccountSetArgs } =
+      azAccountSetCommand(request.subscriptionId);
+
+    const getKeyPath = () => {
+      // Use the same key path as the one generated in setup() so it matches the ssh command that is generated
+      // elsewhere. It'll be an annoying long temporary directory name, but it strictly will work for reproduction. If
+      // additionalData isn't present (which it always should be for the azureSshProvider), we'll use the user's home
+      // directory.
+      if (additionalData) {
+        return path.dirname(additionalData.identityFile);
+      } else {
+        const basePath = process.env.HOME || process.env.USERPROFILE || "";
+        return path.join(basePath, "p0cli-azure-ssh-keys");
+      }
+    };
+
+    const keyPath = getKeyPath();
+
+    const { command: azCertGenExe, args: azCertGenArgs } =
+      azSshCertCommand(keyPath);
+
+    // If additionalData is undefined (which, again, should be never), use the default port for Azure Network Bastion
+    // tunnels instead of generating a random one
+    const { command: azTunnelExe, args: azTunnelArgs } = azBastionTunnelCommand(
+      request,
+      additionalData?.port ?? "50022"
+    );
+
+    return [
+      `${azLoginExe} ${azLoginArgs.join(" ")}`,
+      `${azAccountSetExe} ${azAccountSetArgs.join(" ")}`,
+      `mkdir ${keyPath}`,
+      `${azCertGenExe} ${azCertGenArgs.join(" ")}`,
+      `${azTunnelExe} ${azTunnelArgs.join(" ")}`,
+    ];
+  },
 
   setup: async (request) => {
-    // TODO(ENG-3129): Does this specifically need to be the subscription ID for the Bastion?
+    // The subscription ID here is used to ensure that the user is logged in to the correct tenant/directory.
+    // As long as a subscription ID in the correct tenant is provided, this will work; it need not be the same
+    // subscription as which contains the Bastion host or the target VM.
     await azLogin(request.subscriptionId); // Always re-login to Azure CLI
 
     const { path: keyPath, cleanup: sshKeyPathCleanup } =
@@ -92,19 +131,18 @@ export const azureSshProvider: SshProvider<
 
     return {
       sshOptions: [
-        `IdentityFile ${sshPrivateKeyPath}`,
-        `CertificateFile ${sshCertificateKeyPath}`,
-        "IdentitiesOnly yes",
+        `CertificateFile=${sshCertificateKeyPath}`,
 
         // 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",
+        "StrictHostKeyChecking=no",
+        "UserKnownHostsFile=/dev/null",
       ],
+      identityFile: sshPrivateKeyPath,
       port: tunnelLocalPort,
       teardown,
     };

src/plugins/azure/tunnel.ts

@@ -22,6 +22,26 @@ export type BastionTunnelMeta = {
   tunnelLocalPort: string;
 };
 
+export const azBastionTunnelCommand = (
+  request: AzureSshRequest,
+  port: string
+) => ({
+  command: "az",
+  args: [
+    "network",
+    "bastion",
+    "tunnel",
+    "--ids",
+    request.bastionId,
+    "--target-resource-id",
+    request.instanceId,
+    "--resource-port",
+    "22",
+    "--port",
+    port,
+  ],
+});
+
 const selectRandomPort = (): string => {
   // The IANA ephemeral port range is 49152 to 65535, inclusive. Pick a random value in that range.
   // If the port is in use (unlikely but possible), we can just generate a new value and try again.
@@ -40,25 +60,11 @@ const spawnBastionTunnelInBackground = (
     let stdout = "";
     let stderr = "";
 
-    const child = spawn(
-      "az",
-      [
-        "network",
-        "bastion",
-        "tunnel",
-        "--ids",
-        request.bastionId,
-        "--target-resource-id",
-        request.instanceId,
-        "--resource-port",
-        "22",
-        "--port",
-        port,
-      ],
-      // Spawn the process in detached mode so that it is in its own process group; this lets us kill it and all
-      // descendent processes together.
-      { detached: true }
-    );
+    const { command, args } = azBastionTunnelCommand(request, port);
+
+    // Spawn the process in detached mode so that it is in its own process group; this lets us kill it and all
+    // descendent processes together.
+    const child = spawn(command, args, { detached: true });
 
     child.on("exit", (code) => {
       processExited = true;

src/plugins/ssh/index.ts

@@ -214,12 +214,12 @@ const createCommand = (
   setupData: SshAdditionalSetup | undefined,
   proxyCommand: string[]
 ) => {
-  addCommonArgs(args, proxyCommand);
+  addCommonArgs(args, proxyCommand, setupData);
 
-  const sshOptions = setupData?.sshOptions ?? [];
+  const sshOptionsOverrides = setupData?.sshOptions ?? [];
   const port = setupData?.port;
 
-  const argsOverride = sshOptions.flatMap((opt) => ["-o", opt]);
+  const argsOverride = sshOptionsOverrides.flatMap((opt) => ["-o", opt]);
 
   if ("source" in args) {
     addScpArgs(args);
@@ -260,7 +260,8 @@ const createCommand = (
  */
 const addCommonArgs = (
   args: CommandArgs,
-  sshProviderProxyCommand: string[]
+  sshProviderProxyCommand: string[],
+  setupData: SshAdditionalSetup | undefined
 ) => {
   const sshOptions = args.sshOptions ? args.sshOptions : [];
 
@@ -278,7 +279,8 @@ const addCommonArgs = (
   // Explicitly specify which private key to use to avoid "Too many authentication failures"
   // error caused by SSH trying every available key
   if (!identityFileOptionExists) {
-    sshOptions.push("-i", PRIVATE_KEY_PATH);
+    sshOptions.push("-i", setupData?.identityFile ?? PRIVATE_KEY_PATH);
+
     // Only use the authentication identity specified by -i above
     if (!identitiesOnlyOptionExists) {
       sshOptions.push("-o", "IdentitiesOnly=yes");
@@ -410,7 +412,7 @@ export const sshOrScp = async (args: {
   );
 
   if (cmdArgs.debug) {
-    const reproCommands = sshProvider.reproCommands(request);
+    const reproCommands = sshProvider.reproCommands(request, setupData);
     if (reproCommands) {
       const repro = [
         ...reproCommands,

src/types/ssh.ts

@@ -91,7 +91,10 @@ export type SshProvider<
 
   /** Each element in the returned array is a command that can be run to reproduce the
    * steps of logging in the user to the ssh session. */
-  reproCommands: (request: SR) => string[] | undefined;
+  reproCommands: (
+    request: SR,
+    additionalData?: SshAdditionalSetup
+  ) => string[] | undefined;
 
   /** Unwraps this provider's types */
   requestToSsh: (request: CliPermissionSpec<PR, O>) => SR;