Heuristics Based AV Prototype

local detectThreshold = 50
local heuristicsTriggers = { -- this is a list of function calls that are classified as "dangerous".
    "fs.delete",
    "fs.move",
    "fs.copy",
    "fs.open",
    "os.reboot",
    "os.shutdown",
}

local totalLines = 0
local totalScore = 0
local totalPercent = 0
local scores = {}
local percents = {}
local args = {...}

if args[1] == nil then
    print("USAGE: "..fs.getName(shell.getRunningProgram()).." [file]")
    return
end

for _,v in ipairs(heuristicsTriggers) do
    scores[v] = 0
    percents[v] = 0
end

local handle = io.open(args[1], "r")

for line in handle:lines() do
    if #line > 0 then
        totalLines = totalLines+1
        for _,trigger in ipairs(heuristicsTriggers) do
            for match in string.gmatch(line, "("..trigger..")") do
                scores[trigger] = scores[trigger]+1
                totalScore = totalScore+1
            end
        end
    end
end

totalPercent = math.ceil((totalScore/totalLines) * 100)

term.clear()
term.setCursorPos(1,1)
print("Results of heuristics scan for file: "..args[1])
print("Total (useful) lines: "..totalLines)
print("Total trigger phrase count: "..totalScore)
print("Total trigger phrase percentage count: "..totalPercent.."%")
for i,v in pairs(scores) do
    percents[i] = math.ceil((v/totalLines) * 100)
    print("Percentage found: "..i..": "..percents[i].."% (found "..v.." times)")
end

for i,v in pairs(percents) do
    if v >= detectThreshold then
        print("Caution: This file may be a virus, because it has a lot of "..i.." calls.")
        print("We may be wrong, but it may be a good idea to check this file manually.")
        print("Move to quarantine? (Y/N)")
        write(">")
        local quarantine = false
        if string.upper(string.sub(read(),1,1)) == "Y" then
            quarantine = true
        else
            quarantine = false
        end
        if quarantine then
            if not fs.exists("quarantine") then
                fs.makeDir("quarantine")
            end
            local oldDataHandle = fs.open(args[1], "r")
            local oldData = "--[["..oldDataHandle.readAll().."]]"
            oldDataHandle.close()
            local quarantineHandle = fs.open("quarantine/"..args[1], "w")
            quarantineHandle.write(oldData)
            quarantineHandle.close()
        end
    end
end

if totalPercent >= detectThreshold then
    print("Caution: This file may be a virus, because it contains many \"dangerous\" calls.")
    print("We may be wrong, but it may be a good idea to check this file manually.")
    print("Move to quarantine? (Y/N)")
    write(">")
    local quarantine = false
    if string.upper(string.sub(read(),1,1)) == "Y" then
        quarantine = true
    else
        quarantine = false
    end
    if quarantine then
        if not fs.exists("quarantine") then
            fs.makeDir("quarantine")
        end
        local oldDataHandle = fs.open(args[1], "r")
        local oldData = "--[["..oldDataHandle.readAll().."]]"
        oldDataHandle.close()
        local quarantineHandle = fs.open("quarantine/"..args[1], "w")
        quarantineHandle.write(oldData)
        quarantineHandle.close()
    end
end

print("Large percentages may indicate viruses and/or malicious software.")
print("Note: You may need to check this file manually.")