fix: pii radar (#349)

lunary-ai · Jun 4, 2024 · 91ade93 · 91ade93
1 parent 8640f80
commit 91ade93
Show file tree

Hide file tree

Showing 4 changed files with 91 additions and 54 deletions.
diff --git a/packages/backend/src/api/v1/template-versions.ts b/packages/backend/src/api/v1/template-versions.ts
@@ -24,15 +24,23 @@ versions.get("/latest", async (ctx: Context) => {
   }
 
   const [latestVersion] = await unCameledSql`
-    SELECT t.id::text, t.slug, tv.id::text, tv.content, tv.extra, tv.created_at, tv.version
-    FROM template t
-    INNER JOIN template_version tv ON t.id = tv.template_id
-    WHERE 
+    select 
+      t.id::text, 
+      t.slug, 
+      tv.id::text, 
+      tv.content, 
+      tv.extra, 
+      tv.created_at, 
+      tv.version
+    from 
+      template t
+      inner join template_version tv on t.id = tv.template_id
+    where 
       t.project_id = ${projectId}
-      AND t.slug = ${slug}
-      AND tv.is_draft = false
-    ORDER BY tv.created_at DESC
-    LIMIT 1
+      and t.slug = ${slug}
+      and tv.is_draft = false
+    order by tv.created_at desc
+    limit 1
   `
 
   if (!latestVersion) {

diff --git a/packages/backend/src/checks/index.ts b/packages/backend/src/checks/index.ts
@@ -537,14 +537,51 @@ export const CHECK_RUNNERS: CheckRunner[] = [
     id: "pii",
     async evaluator(run, params) {
       const { field, type, entities } = params
+      const regexResults = {}
 
-      const results = await callML("pii", {
-        texts: getTextsTypes(field, run),
-        entities,
-      })
+      const texts = getTextsTypes(field, run)
+      const text = texts.join(" ")
+      if (!text.length) {
+        return null
+      }
 
-      let passed = false
+      if (entities.includes("email")) {
+        const emailRegex = /[A-Z0-9._%+-]+@[A-Z0-9.-]+\.[A-Z]{2,}/gim
+        regexResults.email = [...new Set(text.match(emailRegex))]
+      }
+
+      if (entities.includes("ip")) {
+        const ipRegex = /(?:\d{1,3}\.){3}\d{1,3}/gim
+        regexResults.ip = [...new Set(text.match(ipRegex))]
+      }
+
+      if (entities.includes("cc")) {
+        const ccRegex = /\d{13,16}/gim
+        regexResults.cc = [...new Set(text.match(ccRegex))]
+      }
 
+      const mlResults = await callML(
+        "pii",
+        {
+          text,
+          entities: entities.filter((entity: string) =>
+            ["person", "location", "org", "misc"].includes(entity),
+          ),
+        },
+        process.env.NEW_ML_URL,
+      )
+
+      const results = { ...regexResults, ...mlResults }
+      console.log(
+        text,
+        entities.filter((entity: string) =>
+          ["person", "location", "org", "misc"].includes(entity),
+        ),
+        mlResults,
+        "\n\n\n\n",
+      )
+
+      let passed = false
       if (type === "contains") {
         passed = Object.keys(results).some(
           (entity: string) => results[entity]?.length > 0,
@@ -555,33 +592,9 @@ export const CHECK_RUNNERS: CheckRunner[] = [
         )
       }
 
-      let labels = {
-        person: "Persons",
-        org: "Organizations",
-        location: "Locations",
-        email: "Emails",
-        phone: "Phone numbers",
-        cc: "Credit card numbers",
-        ssn: "Social security numbers",
-      }
-
-      let reason = `No entities detected among ${entities?.join(", ")}`
-      if (passed) {
-        reason =
-          "Entities detected: \n" +
-          Object.keys(results)
-            .filter((key) => results[key].length > 0)
-            .map(
-              (key: string) =>
-                (labels[key] || key) + ": " + results[key].join(", "),
-            )
-            .join("\n")
-      }
-
       return {
         passed,
-        reason,
-        details: results,
+        details: "",
       }
     },
   },

diff --git a/packages/backend/src/utils/ml.ts b/packages/backend/src/utils/ml.ts
@@ -1,5 +1,9 @@
-export async function callML(method: string, data: any) {
-  const response = await fetch(`${process.env.ML_URL}/${method}`, {
+export async function callML(
+  method: string,
+  data: any,
+  baseUrl: string = process.env.ML_URL,
+) {
+  const response = await fetch(`${baseUrl}/${method}`, {
     method: "POST",
     headers: {
       "Content-Type": "application/json",

diff --git a/packages/shared/checks/index.ts b/packages/shared/checks/index.ts
@@ -613,8 +613,16 @@ export const CHECKS: Check[] = [
         type: "select",
         id: "entities",
         width: 230,
-        // defaultValue: ["person", "location", "email", "cc", "phone", "ssn"],
-        defaultValue: ["person", "location", "org"],
+        defaultValue: [
+          "person",
+          "location",
+          "org",
+          "email",
+          "cc",
+          "phone",
+          "ip",
+          // "ssn",
+        ],
         multiple: true,
         searchable: true,
         options: [
@@ -630,18 +638,22 @@ export const CHECKS: Check[] = [
             label: "Organization",
             value: "org",
           },
-          // {
-          //   label: "Email",
-          //   value: "email",
-          // },
-          // {
-          //   label: "Credit Card",
-          //   value: "cc",
-          // },
-          // {
-          //   label: "Phone",
-          //   value: "phone",
-          // },
+          {
+            label: "Email",
+            value: "email",
+          },
+          {
+            label: "Credit Card",
+            value: "cc",
+          },
+          {
+            label: "Phone",
+            value: "phone",
+          },
+          {
+            label: "IP Address",
+            value: "ip",
+          },
           // {
           //   label: "SSN",
           //   value: "ssn",