wip, adding ollama as an alternative to detect-secrets

1 files changed, 94 insertions, 0 deletions

diff --git a/ollama.lua b/ollama.lua
new file mode 100644
index 0000000..5a27160
--- /dev/null
+++ b/ollama.lua
@@ -0,0 +1,94 @@
+local http_request = require("http.request")
+local libgen = require("posix.libgen")
+
+local base_path = libgen.dirname(arg[0])
+package.path = package.path .. ";" .. base_path .. "/?.lua"
+local json = require("json")
+
+local ollama = {}
+
+function ollama.ollama_req(clipboard_content)
+    local url = "http://172.17.0.1:11434/api/chat"
+    local req = http_request.new_from_uri(url)
+
+    local body = {
+        model = "llama3.1",
+        stream = false,
+        format = "json",
+        messages = {
+            {content = clipboard_content, role = "user"}, {
+                content = [[
+            a public key is not a secret.
+            a base64 encoded string is a not secret.
+            a private key is a seceret.
+            an api key is a secret.
+            a password is a secret.
+            a token is a secret.
+            a long string of random characters is a secret.
+            ]],
+                role = "assistant"
+            }, {
+                content = [[
+                Only answer in json. 
+                The answer must a field named 'isSecret'.
+                The answer must have a field named 'reasoning'.
+                The value of 'isSecret' must be a boolean.
+                The value of reasoning must be a string.
+                The answer must be valid json.
+                ]],
+                role = "assistant"
+            }, {
+                content = [[
+                Now I will give you your task.
+                Look at the user-provided string content. 
+                Is it a secret? answer in json.
+                ]],
+                role = "assistant"
+            }
+        },
+        options = {
+            temperature = 0.5,
+            max_tokens = 10000,
+            top_p = 1.0,
+            frequency_penalty = 0.0,
+            presence_penalty = 0.0
+        }
+    }
+
+    local body_json = json.encode(body)
+    req:set_body(body_json)
+
+    req.headers:upsert(":method", "POST")
+
+    local headers, stream = req:go(10000)
+
+    if headers:get(":status") ~= "200" then return nil end
+
+    local result_body = stream:get_body_as_string()
+
+    return result_body
+end
+
+function ollama.ask_ollama(clipboard_content, count)
+    local true_count = 0
+    local false_count = 0
+
+    for _ = 1, count do
+        local result = ollama.ollama_req(clipboard_content)
+        local result_decoded = json.decode(result)
+        local final_result = json.decode(result_decoded["message"]["content"])
+        if final_result == true then
+            true_count = true_count + 1
+        else
+            false_count = false_count + 1
+        end
+    end
+
+    if true_count > false_count then
+        return true
+    else
+        return false
+    end
+end
+
+return ollama