From b07fdbb5e277620c4c6a605b4d782c9ff44b04e5 Mon Sep 17 00:00:00 2001
From: vindard <17693119+vindard@users.noreply.github.com>
Date: Thu, 16 Nov 2023 14:39:00 -0400
Subject: [PATCH] build(api): add 'audit' to pnpm toolchain

---
 toolchains/workspace-pnpm/BUCK          |  5 ++
 toolchains/workspace-pnpm/macros.bzl    | 66 +++++++++++++++++++++++++
 toolchains/workspace-pnpm/run_audit.py  | 55 +++++++++++++++++++++
 toolchains/workspace-pnpm/toolchain.bzl |  5 ++
 4 files changed, 131 insertions(+)
 create mode 100644 toolchains/workspace-pnpm/run_audit.py

diff --git a/toolchains/workspace-pnpm/BUCK b/toolchains/workspace-pnpm/BUCK
index 830135d6f4..ed62a43d88 100644
--- a/toolchains/workspace-pnpm/BUCK
+++ b/toolchains/workspace-pnpm/BUCK
@@ -42,3 +42,8 @@ export_file(
   name = "run_in_dir.py",
   visibility = ["PUBLIC"],
 )
+
+export_file(
+  name = "run_audit.py",
+  visibility = ["PUBLIC"],
+)
diff --git a/toolchains/workspace-pnpm/macros.bzl b/toolchains/workspace-pnpm/macros.bzl
index 4efee40b59..bf40792efd 100644
--- a/toolchains/workspace-pnpm/macros.bzl
+++ b/toolchains/workspace-pnpm/macros.bzl
@@ -608,6 +608,72 @@ def _npm_test_impl(
         DefaultInfo(default_output = args_file),
     ]
 
+def _audit_impl(ctx: AnalysisContext) -> list[[
+    DefaultInfo,
+    RunInfo,
+    ExternalRunnerTestInfo,
+]]:
+    pnpm_toolchain = ctx.attrs._workspace_pnpm_toolchain[WorkspacePnpmToolchainInfo]
+
+    audit_args = cmd_args()
+    audit_args.add("--ignore-registry-errors")
+
+    run_cmd_args = cmd_args([
+        ctx.attrs._python_toolchain[PythonToolchainInfo].interpreter,
+        pnpm_toolchain.run_audit[DefaultInfo].default_outputs,
+        "--audit-level",
+        ctx.attrs.level,
+        "--",
+        audit_args,
+    ])
+
+    args_file = ctx.actions.write("args.txt", run_cmd_args)
+
+    return inject_test_run_info(
+        ctx,
+        ExternalRunnerTestInfo(
+            type = "audit",
+            command = [run_cmd_args],
+        ),
+    ) + [
+        DefaultInfo(default_output = args_file),
+    ]
+
+_audit = rule(
+    impl = _audit_impl,
+    attrs = {
+        "level": attrs.enum(
+            ["low", "moderate", "high", "critical"],
+            default = "critical"
+        ),
+        "node_modules": attrs.source(
+            doc = """Target which builds `node_modules`.""",
+        ),
+        "_inject_test_env": attrs.default_only(
+            attrs.dep(default = "prelude//test/tools:inject_test_env"),
+        ),
+        "_python_toolchain": attrs.toolchain_dep(
+            default = "toolchains//:python",
+            providers = [PythonToolchainInfo],
+        ),
+        "_workspace_pnpm_toolchain": attrs.toolchain_dep(
+            default = "toolchains//:workspace_pnpm",
+            providers = [WorkspacePnpmToolchainInfo],
+        ),
+    },
+)
+
+def audit(
+        node_modules = ":node_modules",
+        visibility = ["PUBLIC"],
+        **kwargs):
+
+    _audit(
+        node_modules = node_modules,
+        visibility = visibility,
+        **kwargs,
+    )
+
 def eslint_impl(ctx: AnalysisContext) -> list[[
     DefaultInfo,
     RunInfo,
diff --git a/toolchains/workspace-pnpm/run_audit.py b/toolchains/workspace-pnpm/run_audit.py
new file mode 100644
index 0000000000..c372035eab
--- /dev/null
+++ b/toolchains/workspace-pnpm/run_audit.py
@@ -0,0 +1,55 @@
+#!/usr/bin/env python3
+"""
+Runs a program in a directory.
+"""
+import argparse
+import json
+import subprocess
+import sys
+
+def sum_severities(severity_dict, start_level):
+    severity_order = [
+        "low",
+        "moderate",
+        "high",
+        "critical"
+    ]
+
+    start_index = severity_order.index(start_level)
+    return sum(
+        severity_dict[level]
+        for level in severity_order[start_index:]
+    )
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--audit-level",
+        help="Audit severity to print advisories against.",
+    )
+    parser.add_argument(
+        "args",
+        help="Audit arguments",
+        nargs=argparse.REMAINDER,
+    )
+
+    args = parser.parse_args()
+    audit_args = args.args[1:] # ignore '--' separator
+
+    pnpm_cmd = ["pnpm", "audit"]
+    audit_cmd = [*pnpm_cmd, *audit_args]
+    audit_cmd_json_out = [*audit_cmd, "--json"]
+
+    result = subprocess.run(audit_cmd_json_out, stdout=subprocess.PIPE)
+    result_dict = json.loads(result.stdout)
+
+    num_vulns = sum_severities(
+        result_dict["metadata"]["vulnerabilities"],
+        args.audit_level
+    )
+    if num_vulns > 0:
+        printable_result = subprocess.run(audit_cmd, stdout=subprocess.PIPE, text=True)
+        print(printable_result.stdout)
+        sys.exit(1)
+
+    sys.exit(0)
diff --git a/toolchains/workspace-pnpm/toolchain.bzl b/toolchains/workspace-pnpm/toolchain.bzl
index edfc909397..141d3b4e65 100644
--- a/toolchains/workspace-pnpm/toolchain.bzl
+++ b/toolchains/workspace-pnpm/toolchain.bzl
@@ -8,6 +8,7 @@ WorkspacePnpmToolchainInfo = provider(fields = [
   "build_next_build",
   "package_next_bin",
   "run_in_dir",
+  "run_audit",
 ])
 
 def workspace_pnpm_toolchain_impl(ctx) -> list[[DefaultInfo, WorkspacePnpmToolchainInfo]]:
@@ -26,6 +27,7 @@ def workspace_pnpm_toolchain_impl(ctx) -> list[[DefaultInfo, WorkspacePnpmToolch
             build_next_build = ctx.attrs._build_next_build,
             package_next_bin = ctx.attrs._package_next_bin,
             run_in_dir = ctx.attrs._run_in_dir,
+            run_audit = ctx.attrs._run_audit,
         )
     ]
 
@@ -59,6 +61,9 @@ workspace_pnpm_toolchain = rule(
         "_run_in_dir": attrs.dep(
             default = "toolchains//workspace-pnpm:run_in_dir.py",
         ),
+        "_run_audit": attrs.dep(
+            default = "toolchains//workspace-pnpm:run_audit.py",
+        ),
     },
     is_toolchain_rule = True,
 )