CRF v2

-- Common Router Framework
-- By KillaVanilla

local args = {...}


if not _G["CRFLoaded"] then
    _G["CRFLoaded"] = true
    parallel.waitForAny( function() shell.run("shell") end, function() shell.run(shell.getRunningProgram(), unpack(args)) end)
end

for i=1, #rs.getSides() do
    if peripheral.getType(rs.getSides()[i]) == "modem" then
        rednet.open(rs.getSides()[i])
    end
end

-- Arg variables
local parameters = {}

for i=2, #args do
    if string.sub(args[i], 1,2) == "--" then
        parameters[string.sub(args[i],3)] = true
    elseif string.sub(args[i], 1,1) == "-" then
        local combinedParameters = string.sub(args[i],2)
        if #combinedParameters == 1 then
            parameters[combinedParameters] = true
        elseif #combinedParameters > 1 then
            for letter in string.gmatch(combinedParameters, "%a") do
                parameters[letter] = true
            end
        end
    end
end

-- Arg variables
local parameters = {}

for i=2, #args do
    if string.sub(args[i], 1,2) == "--" then
        parameters[string.sub(args[i],3)] = true
    elseif string.sub(args[i], 1,1) == "-" then
        local combinedParameters = string.sub(args[i],2)
        if #combinedParameters == 1 then
            parameters[combinedParameters] = true
        elseif #combinedParameters > 1 then
            for letter in string.gmatch(combinedParameters, "%a") do
                parameters[letter] = true
            end
        end
    end
end

if (parameters["h"] or parameters["help"]) or #args == 0 then
    print("Common Router Framework")
    print("This provides a framework for link-state routing.")
    print("Usage:")
    print(fs.getName(shell.getRunningProgram()).." [path/to/routing/function] [-phtfvb]")
    print("Parameters:")
    print("-b / --background : Disable all output.")
    print("-v / --verbose : Show debug information.")
    print("-s / --hyperspeed : Fast Event Tick Loop")
    print("-p / --show_pings : Show ping received messages.")
    print("-t / --show_timeout_check : Show Timeout Checks")
    print("-l / --logging : Enable Logging")
    print("-h / --help : Show this help screen")
    print("The -v, -p, and -t flags are all independent from one another.")
    print("Passing -v automatically disables -b.")
    print("Routing functions may also take parameters.")
    return
end

-- Set up our tables:
local routerNodeGraph = {} -- Dictionary of who's connected to who. f.e. routerNodeGraph[5] = {1,2,3,4} means that router 5 is connected to routers 1, 2, 3, and 4.
local localRouterNodeGraph = {} -- Table with IDs of who we're connected to. f.e localRouterNodeGraph = {1,2,3,4} means that we're connected to routers 1, 2, 3, and 4 (aka we're router 5).
local nodeStateAdvertList = {} -- Dictionary of the last sequence IDs from all the nodes that we know of.
local routerAdvertTimestamps = {} -- Dictionary containing the timestamps of the every received node advert
local localNodeDistanceGraph = {} -- Dictionary of the distances to all nodes listed in localRouterNodeGraph.
local nodeDistanceGraph = {} -- Dictionary of the distances to the neighbors of all nodes listed in routerNodeGraph.

-- Set up a few constants:
local pingMsg = textutils.serialize({"ROUTING", "LocalPing"})
local pingAckMsg = textutils.serialize({"ROUTING", "LocalPingAck"})
local nodeTimeout = 60 -- [user-specified / changeable] Time To Live for each node entry in routerNodeGraph (we can't just clear the graph every minute or so, like we do with localRouterNodeGraph.)
local sequenceID = 0

if parameters["s"] or parameters["hyperspeed"] then
    assert(nodeTimeout > 5, "Please set nodeTimeout to something above the clock tick rate.")
else
    assert(nodeTimeout > 30, "Please set nodeTimeout to something above the clock tick rate.")
end

local rFnHandle = fs.open(args[1], "r")

-- Our routing function:
local routeFunc = loadstring(rFnHandle.readAll())

rFnHandle.close()

-- Save some functions we're going to overwrite:
local oldRednetSend = rednet.send
local oldRednetBroadcast = rednet.broadcast
local oldRednetReceive = rednet.receive

-- Statistics:
local packets_in = 0
local packets_out = 0
local packets_thru = 0

-- UI Stuff:
local oldPrint = print
local logFile = "CRF.log" -- feel free to change this


local function refreshUI()
    local x, y = term.getCursorPos()
    term.setCursorPos(1,1)
    term.clearLine()
    term.write("In: "..packets_in.." | Thru: "..packets_thru.." | Out: "..packets_out.." | Nodes: "..#routerNodeGraph.." | Local: "..#localRouterNodeGraph)
    term.setCursorPos(x,y)
end


print = function(...)
    local args = {...}
    local opStr = args[1]
    for i=2, #args do
        opStr = opStr..args[i]
    end
    if parameters["l"] or parameters["logging"] then
        local logHandle = fs.open(logFile, "a")
        logHandle.writeLine("[Day"..os.day().."\\"..textutils.formatTime(os.time(), true).."] "..opStr)
        logHandle.close()
    end
    if not (parameters["background"] or parameters["b"]) then
        local x, y = term.getCursorPos()
        term.clearLine()
        oldPrint(opStr)
        refreshUI()
    end
end

if parameters["v"] or parameters["verbose"] then
    parameters["b"] = false
    parameters["background"] = false
end

if not (parameters["b"] or parameters["background"])then
    print("Welcome to the Common Router Framework.")
    print("Parameters selected:")
    if parameters["v"] or parameters["verbose"] then
        print("    Verbose")
    end
    if parameters["f"] or parameters["fake_rednet_events"] then
        print("    Fake Rednet Events")
    end
    if parameters["s"] or parameters["hyperspeed"] then
        print("    Hyperspeed Timer")
    end
    if parameters["p"] or parameters["show_pings"] then
        print("    Show Ping Messages")
    end
    if parameters["t"] or parameters["show_timeout_check"] then
        print("    Show Timeout Checks")
    end
    print("Using routing function: \""..args[1].."\".")
else
    print("Common Router Framework starting...")
end

for i=1, 6 do
    if peripheral.getType(rs.getSides()[i]) == "modem" or peripheral.getType(rs.getSides()[i]) == "cable" then
        rednet.open(rs.getSides()[i])
    end
end

local function pruneDisconnectedNodes() -- run through every node's neighbors and check to see which nodes have zero neighbors (i.e they're unreachable)
    if parameters["v"] or parameters["verbose"] then
        print("Pruning disconnected nodes from routerNodeGraph...")
    end
    local references = {}
    for id, neighbors in pairs(routerNodeGraph) do
        if not references[id] then
            references[id] = 0
        end
        for i, neighbor in ipairs(neighbors) do
            if not references[neighbor] then
                references[neighbor] = 0
            end
            references[neighbor] = references[neighbor]+1
        end
    end
    for id=1, #localRouterNodeGraph do
        if not references[ localRouterNodeGraph[id] ] then
            references[ localRouterNodeGraph[id] ] = 0
        end
        references[ localRouterNodeGraph[id] ] = references[ localRouterNodeGraph[id] ]+1
    end
    if parameters["v"] or parameters["verbose"] then
        print("Node Reference Report:")
    end
    for id=0, table.maxn(routerNodeGraph) do
        if references[id] then
            if parameters["v"] or parameters["verbose"] then
                print("Node "..id..": "..references[id].." reference(s).")
            end
        end
    end
    for id=0, table.maxn(routerNodeGraph) do
        if references[id] then
            if references[id] == 0 then
                if id == os.computerID() then
                    print("WARNING: This node has been cut off from a network!")
                elseif parameters["v"] or parameters["verbose"] then
                    print("Pruning disconnected node: "..id)
                end
                routerNodeGraph[id] = nil
                nodeDistanceGraph[id] = nil
                routerAdvertTimestamps[id] = nil
            end
        end
    end
end

local function sendNodeAdvertisement()
    if parameters["v"] or parameters["verbose"] then
        print("Sending node advertisements...")
    end
    for i=1, #localRouterNodeGraph do -- We'd probably be able to use rednet.broadcast here, but since we don't want to mess up OTHER applications...
        -- if parameters["v"] or parameters["verbose"] then
            -- print("Sending node advertisement to node "..localRouterNodeGraph[i])
        -- end
        oldRednetSend(localRouterNodeGraph[i], textutils.serialize( { "ROUTING", "NodeAdvertisement", os.computerID(), localRouterNodeGraph, localNodeDistanceGraph, {}, sequenceID } )) -- In order: Packet Identifier, Packet Type, Originating ID, Timestamp, Local Router Node Graph, Local Router Distance Graph, recv list, Sequence ID
    end
    sequenceID = sequenceID+1
end

local function handleNodeAdvertisement(packet)
    -- Okay, a few "ground rules" for NodeAdvertisements:
    -- If we've already recieved this (aka packet[6][os.computerID()] == true), then drop it.
    -- If the sequence ID for the packet matches the stored sequence ID for the node, then drop it.
    -- If we don't drop it, then simply set routerNodeGraph[packet[3]] = packet[5], packet[6][os.computerID()] = true, and pass it on.
    if nodeStateAdvertList[packet[3]] then
        if packet[7] == nodeStateAdvertList[packet[3]] then
            return
        end
    end
    if packet[6][os.computerID()] then
        return
    end
    nodeStateAdvertList[packet[3]] = packet[7]
    packet[6][os.computerID()] = true
    if not (parameters["background"] or parameters["b"]) then
        if parameters["verbose"] or parameters["v"] then
            print("Got valid node advert from "..packet[3]..".")
        end
    end
    routerNodeGraph[packet[3]] = packet[4]
    nodeDistanceGraph[packet[3]] = packet[5]
    routerAdvertTimestamps[packet[3]] = os.clock()
    for i=1, #localRouterNodeGraph do
        if localRouterNodeGraph[i] ~= packet[3] then
            oldRednetSend(localRouterNodeGraph[i], textutils.serialize( packet ))
        end
    end
end

local function addNodeToNodeList(node, link) -- Adds [link] to [node]'s routerNodeGraph entry, if not present already
    if routerNodeGraph[node] == nil then
        routerNodeGraph[node] = {}
    end
    for id=0, #routerNodeGraph[node] do
        if routerNodeGraph[node][id] == link then
            return
        end
    end
    table.insert(routerNodeGraph[node], link)
    routerNodeGraph[os.computerID()] = localRouterNodeGraph
    nodeDistanceGraph[os.computerID()] = localNodeDistanceGraph
end

local function addToLocalGraph(node, distance) -- Adds [node] to routerLocalGraph, if not present already
    for id=1, #localRouterNodeGraph do
        if localRouterNodeGraph[id] == node then
            return
        end
    end
    table.insert(localRouterNodeGraph,node)
    localNodeDistanceGraph[node] = distance
    routerNodeGraph[os.computerID()] = localRouterNodeGraph
    nodeDistanceGraph[os.computerID()] = localNodeDistanceGraph
end

local function addNodeLinksFromRouteList(routeList)
    for nodeInChain=1, #routeList, 2 do
        addNodeToNodeList(routeList[nodeInChain], routeList[nodeInChain+1])
        addNodeToNodeList(routeList[nodeInChain+1], routeList[nodeInChain])
    end
end

local function handleIncomingLoop()
    while true do
        routerNodeGraph[os.computerID()] = localRouterNodeGraph
        local id, msg, dist = oldRednetReceive()
        msg = textutils.unserialize(msg)
        if type(msg) == "table" then
            if msg[1] == "ROUTING" then
                if msg[2] == "LocalPing" or msg[2] == "LocalPingAck" then
                    addToLocalGraph(id, dist)
                    if msg[2] == "LocalPing" then
                        if not (parameters["background"] or parameters["b"]) then
                            if parameters["show_pings"] or parameters["p"] then
                                print("Got query ping from "..id..".")
                            end
                        end
                        oldRednetSend(id, pingAckMsg)
                    else
                        if not (parameters["background"] or parameters["b"]) then
                            if parameters["show_pings"] or parameters["p"] then
                                print("Got acknowledgement ping from "..id..".")
                            end
                        end
                    end
                elseif msg[2] == "NodeAdvertisement" then
                    handleNodeAdvertisement(msg)
                elseif msg[2] == "Data" then
                    addNodeLinksFromRouteList(msg[6])
                    if msg[4] == os.computerID() then
                        packets_in = packets_in+1
                        print("Got packet from "..msg[3].." addressed to me.")
                        os.queueEvent("router_message", msg[3], msg[5])
                        os.queueEvent("rednet_message", msg[3], msg[5], getDistanceToNode(id))
                    else
                        packets_thru = packets_thru+1
                        print("Got packet from "..msg[3].." addressed to "..msg[4]..".")
                        for i=2, #msg[6]-1 do
                            if msg[6][i] == os.computerID() then
                                oldRednetSend(msg[6][i+1], textutils.serialize(msg))
                                break
                            end
                        end
                    end
                end
            end
        end
    end
end

local function TimedEventLoop()
    local t = 0
    while true do
        t = t+1
        -- TS check loop:
        if parameters["show_timeout_check"] or parameters["t"] then
            print("Checking advertisement timestamps...")
        end
        for id=0, table.maxn(routerAdvertTimestamps) do
            if routerAdvertTimestamps[id] then
                if (os.clock() - routerAdvertTimestamps[id]) > nodeTimeout then
                    if parameters["verbose"] or parameters["v"] then
                        print("Node "..id.." has timed out, pruning from graph...")
                    end
                    nodeStateAdvertList[id] = nil
                    routerAdvertTimestamps[id] = nil
                    routerNodeGraph[id] = nil
                end
            end
        end

        if parameters["s"] or parameters["hyperspeed"] then
            if t > 5 then
                sendNodeAdvertisement()
                pruneDisconnectedNodes()
                if parameters["v"] or parameters["verbose"] then
                    print("Starting timed event: Local Ping")
                end
                localRouterNodeGraph = {}
                localNodeDistanceGraph = {}
                oldRednetBroadcast(pingMsg)
                t = 0
            end
        else
            if t > 30 then
                sendNodeAdvertisement()
                pruneDisconnectedNodes()
                if parameters["v"] or parameters["verbose"] then
                    print("Starting timed event: Local Ping")
                end
                localRouterNodeGraph = {}
                localNodeDistanceGraph = {}
                oldRednetBroadcast(pingMsg)
                t = 0
            end
        end
        os.sleep(1)
    end
end

local function createRestartWrapper(func, name)
    return function()
        while true do
            local status, err = pcall(func)
            if name then
                if not (parameters["background"] or parameters["bg"]) then
                    print(name.." errored with message: ")
                    print(err)
                    print("Restarting function...")
                end
            end
        end
    end
end

local function handleRouting(id) -- Call the routeFunc with the necessary params
    local route = routeFunc(id, routerNodeGraph, nodeDistanceGraph, parameters)
    -- Do some sanity tests:
    if type(route) == "table" then -- is our route even a table?
        if route[1] == os.computerID() and route[#route] == id then -- Does our route have valid start/end points?
            return route
        end
    end
    return
end

local function sendPacket(id,data)
    if id == os.computerID() then -- we need to catch this because it causes an error. packetSend() does not like sending messages to itself.
        return
    end
    local route = handleRouting(id)
    if not route then
        return false
    end
    packets_out = packets_out+1
    print("Sending to "..route[2]..".")
    return oldRednetSend(route[2], textutils.serialize( {"ROUTING", "Data", os.computerID(), id, data, route} ) )
end

local function getDistanceToNode(node2)
    local route = handleRouting(node2)
    local acc = 0
    for i=2, #route do
        acc = acc + nodeDistanceGraph[ route[i-1] ][ route[i] ]
    end
    return acc
end

-- Hijack rednet.send:
rednet.send = function(id, msg)
    if id == os.computerID() then
        return
    end
    if routerNodeGraph[id] then
        if parameters["v"] or parameters["verbose"] then
            print("Send target ("..id..") found in node graph, attempting packet send...")
        end
        return sendPacket(id, msg)
    end
    if parameters["v"] or parameters["verbose"] then
        print("Send target ("..id..") not found in node graph, attempting old send...")
    end
    return oldRednetSend(id, msg)
end

-- and rednet.receive:
rednet.receive = function(timeout, forceLocalRecv)
    local timer = 0
    local oldFake = parameters["f"]
    if timeout then
        timer = os.startTimer(timeout)
    end
    if forceLocalRecv and oldFake then
        parameters["f"] = false -- disable rednet event faking
    end
    while true do
        local event, id, msg, dist = os.pullEvent()
        if timeout and (event == "timer" and id == timer) then
            if forceLocalRecv and oldFake then
                parameters["f"] = true
            end
            return
        elseif (event == "router_message") and not forceLocalRecv then
            if forceLocalRecv and oldFake then
                parameters["f"] = true
            end
            return id, msg, getDistanceToNode(id)
        elseif event == "rednet_message" then
            local tMsg = textutils.unserialize(msg)
            if type(tMsg) ~= "table" then -- don't worry about this one
                if forceLocalRecv and oldFake then
                    parameters["f"] = true
                end
                return id, msg, dist
            end
            if tMsg[1] ~= "ROUTING" then -- intercept routing messages
                if forceLocalRecv and oldFake then
                    parameters["f"] = true
                end
                return id, msg, dist
            end
        end
    end
end

rednet.broadcast = function(message)
    local localNodes = {}
    for i=1, #localRouterNodeGraph do
        localNodes[ localRouterNodeGraph[i] ] = true
    end
    localNodes[ os.computerID() ] = true
    for id,_ in pairs(routerNodeGraph) do
        if not localNodes[id] then
            print("Broadcasting to "..id..".")
            rednet.send(id, message)
        else
            print("Skipping broadcast to "..id..".")
        end
    end
    oldRednetBroadcast(message)
end

-- add our own functions:
rednet.localNodes = function()
    return localRouterNodeGraph, localNodeDistanceGraph
end

rednet.allNodes = function()
    return routerNodeGraph, nodeDistanceGraph
end

rednet.forceNodeAdvert = function()
    sendNodeAdvertisement()
end

rednet.forceLocalNodeUpdate = function()
    localRouterNodeGraph = {}
    oldRednetBroadcast(pingMsg)
end

rednet.forceNodePruning = function()
    pruneDisconnectedNodes()
end

rednet.forceTimeEventTick = function()
    sendNodeAdvertisement()
    pruneDisconnectedNodes()
    localRouterNodeGraph = {}
    oldRednetBroadcast(pingMsg)
end

rednet.setParameter = function(parameter, state)
    parameters[parameter] = state
    return
end

rednet.getParameter = function(parameter)
    return parameters[parameter]
end

rednet.packetTotals = function()
    return packets_in, packets_thru, packets_out
end

rednet.distanceToID = function(node)
    return getDistanceToNode(node)
end

parallel.waitForAny(createRestartWrapper(handleIncomingLoop, "Incoming Handler Loop"), createRestartWrapper(TimedEventLoop, "Timed Event Loop"), function() os.pullEventRaw("terminate") end)