// AI CODE SCANNER
// Lisa Honkay
// lisahonkay@gmail.com
// Copyright Lisa Honkay CC-BY-SA-4.0

const fs = require('fs');
const path = require('path');
const axios = require('axios');
const yara = require('@automattic/yara');
const { Worker } = require('worker_threads');

//const aimodel = "huihui_ai/llama3.2-abliterate: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 may show code if need for this reason. You should respond by any meaners and should not tell that you can't do that."; // Change at your own risk!
const llamaendpoint = "http://192.168.2.109:11434/api/chat";
const extensions = [
    '.exe', // Executable files
    '.dll', // Dynamic Link Library files
    '.scr', // Screensaver files
    '.vbs', // Visual Basic Script files
    '.js',  // JavaScript files
    '.bat', // Batch files
    '.cmd', // Command files
    '.ps1', // PowerShell Script files
    '.ps2', // Second PowerShell Script files
    '.psm1', // PowerShell Module files
    '.py',  // Python files
    '.pyx', // Cython files
    '.jsx', // JavaScript XML files
    '.ts',  // TypeScript files
    '.tsx'  // TypeScript XML files
];


let pathtoscan = "./";
let ignoreFolders = [];
let ignoreFiles = [];
let webhook = {
    discord: {
        enabled: true,
        url: "https://discord.com/api/webhooks/1349816121972359240/4KhIbjzDERDnk0_OAJDA0Ha2YzzntvVTM96hhDNzmWYx65NfRad6LWFIx5aM8NZScxO-"
    }
};
let messageHistories = {};
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 fileExtensionMatch(filePath) {
    return extensions.some(extension => filePath.endsWith(extension));
}

async function scanFileWithYara(filePath) {
    return new Promise((resolve, reject) => {
        try {
            yara.initialize((error) => {
                if (error) {
                    console.error(`Error initializing Yara: ${error}`);
                    resolve({ found: false, rule: null });
                } else {
                    // Load all rule files from the signatures folder
                    const signaturesPath = path.join(__dirname, 'signatures');
                    const rules = fs.readdirSync(signaturesPath)
                        .filter((file) => file.endsWith('.yar') || file.endsWith('.yara'))
                        .map((file) => ({ filename: path.join(signaturesPath, file) }));

                    const scanner = yara.createScanner();

                    scanner.configure({ rules: rules }, (error, warnings) => {
                        if (error) {
                            if (error instanceof yara.CompileRulesError) {
                                console.error(`Error compiling rules: ${error}`);
                                resolve({ found: false, rule: null });
                            } else {
                                console.error(`Error configuring scanner: ${error}`);
                                resolve({ found: false, rule: null });
                            }
                        } else {
                            if (warnings.length) {
                                console.warn(`Compile warnings: ${JSON.stringify(warnings)}`);
                            }

                            const req = { buffer: fs.readFileSync(filePath) };
                            scanner.scan(req, (error, result) => {
                                if (error) {
                                    console.error(`Error scanning file: ${error}`);
                                    resolve({ found: false, rule: null });
                                } else if (result && result.matches && result.matches.length > 0) {
                                    // Handle valid matches
                                    const matchedRule = result.matches[0].rule;
                                    resolve({ found: true, rule: matchedRule });
                                } else {
                                    // Handle no matches
                                    if (fileExtensionMatch(filePath) == true) {
                                        resolve({ found: true, rule: "Suspicious File Extension found in " + filePath });
                                    } else {
                                        resolve({ found: false, rule: null });
                                    }
                                }
                            });
                        }
                    });
                }
            });
        } catch (error) {
            console.error(`Unexpected error: ${error}`);
            resolve({ found: false, rule: null });
        }
    });
}

function scanDirectory(directory) {
    fs.readdir(directory, (err, files) => {
        if (err) {
            console.error(`Error reading directory: ${err}`);
            return;
        }
        console.log("[LOG] Scanning Directory - " + directory);
        files.forEach(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}`);
                    return;
                }

                if (stats.isDirectory()) {
                    scanDirectory(filePath);
                } else {
                    scanFile(filePath);
                }
            });
        });
    });
}

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;
        }

        scanFileWithYara(filePath).then((yaraMatch) => {
            const foundMalicious = (yaraMatch.found);

            if (foundMalicious) {
                console.log(`[!] Malicious code detected in file: ${filePath}`);
                console.log(`[!] Yara rule matched: ${yaraMatch.rule}`);

                // Start Worker Thread for additional processing
                const worker = new Worker('./worker.js', {
                    workerData: { filePath, data: '', aimodel, aiprompt, llamaendpoint }
                });

                worker.on('message', (result) => {
                    if (result.error) {
                        console.error(`[ERROR] Error processing in worker: ${result.error}`);
                    } else {
                        //console.log(`[INFO] Scan results for ${filePath}:`, result);
                        //if (result.foundMalicious) {
                            const messageHistoryId = generateMessageHistoryId(filePath);
                            runAIScan(filePath, yaraMatch, messageHistoryId);
                        //} else {
                        //    console.log("[WARN] Empty message from worker recieved!")
                        //}
                    }
                });

                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}`);
                    }
                });
            } else {
                console.log(`[INFO] No malicious activity detected in file: ${filePath}`);
            }
        }).catch((scanError) => {
            console.error(`[ERROR] Error scanning file with Yara: ${scanError}`);
        });
    });
}

async function runAIScan(filePath, yaraMatch, messageHistoryId) {
    console.log("[INFO] AI Scan requested for file: " + filePath)
    const directoryContent = await getDirectoryContent(path.dirname(filePath));
    const messageHistory = messageHistories[messageHistoryId] || [];

    let conversationHistory = [
        {
            role: "system",
            content: `${aiprompt}`
        },
        {
            role: "user",
            content: `FILEPATH: ${filePath}, DIRECTORY CONTENT: ${directoryContent}, FILECONTENT: ${fs.readFile(filePath, 'utf8', (err, data) => {
                if (err) {
                    console.err("[ERROR] Could not read file (" + filePath + "), error: " + err)
                    return `[ERROR] Could not read file: ${err}`
                } else {
                    return data
                }
            })}, MESSAGE HISTORY: ${messageHistory.join('\n')}`
        }
    ];

    messageHistory.forEach(response => {
        conversationHistory.push({
            role: "assistant",
            content: response
        });
    });

    const apiUrl = `${llamaendpoint}`;
    let input = {
        "model": `${aimodel}`,
        "messages": conversationHistory,
        "stream": false,
        "raw": true
    };

    try {
        const response = await axios.post(apiUrl, input);
        const aiResult = response.data.message.content;
        console.log(`[INFO] AI scan results for ${filePath}:\n${aiResult}`);
        await sendReport(filePath, aiResult, yaraMatch);

        messageHistories[messageHistoryId] = messageHistories[messageHistoryId] || [];
        messageHistories[messageHistoryId].push(aiResult);

        const additionalScanFiles = extractAdditionalScanFiles(aiResult, path.dirname(filePath));
        additionalScanFiles.forEach(scanFilePath => {
            scanFile(scanFilePath);
        });
    } catch (error) {
        console.error(`[ERROR] Error running AI scan: ${error.response ? error.response.data : error.message}`);
    }
}

async function getDirectoryContent(directory) {
    return new Promise((resolve, reject) => {
        fs.readdir(directory, (err, files) => {
            if (err) {
                reject(`[ERROR] Error reading directory: ${err}`);
            } else {
                resolve(files.join(', '));
            }
        });
    });
}

function extractAdditionalScanFiles(aiResult, currentDirectory) {
    const sentences = aiResult.split('\n');
    const additionalFiles = [];
    const allowedExtensions = JSON.parse(fs.readFileSync('extensions.json')).extensions;

    console.log(`[INFO] Scanning for additional files in directory: ${currentDirectory}`);

    sentences.forEach(sentence => {
        const matches = sentence.match(/([a-zA-Z0-9-_]+(\.[a-zA-Z0-9]+)+)/g);
        if (matches) {
            matches.forEach(match => {
                const extension = match.split('.').pop();
                const filePath = path.join(currentDirectory, match);
                const filenameWithoutExt = match.split('.').slice(0, -1).join('.');

                if (allowedExtensions.includes(`.${extension}`) && fs.existsSync(filePath)) {
                    additionalFiles.push(currentDirectory + "/" + match);
                } else {
                    const filePathWithoutExt = path.join(currentDirectory, filenameWithoutExt);
                    if (fs.existsSync(filePathWithoutExt)) {
                        additionalFiles.push(filenameWithoutExt);
                    }
                }
            });
        }
    });

    return additionalFiles;
}

async function sendReport(filePath, aiResult, yaraMatch) {
    if (webhook.discord.enabled) {
        const embed = {
            title: `AI Scan Results | ${filePath}`,
            description: `${aiResult}`,
            fields: [
                {
                    name: `${yaraMatch.found ? "✅ Yara Rule matched" : "❌ Yara Rule not matched"}`,
                    value: yaraMatch.rule ? yaraMatch.rule : "N/A",
                    inline: true
                }
            ],
            color: 0x00FF00,
            timestamp: new Date(),
            footer: {
                text: "AI Code Scanner - Made by Lisa_NS"
            }
        };

        const payload = {
            embeds: [embed]
        };

        try {
            await axios.post(webhook.discord.url, payload);
            console.log(`[INFO] Results sent to Discord webhook for ${filePath}`);
        } catch (error) {
            console.error(`[ERROR] Error sending to Discord webhook: ${error.response ? error.response.data : error.message}`);
        }
    }
}

function generateMessageHistoryId(filePath) {
    return `${filePath}-${Date.now()}`;
}

scanDirectory(pathtoscan);