automated-clean.js

import fs, { existsSync, mkdirSync, rmSync } from "fs";
import { loadDir, buildFullAst, findRootExports, taintPath } from "./index.js";
import { buildDepGraph } from "./depGraph.js";
import { getTaintGraph } from "./taintGraph.js";
import { join } from "path";
import * as child_process from "child_process";
import * as readline from "readline";
import { allChildNodes } from "./utils.js";
import { generateReport } from "./generateReport.js";

const TIMEOUT = 30000;

const reposDirName = "./repos"; // create this dir before executing
const astDirName = "./ast"; // create this dir before executing

const lastTaintGraph = {
    name: "",
    graph: null,
};

const startTime = Date.now();
const lastIndex = 0;  // put index of obj from which you want to start the analysis (in case the execution was interrupted mid-file)

async function main(){
    const text = "[" + fs.readFileSync(process.argv[2]).toString().split("\n").filter(Boolean).join(",") + "]"; // input results from semgrep (list of objs {function, filelink, matches})
    const db = JSON.parse(text).slice(lastIndex);

    await forEachAsync(db, processFunc, 1);
}

main();

function forEachAsync(array, func, concurrency=0) {
    let nextUp = concurrency > 0
        ? Math.min(concurrency, array.length)
        : array.length;

    return Promise.all(array.slice(0, nextUp)
        .map((item, i) => run(i))
    );

    async function run(index) {
        while(index < array.length) {
            try {
                await func(array[index], index, array)
            } catch (error) {
                console.error(error)
            } finally {
                index = nextUp;
                nextUp += 1;
            }
        }
    }
}


async function processFunc(funcObj, i) {
    try {
        log("Starting...", i);
        const {local, github} = downloadProjectDir(funcObj.fileLink, i);   // specific for projects with github file link
        const astName = funcObj.fileLink.split("/").slice(3,6).join("/");  // adapt to contain unique name of the project (if not npm)

        if(lastTaintGraph.name !== astName){
            log("Generating asts...", i);
            // console.time(`created ast and graphs`)
            const asts = await loadDir(local, astName);
            const nodes = allChildNodes(asts).length;

            if (nodes > 1000000) {
                throw new Error(`Maximum project AST exceeded: ${nodes} nodes total`);
            }
            log("Asts done", i);

            log("Adding import/export info...", i);
            const depGraph = buildDepGraph(asts); //=> builds obj Graph with import-export info
            const fullAst = buildFullAst(asts, depGraph); // => combines DepGraph (FDG) with file ASTs and produces 1 extended AST
            const rootExports = findRootExports(fullAst); // specific to npm packages => root exports = sources. For other projects identify own "source" nodes
            log(`Found ${rootExports.length} root exports`, i) // adapt for "source nodes"

            const taints = rootExports.map(node => ({
                node,
                by: [],
                kind: "args",
                argIndex: -1,
                rootTaint: true       // mark nodes selected as "sources" as rootTaint
            }));

            lastTaintGraph.name = astName;
            
            log(`Import/export info complete.`, i);

            log(`Generating taint graph...`, i);
            lastTaintGraph.graph = getTaintGraph(fullAst, taints);
            // console.timeEnd(`created ast and graphs`);
        }
        const taintGraph = lastTaintGraph.graph;
        log(`Taint graph generated.`, i);
        log(`${taintGraph.size} nodes tainted`, i);
   
        log(`Getting sink local coordinates...`, i);
        // console.time("get coordinates from semgrep");
        const fullFuncObj = getUniqueCoords(funcObj); //reduces semgrep matches to unique ones, extends obj with reachablePaths: [] 
        // console.timeEnd("get coordinates from semgrep");
        const filePath = join(   
            local,
            fullFuncObj.fileLink.split("/").slice(6).join("/")
        );
        log(`Sink local coordinates retrieved.`, i);

        log(`Getting function coordinates...`, i);
        const funcPos = await getCoordsInFile(filePath, funcObj.function);
        log(`Function coordinates retrieved.`, i);
        
        
        //get sink
        log(`Generating sinks...`, i);
        console.time("get sinks");
        const sinks = [];
        fullFuncObj.matches.forEach(match => {
            const matchpos = {
                start: {
                    line: match.loc.start.line + funcPos.line - 1,
                    column: match.loc.start.line === 1
                        ? match.loc.start.column + funcPos.column
                        : match.loc.start.column
                },
                end: {
                    line: match.loc.end.line + funcPos.line - 1,
                    column: match.loc.end.line === 1
                        ? match.loc.end.column + funcPos.column
                        : match.loc.end.column
                },
            };
            sinks.push(getSink(
                taintGraph,
                filePath,
                matchpos.start.line,
                matchpos.start.column - 1,
                matchpos.end.line,
                matchpos.end.column - 1
            ));
        });
        log(`Sinks generated.`, i);

        sinks.forEach(vulnSink => {
            if (vulnSink) {
                log("!!!!!!!!!!!!!Sink found!!!!!!!!!!!!!!!", i);
                // const temp = taintPath(taintGraph.get(vulnSink)[0], github);
                fullFuncObj.reachableSinkPaths.push(taintPath(taintGraph.get(vulnSink)[0], github)); //github - link to a code line with tainted node
            } else {
                // log("Path to the specified sink not found");
            }
        });
        console.timeEnd("get sinks");

        // log(fullFuncObj.reachableSinkPaths.map(Boolean), i);
        log(`Total elapsed: ${formatMs(Date.now() - startTime)}`, i);
        if(fullFuncObj.reachableSinkPaths.length) {
            fs.appendFileSync("./taintResults.json", JSON.stringify(fullFuncObj))  // resulting object with taint path
            fs.appendFileSync("./taintResultsReports.txt", generateReport(fullFuncObj) + "\n\n----------------------------------\n\n"); // readable report with information on exploitability
        }
    } catch (error) {
        console.error(error);
        
        fs.appendFileSync("./errors.json", JSON.stringify({i: i + lastIndex, message: error.message, stack: error.stack}) + "\n");
    } finally {
    }
}



function downloadProjectDir(githubLink, i) {  
    if (!githubLink.startsWith("https://raw.githubusercontent.com/")) {
        throw new Error(`Only githubusercontent urls are supported! Got ${githubLink}`)
    }
    
    const [user,repo,commit] = githubLink.split("/").slice(3,6);
    log(`Downloading ${user}/${repo}`, i);
    if (!existsSync(join(reposDirName, user, repo))) {
        console.time("folder downloaded");
        try {
            rmSync(reposDirName, {recursive: true, force: true});
        } catch (error) {
            console.error(error);
        }
        try {
            rmSync(astDirName, {recursive: true, force: true});
        } catch (error) {
            console.error(error);
        }
        try {
            mkdirSync(join(reposDirName, user), {recursive: true});
        } catch (error) {
            console.error(error);
        }
        child_process.execSync(`git clone https://github.com/${user}/${repo} --depth 1`, {
            cwd: join(reposDirName, user),
            stdio: 'ignore',
            timeout: TIMEOUT,
            killSignal: "SIGKILL"
        });
        console.timeEnd("folder downloaded");
    } else {
        log("Folder already exisits.", i);
    }
    log(`Checking out ${commit}...`, i);
    child_process.execSync(`git fetch --depth 1 origin ${commit} && git checkout FETCH_HEAD`, {
        cwd: join(reposDirName, user, repo),
        stdio: 'ignore',
        timeout: TIMEOUT,
        killSignal: "SIGKILL"
    });
    return {
        local: join(reposDirName, user, repo),
        github: join("https://github.com/", user, repo, "/blob/", commit),
    };
}


function getUniqueCoords(funcObj) {
    const usefulSemgrep = funcObj.matches.reduce((acc, item) => {  //reduce duplicate matches (1 place by multiple rules) 
        const key = JSON.stringify(item.loc);
        acc[key] = acc[key] || {
            ...item,
            ruleID: undefined,
            ruleIds: [],
        };
        acc[key].ruleIds.push(item.ruleID);
        return acc;
    }, {});
    
    return {...funcObj, matches: Object.values(usefulSemgrep), reachableSinkPaths: []}
}

function getSink(taintGraph, 
    targetFile,
    targetStartLine,
    targetStartCol,
    targetEndLine,
    targetEndCol
    ){
    const vulnSink = Array.from(taintGraph.keys()).find(node => 
        node.file === targetFile &&
        (
            (node.loc.start?.line === +targetStartLine && node.loc.start?.column === +targetStartCol) ||
            (node.loc.end?.line === +targetEndLine && node.loc.end?.column === +targetEndCol)
        )
    );
    return vulnSink;
}

function formatMs(ms) {
    var date = new Date(ms);
    var timeString = date.toISOString().slice(11, 19);
    return timeString;
}


function log(str, i) {
    console.log(`${new Date().toLocaleTimeString()} [${i + lastIndex}]: ${str}`);
}

async function getCoordsInFile(file, str) {
    const linesToSearch = str.split(/\r?\n/);
    let currentCompareLine = 0;
    let currentLine = 1;
    let line = 0;
    let column = 0;

    const fileStream = fs.createReadStream(file);
    const rl = readline.createInterface({
      input: fileStream,
      crlfDelay: Infinity
    });
    for await (const l of rl) {
        if (linesToSearch.length === 1 && l.includes(linesToSearch[currentCompareLine])) {
            return {
                line: currentLine,
                column: l.indexOf(linesToSearch[currentCompareLine]) + 1
            };
        } else if (!currentCompareLine && l.endsWith(linesToSearch[currentCompareLine])) {
            line = currentLine;
            column = l.indexOf(linesToSearch[currentCompareLine]) + 1;
            currentCompareLine++;
        } else if (currentCompareLine < linesToSearch.length - 1 && l === linesToSearch[currentCompareLine]) {
            currentCompareLine++;
        } else if (currentCompareLine === linesToSearch.length - 1 && l.startsWith(linesToSearch[currentCompareLine])) {
            return {
                line,
                column
            };
        } else {
            line = 0;
            column = 0;
            currentCompareLine = 0;
        }
        currentLine++;
    }
    return null;
}



// const vuln =     [{
//     "function": "function parseQuery(qs) {\n  return qs\n    .replace('?', '')\n    .split('&')\n    .reduce(function(obj, pair) {\n      var i = pair.indexOf('=');\n      var key = pair.slice(0, i);\n      var val = pair.slice(++i);\n\n      // Due to how the URLSearchParams API treats spaces\n      obj[key] = decodeURIComponent(val.replace(/\\+/g, '%20'));\n\n      return obj;\n    }, {});\n}",
//     "fileLink": "https://raw.githubusercontent.com/mochajs/mocha/238268dfede04fc68ae1902fac5fdda8f098432f/lib/browser/parse-query.js",
//     "protoPollution": true
// },
// {
//     "function": "function arrayLikeKeys(value, inherited) {\n  const isArr = Array.isArray(value)\n  const isArg = !isArr && isArguments(value)\n  const isBuff = !isArr && !isArg && isBuffer(value)\n  const isType = !isArr && !isArg && !isBuff && isTypedArray(value)\n  const skipIndexes = isArr || isArg || isBuff || isType\n  const length = value.length\n  const result = new Array(skipIndexes ? length : 0)\n  let index = skipIndexes ? -1 : length\n  while (++index < length) {\n    result[index] = `${index}`\n  }\n  for (const key in value) {\n    if ((inherited || hasOwnProperty.call(value, key)) &&\n        !(skipIndexes && (\n        // Safari 9 has enumerable `arguments.length` in strict mode.\n          (key === 'length' ||\n           // Skip index properties.\n           isIndex(key, length))\n        ))) {\n      result.push(key)\n    }\n  }\n  return result\n}",
//     "fileLink": "https://raw.githubusercontent.com/lodash/lodash/86a852fe763935bb64c12589df5391fd7d3bb14d/.internal/arrayLikeKeys.js",
//     "protoPollution": true,
//     "info": "fixed in new versions, still present in master"
// },
// {
//     "function": "function invert(object) {\n  const result = {}\n  Object.keys(object).forEach((key) => {\n    let value = object[key]\n    if (value != null && typeof value.toString !== 'function') {\n      value = toString.call(value)\n    }\n    result[value] = key\n  })\n  return result\n}",
//     "fileLink": "https://raw.githubusercontent.com/lodash/lodash/86a852fe763935bb64c12589df5391fd7d3bb14d/invert.js",
//     "protoPollution": true,
//     "info": "fixed in new versions, still present in master"
// },
// ]