Merge pull request #152 from AikidoSec/pypi-feature-flag

Add feature flag in setup for python support.

packages/safe-chain/bin/safe-chain.js

@@ -6,6 +6,7 @@ import { ui } from "../src/environment/userInteraction.js";
 import { setup } from "../src/shell-integration/setup.js";
 import { teardown } from "../src/shell-integration/teardown.js";
 import { setupCi } from "../src/shell-integration/setup-ci.js";
+import { initializeCliArguments } from "../src/config/cliArguments.js";
 
 if (process.argv.length < 3) {
   ui.writeError("No command provided. Please provide a command to execute.");
@@ -14,6 +15,8 @@ if (process.argv.length < 3) {
   process.exit(1);
 }
 
+initializeCliArguments(process.argv);
+
 const command = process.argv[2];
 
 if (command === "help" || command === "--help" || command === "-h") {
@@ -56,6 +59,11 @@ function writeHelp() {
       "safe-chain setup"
     )}: This will setup your shell to wrap safe-chain around npm, npx, yarn, pnpm, pnpx, bun, bunx, pip and pip3.`
   );
+  ui.writeInformation(
+    `    ${chalk.yellow(
+      "--include-python"
+    )}: Experimental: include Python package managers (pip, pip3) in the setup.`
+  );
   ui.writeInformation(
     `- ${chalk.cyan(
       "safe-chain teardown"
@@ -67,9 +75,14 @@ function writeHelp() {
     )}: This will setup safe-chain for CI environments by creating shims and modifying the PATH.`
   );
   ui.writeInformation(
-    `- ${chalk.cyan(
-      "safe-chain --version"
-    )} (or ${chalk.cyan("-v")}): Display the current version of safe-chain.`
+    `    ${chalk.yellow(
+      "--include-python"
+    )}: Experimental: include Python package managers (pip, pip3) in the setup.`
+  );
+  ui.writeInformation(
+    `- ${chalk.cyan("safe-chain --version")} (or ${chalk.cyan(
+      "-v"
+    )}): Display the current version of safe-chain.`
   );
   ui.emptyLine();
 }

packages/safe-chain/src/config/cliArguments.js

@@ -1,8 +1,9 @@
 /**
- * @type {{loggingLevel: string | undefined}}
+ * @type {{loggingLevel: string | undefined, includePython: boolean}}
  */
 const state = {
   loggingLevel: undefined,
+  includePython: false,
 };
 
 const SAFE_CHAIN_ARG_PREFIX = "--safe-chain-";
@@ -27,6 +28,7 @@ export function initializeCliArguments(args) {
   }
 
   setLoggingLevel(safeChainArgs);
+  setIncludePython(args);
 
   return remainingArgs;
 }
@@ -64,3 +66,31 @@ function setLoggingLevel(args) {
 export function getLoggingLevel() {
   return state.loggingLevel;
 }
+
+/**
+ * @param {string[]} args
+ */
+function setIncludePython(args) {
+  // This flag doesn't have the --safe-chain- prefix because
+  // it is only used for the safe-chain command itself and
+  // not when wrapped around package manager commands.
+  state.includePython = hasFlagArg(args, "--include-python");
+}
+
+export function includePython() {
+  return state.includePython;
+}
+
+/**
+ * @param {string[]} args
+ * @param {string} flagName
+ * @returns {boolean}
+ */
+function hasFlagArg(args, flagName) {
+  for (const arg of args) {
+    if (arg.toLowerCase() === flagName.toLowerCase()) {
+      return true;
+    }
+  }
+  return false;
+}

packages/safe-chain/src/shell-integration/helpers.js

@@ -2,28 +2,30 @@ import { spawnSync } from "child_process";
 import * as os from "os";
 import fs from "fs";
 import path from "path";
+import { ECOSYSTEM_JS, ECOSYSTEM_PY } from "../config/settings.js";
 
 /**
  * @typedef {Object} AikidoTool
  * @property {string} tool
  * @property {string} aikidoCommand
+ * @property {string} ecoSystem
  */
 
 /**
  * @type {AikidoTool[]}
  */
 export const knownAikidoTools = [
-  { tool: "npm", aikidoCommand: "aikido-npm" },
-  { tool: "npx", aikidoCommand: "aikido-npx" },
-  { tool: "yarn", aikidoCommand: "aikido-yarn" },
-  { tool: "pnpm", aikidoCommand: "aikido-pnpm" },
-  { tool: "pnpx", aikidoCommand: "aikido-pnpx" },
-  { tool: "bun", aikidoCommand: "aikido-bun" },
-  { tool: "bunx", aikidoCommand: "aikido-bunx" },
-  { tool: "pip", aikidoCommand: "aikido-pip" },
-  { tool: "pip3", aikidoCommand: "aikido-pip3" },
-  { tool: "python", aikidoCommand: "aikido-python" },
-  { tool: "python3", aikidoCommand: "aikido-python3" },
+  { tool: "npm", aikidoCommand: "aikido-npm", ecoSystem: ECOSYSTEM_JS },
+  { tool: "npx", aikidoCommand: "aikido-npx", ecoSystem: ECOSYSTEM_JS },
+  { tool: "yarn", aikidoCommand: "aikido-yarn", ecoSystem: ECOSYSTEM_JS },
+  { tool: "pnpm", aikidoCommand: "aikido-pnpm", ecoSystem: ECOSYSTEM_JS },
+  { tool: "pnpx", aikidoCommand: "aikido-pnpx", ecoSystem: ECOSYSTEM_JS },
+  { tool: "bun", aikidoCommand: "aikido-bun", ecoSystem: ECOSYSTEM_JS },
+  { tool: "bunx", aikidoCommand: "aikido-bunx", ecoSystem: ECOSYSTEM_JS },
+  { tool: "pip", aikidoCommand: "aikido-pip", ecoSystem: ECOSYSTEM_PY },
+  { tool: "pip3", aikidoCommand: "aikido-pip3", ecoSystem: ECOSYSTEM_PY },
+  { tool: "python", aikidoCommand: "aikido-python", ecoSystem: ECOSYSTEM_PY },
+  { tool: "python3", aikidoCommand: "aikido-python3", ecoSystem: ECOSYSTEM_PY },
   // When adding a new tool here, also update the documentation for the new tool in the README.md
 ];
 

packages/safe-chain/src/shell-integration/setup-ci.js

@@ -1,10 +1,12 @@
 import chalk from "chalk";
 import { ui } from "../environment/userInteraction.js";
-import { knownAikidoTools, getPackageManagerList } from "./helpers.js";
+import { getPackageManagerList, knownAikidoTools } from "./helpers.js";
 import fs from "fs";
 import os from "os";
 import path from "path";
 import { fileURLToPath } from "url";
+import { includePython } from "../config/cliArguments.js";
+import { ECOSYSTEM_PY } from "../config/settings.js";
 
 /**
  * Loops over the detected shells and calls the setup function for each.
@@ -53,7 +55,7 @@ function createUnixShims(shimsDir) {
 
   // Create a shim for each tool
   let created = 0;
-  for (const toolInfo of knownAikidoTools) {
+  for (const toolInfo of getToolsToSetup()) {
     const shimContent = template
       .replaceAll("{{PACKAGE_MANAGER}}", toolInfo.tool)
       .replaceAll("{{AIKIDO_COMMAND}}", toolInfo.aikidoCommand);
@@ -66,9 +68,7 @@ function createUnixShims(shimsDir) {
     created++;
   }
 
-  ui.writeInformation(
-    `Created ${created} Unix shim(s) in ${shimsDir}`
-  );
+  ui.writeInformation(`Created ${created} Unix shim(s) in ${shimsDir}`);
 }
 
 /**
@@ -96,19 +96,17 @@ function createWindowsShims(shimsDir) {
 
   // Create a shim for each tool
   let created = 0;
-  for (const toolInfo of knownAikidoTools) {
+  for (const toolInfo of getToolsToSetup()) {
     const shimContent = template
       .replaceAll("{{PACKAGE_MANAGER}}", toolInfo.tool)
       .replaceAll("{{AIKIDO_COMMAND}}", toolInfo.aikidoCommand);
 
-  const shimPath = `${shimsDir}/${toolInfo.tool}.cmd`;
+    const shimPath = `${shimsDir}/${toolInfo.tool}.cmd`;
     fs.writeFileSync(shimPath, shimContent, "utf-8");
     created++;
   }
 
-  ui.writeInformation(
-    `Created ${created} Windows shim(s) in ${shimsDir}`
-  );
+  ui.writeInformation(`Created ${created} Windows shim(s) in ${shimsDir}`);
 }
 
 /**
@@ -145,3 +143,11 @@ function modifyPathForCi(shimsDir) {
     ui.writeInformation("##vso[task.prependpath]" + shimsDir);
   }
 }
+
+function getToolsToSetup() {
+  if (includePython()) {
+    return knownAikidoTools;
+  } else {
+    return knownAikidoTools.filter((tool) => tool.ecoSystem !== ECOSYSTEM_PY);
+  }
+}

packages/safe-chain/src/shell-integration/setup.js

@@ -6,6 +6,7 @@ import fs from "fs";
 import os from "os";
 import path from "path";
 import { fileURLToPath } from "url";
+import { includePython } from "../config/cliArguments.js";
 
 /**
  * Loops over the detected shells and calls the setup function for each.
@@ -104,7 +105,11 @@ function copyStartupFiles() {
     // Use absolute path for source
     const __filename = fileURLToPath(import.meta.url);
     const __dirname = path.dirname(__filename);
-    const sourcePath = path.resolve(__dirname, "startup-scripts", file);
+    const sourcePath = path.resolve(
+      __dirname,
+      includePython() ? "startup-scripts/include-python" : "startup-scripts",
+      file
+    );
     fs.copyFileSync(sourcePath, targetPath);
   }
 }

packages/safe-chain/src/shell-integration/startup-scripts/include-python/init-fish.fish

@@ -0,0 +1,88 @@
+function printSafeChainWarning
+    set original_cmd $argv[1]
+    
+    # Fish equivalent of ANSI color codes: yellow background, black text for "Warning:"
+    set_color -b yellow black
+    printf "Warning:"
+    set_color normal
+    printf " safe-chain is not available to protect you from installing malware. %s will run without it.\n" $original_cmd
+    
+    # Cyan text for the install command
+    printf "Install safe-chain by using "
+    set_color cyan
+    printf "npm install -g @aikidosec/safe-chain"
+    set_color normal
+    printf ".\n"
+end
+
+function wrapSafeChainCommand
+    set original_cmd $argv[1]
+    set aikido_cmd $argv[2]
+    set cmd_args $argv[3..-1]
+    
+    if type -q $aikido_cmd
+        # If the aikido command is available, just run it with the provided arguments
+        $aikido_cmd $cmd_args
+    else
+        # If the aikido command is not available, print a warning and run the original command
+        printSafeChainWarning $original_cmd
+        command $original_cmd $cmd_args
+    end
+end
+
+function npx
+    wrapSafeChainCommand "npx" "aikido-npx" $argv
+end
+
+function yarn
+    wrapSafeChainCommand "yarn" "aikido-yarn" $argv
+end
+
+function pnpm
+    wrapSafeChainCommand "pnpm" "aikido-pnpm" $argv
+end
+
+function pnpx
+    wrapSafeChainCommand "pnpx" "aikido-pnpx" $argv
+end
+
+function bun
+    wrapSafeChainCommand "bun" "aikido-bun" $argv
+end
+
+function bunx
+    wrapSafeChainCommand "bunx" "aikido-bunx" $argv
+end
+
+function npm
+    # If args is just -v or --version and nothing else, just run the `npm -v` command
+    # This is because nvm uses this to check the version of npm
+    set argc (count $argv)
+    if test $argc -eq 1
+        switch $argv[1]
+            case "-v" "--version"
+                command npm $argv
+                return
+        end
+    end
+
+    wrapSafeChainCommand "npm" "aikido-npm" $argv
+end
+
+function pip
+    wrapSafeChainCommand "pip" "aikido-pip" $argv
+end
+
+function pip3
+    wrapSafeChainCommand "pip3" "aikido-pip3" $argv
+end
+
+# `python -m pip`, `python -m pip3`.
+function python
+    wrapSafeChainCommand "python" "aikido-python" $argv
+end
+
+# `python3 -m pip`, `python3 -m pip3'.
+function python3
+    wrapSafeChainCommand "python3" "aikido-python3" $argv
+end

packages/safe-chain/src/shell-integration/startup-scripts/include-python/init-posix.sh

@@ -0,0 +1,80 @@
+
+function printSafeChainWarning() {
+  # \033[43;30m is used to set the background color to yellow and text color to black
+  # \033[0m is used to reset the text formatting
+  printf "\033[43;30mWarning:\033[0m safe-chain is not available to protect you from installing malware. %s will run without it.\n" "$1"
+  # \033[36m is used to set the text color to cyan
+  printf "Install safe-chain by using \033[36mnpm install -g @aikidosec/safe-chain\033[0m.\n"
+}
+
+function wrapSafeChainCommand() {
+  local original_cmd="$1"
+  local aikido_cmd="$2"
+
+  # Remove the first 2 arguments (original_cmd and aikido_cmd) from $@
+  # so that "$@" now contains only the arguments passed to the original command
+  shift 2
+
+  if command -v "$aikido_cmd" > /dev/null 2>&1; then
+    # If the aikido command is available, just run it with the provided arguments
+    "$aikido_cmd" "$@"
+  else
+    # If the aikido command is not available, print a warning and run the original command
+    printSafeChainWarning "$original_cmd"
+
+    command "$original_cmd" "$@"
+  fi
+}
+
+function npx() {
+  wrapSafeChainCommand "npx" "aikido-npx" "$@"
+}
+
+function yarn() {
+  wrapSafeChainCommand "yarn" "aikido-yarn" "$@"
+}
+
+function pnpm() {
+  wrapSafeChainCommand "pnpm" "aikido-pnpm" "$@"
+}
+
+function pnpx() {
+  wrapSafeChainCommand "pnpx" "aikido-pnpx" "$@"
+}
+
+function bun() {
+  wrapSafeChainCommand "bun" "aikido-bun" "$@"
+}
+
+function bunx() {
+  wrapSafeChainCommand "bunx" "aikido-bunx" "$@"
+}
+
+function npm() {
+  if [[ "$1" == "-v" || "$1" == "--version" ]] && [[ $# -eq 1 ]]; then
+    # If args is just -v or --version and nothing else, just run the npm version command
+    # This is because nvm uses this to check the version of npm
+    command npm "$@"
+    return
+  fi
+
+  wrapSafeChainCommand "npm" "aikido-npm" "$@"
+}
+
+function pip() {
+  wrapSafeChainCommand "pip" "aikido-pip" "$@"
+}
+
+function pip3() {
+  wrapSafeChainCommand "pip3" "aikido-pip3" "$@"
+}
+
+# `python -m pip`, `python -m pip3`.
+function python() {
+  wrapSafeChainCommand "python" "aikido-python" "$@"
+}
+
+# `python3 -m pip`, `python3 -m pip3'.
+function python3() {
+  wrapSafeChainCommand "python3" "aikido-python3" "$@"
+}

packages/safe-chain/src/shell-integration/startup-scripts/include-python/init-pwsh.ps1

@@ -0,0 +1,107 @@
+function Write-SafeChainWarning {
+    param([string]$Command)
+    
+    # PowerShell equivalent of ANSI color codes: yellow background, black text for "Warning:"
+    Write-Host "Warning:" -BackgroundColor Yellow -ForegroundColor Black -NoNewline
+    Write-Host " safe-chain is not available to protect you from installing malware. $Command will run without it."
+    
+    # Cyan text for the install command
+    Write-Host "Install safe-chain by using " -NoNewline
+    Write-Host "npm install -g @aikidosec/safe-chain" -ForegroundColor Cyan -NoNewline
+    Write-Host "."
+}
+
+function Test-CommandAvailable {
+    param([string]$Command)
+    
+    try {
+        Get-Command $Command -ErrorAction Stop | Out-Null
+        return $true
+    }
+    catch {
+        return $false
+    }
+}
+
+function Invoke-RealCommand {
+    param(
+        [string]$Command,
+        [string[]]$Arguments
+    )
+    
+    # Find the real executable to avoid calling our wrapped functions
+    $realCommand = Get-Command -Name $Command -CommandType Application | Select-Object -First 1
+    if ($realCommand) {
+        & $realCommand.Source @Arguments
+    }
+}
+
+function Invoke-WrappedCommand {
+    param(
+        [string]$OriginalCmd,
+        [string]$AikidoCmd,
+        [string[]]$Arguments
+    )
+
+    if (Test-CommandAvailable $AikidoCmd) {
+        & $AikidoCmd @Arguments
+    }
+    else {
+        Write-SafeChainWarning $OriginalCmd
+        Invoke-RealCommand $OriginalCmd $Arguments
+    }
+}
+
+function npx {
+    Invoke-WrappedCommand "npx" "aikido-npx" $args
+}
+
+function yarn {
+    Invoke-WrappedCommand "yarn" "aikido-yarn" $args
+}
+
+function pnpm {
+    Invoke-WrappedCommand "pnpm" "aikido-pnpm" $args
+}
+
+function pnpx {
+    Invoke-WrappedCommand "pnpx" "aikido-pnpx" $args
+}
+
+function bun {
+    Invoke-WrappedCommand "bun" "aikido-bun" $args
+}
+
+function bunx {
+    Invoke-WrappedCommand "bunx" "aikido-bunx" $args
+}
+
+function npm {
+    # If args is just -v or --version and nothing else, just run the npm version command
+    # This is because nvm uses this to check the version of npm
+    if (($args.Length -eq 1) -and (($args[0] -eq "-v") -or ($args[0] -eq "--version"))) {
+        Invoke-RealCommand "npm" $args
+        return
+    }
+    
+    Invoke-WrappedCommand "npm" "aikido-npm" $args
+}
+
+function pip {
+    Invoke-WrappedCommand "pip" "aikido-pip" $args
+}
+
+function pip3 {
+    Invoke-WrappedCommand "pip3" "aikido-pip3" $args
+}
+
+# `python -m pip`, `python -m pip3`.
+function python {
+    Invoke-WrappedCommand 'python' 'aikido-python' $args
+}
+
+# `python3 -m pip`, `python3 -m pip3'.
+function python3 {
+    Invoke-WrappedCommand 'python3' 'aikido-python3' $args
+}
+

packages/safe-chain/src/shell-integration/startup-scripts/init-fish.fish

@@ -68,21 +68,3 @@ function npm
 
     wrapSafeChainCommand "npm" "aikido-npm" $argv
 end
-
-function pip
-    wrapSafeChainCommand "pip" "aikido-pip" $argv
-end
-
-function pip3
-    wrapSafeChainCommand "pip3" "aikido-pip3" $argv
-end
-
-# `python -m pip`, `python -m pip3`.
-function python
-    wrapSafeChainCommand "python" "aikido-python" $argv
-end
-
-# `python3 -m pip`, `python3 -m pip3'.
-function python3
-    wrapSafeChainCommand "python3" "aikido-python3" $argv
-end

packages/safe-chain/src/shell-integration/startup-scripts/init-posix.sh

@@ -60,21 +60,3 @@ function npm() {
 
   wrapSafeChainCommand "npm" "aikido-npm" "$@"
 }
-
-function pip() {
-  wrapSafeChainCommand "pip" "aikido-pip" "$@"
-}
-
-function pip3() {
-  wrapSafeChainCommand "pip3" "aikido-pip3" "$@"
-}
-
-# `python -m pip`, `python -m pip3`.
-function python() {
-  wrapSafeChainCommand "python" "aikido-python" "$@"
-}
-
-# `python3 -m pip`, `python3 -m pip3'.
-function python3() {
-  wrapSafeChainCommand "python3" "aikido-python3" "$@"
-}

packages/safe-chain/src/shell-integration/startup-scripts/init-pwsh.ps1

@@ -86,22 +86,3 @@ function npm {
     
     Invoke-WrappedCommand "npm" "aikido-npm" $args
 }
-
-function pip {
-    Invoke-WrappedCommand "pip" "aikido-pip" $args
-}
-
-function pip3 {
-    Invoke-WrappedCommand "pip3" "aikido-pip3" $args
-}
-
-# `python -m pip`, `python -m pip3`.
-function python {
-    Invoke-WrappedCommand 'python' 'aikido-python' $args
-}
-
-# `python3 -m pip`, `python3 -m pip3'.
-function python3 {
-    Invoke-WrappedCommand 'python3' 'aikido-python3' $args
-}
-