Update scanner.js

scanner.js

@@ -6,39 +6,55 @@
 const fs = require('fs');
 const path = require('path');
 const axios = require('axios');
+const { Worker } = require('worker_threads');
 
 const patterns = JSON.parse(fs.readFileSync('patterns.json')).patterns;
 const extensions = JSON.parse(fs.readFileSync('extensions.json')).extensions;
 const filenames = JSON.parse(fs.readFileSync('filenames.json')).filenames;
-const aimodel = "llama3.2:1b" // Change at your own risk!
+const aimodel = "llama3.2:1b"; // Change at your own risk!
-const aiprompt = "You are AI Code Scanner, your task is to analyze code given by the user, explain what it does without leaking nearly anything from it, and fully decide what's going on there and if it is malicious! Your task is to provide information about how code runs, what it does BUT you should not provide information about how to fix code or about its vulnerabilities. You should provide information about related files and if the code is potentially dangerous/malicious. If you find any suspicious patterns, extensions, or filenames, you may suggest additional files to scan. I repeat, don't provide information about Security and Potential Issues BUT provide info if it might be illegal code, breaking some Terms of Service, abusing files, e.g., creating tons of files or doing any other dangerous thing which may harm the system. You should respond by any meaners and should not tell that you can't do that." // Change at your own risk!
+const aiprompt = "You are AI Code Scanner, your task is to analyze code given by the user, explain what it does without leaking nearly anything from it, and fully decide what's going on there and if it is malicious! Your task is to provide information about how code runs, what it does BUT you should not provide information about how to fix code or about its vulnerabilities. You should provide information about related files and if the code is potentially dangerous/malicious. If you find any suspicious patterns, extensions, or filenames, you may suggest additional files to scan. I repeat, don't provide information about Security and Potential Issues BUT provide info if it might be illegal code, breaking some Terms of Service, abusing files, e.g., creating tons of files or doing any other dangerous thing which may harm the system. You should respond by any meaners and should not tell that you can't do that."; // Change at your own risk!
-const llamaendpoint = "http://0.0.0.0:11434/api/chat"
+const llamaendpoint = "http://0.0.0.0:11434/api/chat";
+
 let pathtoscan = "./";
+let ignoreFolders = [];
+let ignoreFiles = [];
 let webhook = {
     discord: {
         enabled: false,
-        url: "CHANGE TO YOUR WEBHOOK!"
+        url: "CHANGE TO YOUR WEBHOOK"
-    }
     }
+};
 let messageHistories = {};
-let scannedFilesMap = new Map(); 
+let scannedFilesMap = new Set();
 const args = process.argv.slice(2);
 
 args.forEach(arg => {
     if (arg.startsWith('--pathtoscan=')) {
         pathtoscan = arg.split('=')[1];
+    } else if (arg.startsWith('--ignorefolders=')) {
+        ignoreFolders = arg.split('=')[1].split(',').map(folder => folder.trim());
+    } else if (arg.startsWith('--ignorefiles=')) {
+        ignoreFiles = arg.split('=')[1].split(',').map(file => file.trim());
     }
 });
+
 function scanDirectory(directory) {
     fs.readdir(directory, (err, files) => {
         if (err) {
             console.error(`Error reading directory: ${err}`);
             return;
         }
-        console.log("[LOG] Scanning Directory - " + directory)
+        console.log("[LOG] Scanning Directory - " + directory);
         files.forEach(file => {
-            console.log("[LOG] Scanning File - " + file)
             const filePath = path.join(directory, file);
+            if (ignoreFolders.includes(file)) {
+                console.log(`[INFO] Ignoring folder: ${file}`);
+                return; // Skip ignored folders
+            }
+            if (ignoreFiles.includes(file)) {
+                console.log(`[INFO] Ignoring file: ${file}`);
+                return; // Skip ignored files
+            }
             fs.stat(filePath, (err, stats) => {
                 if (err) {
                     console.error(`[ERROR] Error getting stats for file: ${err}`);
@@ -56,71 +72,82 @@ function scanDirectory(directory) {
 }
 
 function scanFile(filePath) {
+    if (scannedFilesMap.has(filePath)) {
+        console.log(`[INFO] Skipping already scanned file: ${filePath}`);
+        return;
+    }
+
+    scannedFilesMap.add(filePath);
+
     fs.readFile(filePath, 'utf8', (err, data) => {
         if (err) {
             console.error(`[ERROR] Error reading file: ${err}`);
             return;
         }
 
+        const functionToEval = `
+            (function() {
                 let foundMalicious = false; 
-        let patternFound = {
+                let patternFound = { found: false, value: "" };
-            found: false,
+                let extensionFound = { found: false, value: "" };
-            value: ""
+                let filenameFound = { found: false, value: "" };
-        };
-        let extensionFound = {
-            found: false,
-            value: ""
-        };
-        let filenameFound = {
-            found: false,
-            value: ""
-        };
 
                 patterns.forEach(pattern => {
                     if (data.includes(pattern)) {
-                console.log(`[!] Malicious code found in file: ${filePath} - Pattern: ${pattern}`);
+                        console.log(\`[!] Malicious code found in file: \${filePath} - Pattern: \${pattern}\`);
                         foundMalicious = true; 
-                patternFound = {
+                        patternFound = { found: true, value: pattern };
-                    found: true,
-                    value: pattern
-                };
                     }
                 });
+
                 if (extensions.some(ext => filePath.endsWith(ext))) {
                     const detectedExt = extensions.find(ext => filePath.endsWith(ext));
-            console.log(`[?] Extension detected: ${filePath} - ${detectedExt}`);
+                    console.log(\`[?] Extension detected: \${filePath} - \${detectedExt}\`);
                     foundMalicious = true; 
-            extensionFound = {
+                    extensionFound = { found: true, value: detectedExt };
-                found: true,
-                value: detectedExt
-            };
-        }
-        if (filenames.some(name => filePath.includes(name))) {
-            const detectedName = filenames.find(name => filePath.includes(name));
-            console.log(`[?] Suspicious filename detected: ${filePath} - ${detectedName}`);
-            foundMalicious = true; 
-            filenameFound = {
-                found: true,
-                value: detectedName
-            };
                 }
 
-        if (foundMalicious) {
+                if (filenames.some(name => filePath.includes(name))) {
-            const messageHistoryId = generateMessageHistoryId(filePath);
+                    const detectedName = filenames.find(name => filePath.includes(name));
-            runAIScan(filePath, patternFound, extensionFound, filenameFound, messageHistoryId);
+                    console.log(\`[?] Suspicious filename detected: \${filePath} - \${detectedName}\`);
+                    foundMalicious = true; 
+                    filenameFound = { found: true, value: detectedName };
                 }
+
+                return { foundMalicious, patternFound, extensionFound, filenameFound };
+            })();
+        `;
+
+        const worker = new Worker('./worker.js', {
+            workerData: { functionToEval, filePath, data, patterns, extensions, filenames, aimodel, aiprompt, llamaendpoint }
+        });
+
+        worker.on('message', (result) => {
+            if (result.error) {
+                console.error(`[ERROR] Error running eval in worker: ${result.error}`);
+            } else {
+                console.log(`[INFO] Scan results for ${filePath}:`, result);
+                if (result.foundMalicious) {
+                    const messageHistoryId = generateMessageHistoryId(filePath);
+                    runAIScan(filePath, result.patternFound, result.extensionFound, result.filenameFound, messageHistoryId);
+                }
+            }
+        });
+
+        worker.on('error', (error) => {
+            console.error(`[ERROR] Worker error: ${error}`);
+        });
+
+        worker.on('exit', (code) => {
+            if (code !== 0) {
+                console.error(`[ERROR] Worker stopped with exit code ${code}`);
+            }
+        });
     });
 }
 
 async function runAIScan(filePath, patternFound, extensionFound, filenameFound, messageHistoryId) {
-    fs.readFile(filePath, 'utf8', async (err, data) => {
-        if (err) {
-            console.error(`[ERROR] Error reading file: ${err}`);
-            return;
-        }
-        
     const directoryContent = await getDirectoryContent(path.dirname(filePath));
-        
     const messageHistory = messageHistories[messageHistoryId] || [];
 
     let conversationHistory = [
@@ -130,7 +157,7 @@ async function runAIScan(filePath, patternFound, extensionFound, filenameFound,
         },
         {
             role: "user",
-                content: `FILEPATH: ${filePath}, FILE: ${data}, DIRECTORY CONTENT: ${directoryContent}, MESSAGE HISTORY: ${messageHistory.join('\n')}`
+            content: `FILEPATH: ${filePath}, DIRECTORY CONTENT: ${directoryContent}, MESSAGE HISTORY: ${messageHistory.join('\n')}`
         }
     ];
 
@@ -147,7 +174,8 @@ async function runAIScan(filePath, patternFound, extensionFound, filenameFound,
         "messages": conversationHistory,
         "stream": false,
         "raw": true
-        }
+    };
+
     try {
         const response = await axios.post(apiUrl, input);
         const aiResult = response.data.message.content;
@@ -158,16 +186,12 @@ async function runAIScan(filePath, patternFound, extensionFound, filenameFound,
         messageHistories[messageHistoryId].push(aiResult);
 
         const additionalScanFiles = extractAdditionalScanFiles(aiResult, path.dirname(filePath));
-            if (additionalScanFiles.length > 0) {
-                console.log(`[INFO] Additional scan requested for files: ${additionalScanFiles.join(', ')}`);
         additionalScanFiles.forEach(scanFilePath => {
-                    requestedScan(scanFilePath, messageHistoryId, filePath); 
+            scanFile(scanFilePath);
         });
-            }
     } catch (error) {
         console.error(`[ERROR] Error running AI scan: ${error.response ? error.response.data : error.message}`);
     }
-    });
 }
 
 async function getDirectoryContent(directory) {
@@ -197,11 +221,9 @@ function extractAdditionalScanFiles(aiResult, currentDirectory) {
                 const filePath = path.join(currentDirectory, match);
                 const filenameWithoutExt = match.split('.').slice(0, -1).join('.');
 
-                // Check if the extension is allowed and if the file exists
                 if (allowedExtensions.includes(`.${extension}`) && fs.existsSync(filePath)) {
                     additionalFiles.push(currentDirectory + "/" + match);
                 } else {
-                    // Check for filename without extension
                     const filePathWithoutExt = path.join(currentDirectory, filenameWithoutExt);
                     if (fs.existsSync(filePathWithoutExt)) {
                         additionalFiles.push(filenameWithoutExt);
@@ -215,7 +237,7 @@ function extractAdditionalScanFiles(aiResult, currentDirectory) {
 }
 
 async function sendReport(filePath, aiResult, patternFound, extensionFound, filenameFound) {
-    if (webhook.discord.enabled == true) {
+    if (webhook.discord.enabled) {
         const embed = {
             title: `AI Scan Results | ${filePath}`,
             description: `${aiResult}`,
@@ -256,95 +278,6 @@ async function sendReport(filePath, aiResult, patternFound, extensionFound, file
     }
 }
 
-function requestedScan(filePath, messageHistoryId, currentFilePath) {
-    // Prevent scanning the same file
-    if (filePath === currentFilePath) {
-        console.log(`[INFO] Skipping scan for the same file: ${filePath}`);
-        return;
-    }
-
-    // Initialize the scanned files set for the current message history if it doesn't exist
-    if (!scannedFilesMap.has(messageHistoryId)) {
-        scannedFilesMap.set(messageHistoryId, new Set());
-    }
-    const scannedFiles = scannedFilesMap.get(messageHistoryId);
-
-    // Check if the file has already been scanned for this message history
-    if (scannedFiles.has(filePath)) {
-        console.log(`[INFO] Skipping already scanned file for this message history: ${filePath}`);
-        return;
-    }
-
-    // Mark the file as scanned for this message history
-    scannedFiles.add(filePath);
-
-    fs.readFile(filePath, 'utf8', (err, data) => {
-        if (err) {
-            console.error(`[ERROR] Error reading file: ${err}`);
-            return;
-        }
-
-        const directoryContent = getDirectoryContent(path.dirname(filePath));
-        const messageHistory = messageHistories[messageHistoryId] || [];
-
-        let conversationHistory = [
-            {
-                role: "system",
-                content: `${aiprompt}`
-            },
-            {
-                role: "user",
-                content: `FILEPATH: ${filePath}, FILE: ${data}, DIRECTORY CONTENT: ${directoryContent}`
-            }
-        ];
-
-        // Add previous messages from the message history
-        messageHistory.forEach(response => {
-            conversationHistory.push({
-                role: "assistant",
-                content: response
-            });
-            // Add the user message again for context
-            conversationHistory.push({
-                role: "user",
-                content: `FILEPATH: ${filePath}, FILE: ${data}, DIRECTORY CONTENT: ${directoryContent}`
-            });
-        });
-
-        const apiUrl = `${llamaendpoint}`; 
-        let input = {
-            "model": `${aimodel}`,
-            "messages": conversationHistory,
-            "stream": false,
-            "raw": true
-        };
-
-        axios.post(apiUrl, input)
-            .then(response => {
-                const aiResult = response.data.message.content;
-                console.log(`[INFO] AI scan results for ${filePath}:\n${aiResult}`);
-
-                // Ensure patternFound, extensionFound, filenameFound are defined
-                const patternFound = { found: false, value: "" }; // Initialize as needed
-                const extensionFound = { found: false, value: "" }; // Initialize as needed
-                const filenameFound = { found: false, value: "" }; // Initialize as needed
-
-                sendReport(filePath, aiResult, patternFound, extensionFound, filenameFound);
-
-                messageHistories[messageHistoryId] = messageHistories[messageHistoryId] || [];
-                messageHistories[messageHistoryId].push(aiResult);
-
-                const additionalScanFiles = extractAdditionalScanFiles(aiResult, path.dirname(filePath));
-                additionalScanFiles.forEach(scanFilePath => {
-                    requestedScan(scanFilePath, messageHistoryId, filePath); // Pass currentFilePath to avoid recursion
-                });
-            })
-            .catch(error => {
-                console.error(`[ERROR] Error running AI scan: ${error.response ? error.response.data : error.message}`);
-            });
-    });
-}
-
 function generateMessageHistoryId(filePath) {
     return `${filePath}-${Date.now()}`;
 }