Secure Network Router

-- Load APIs
os.loadAPI("disk/secunet/apis/uuid")
os.loadAPI("disk/secunet/apis/base64")
os.loadAPI("disk/secunet/apis/aeslua")

-- Variables
local userdata = {} -- (Array of usernames as keys with values as table of userdata [HashMap<String, HashMap<String, String>>])
local usernameByPin = {} -- Look up username by pin
local hosts = {} -- Similar to DNS : Lookup usrname by hostname
local reverseHosts = {} -- Reverse hosts
local modem = peripheral.wrap("top")

-- Thanks to Lyqyd for this function
local function split(input)

    local results = {}

    for regexmatch in string.gmatch(input, "[^ ]+") do
        table.insert(results, regexmatch)
    end

    return results
end

-- Thanks to Lyqyd for this function
local function splitTab(input)

    local results = {}

    for regexmatch in string.gmatch(input, "[^\t]+") do -- edit
        table.insert(results, regexmatch)
    end

    return results
end


-- Generates IV table
local function generateIV()

    local index = 1
    local iv = {}

    repeat

    iv[index] = math.random(1, 255)
    index = index + 1

    until index == 17

    return iv

end

-- Encrypt
local function encrypt(key, message, iv)
    return aeslua.encrypt(key, message, aeslua.AES256, aeslua.CBCMODE, iv)
end

-- Decrypt
local function decrypt(key, message, iv)
    return aeslua.decrypt(key, message, aeslua.AES256, aeslua.CBCMODE, iv)
end

-- Saves user data
local function saveUserdata(password)

    local userfile = assert(fs.open(shell.dir() .. "/users/server.dat", "w"))
    local iv = generateIV()
    local data = textutils.serialize(iv) .. "\t".. base64.enc(encrypt(password, textutils.serialize(userdata), iv))
    userfile.write(data)
    userfile.close()

end

-- Open and decrypt users' login details
local function getUserdata(password)
    if fs.exists(shell.dir() .. "/users/server.dat") ~= true then return end
    local userfile = fs.open(shell.dir() .. "/users/server.dat", "r")
    local userfileData = userfile.readAll()
    local userfileIvSerialized = splitTab(userfileData)[1]
    local userfileDataSplit = splitTab(userfileData)
    table.remove(userfileDataSplit, 1)
    local userfileDataEncrypted = base64.dec(table.concat(userfileDataSplit, " "))
    local iv = textutils.unserialize(userfileIvSerialized)
    local message = userfileDataEncrypted
    local cleartextDetails = decrypt(password, message, iv)
    if cleartextDetails == nil then
        error()
    end
    userfile.close()
    userdata = textutils.unserialize(cleartextDetails)
    print("Successfully loaded " .. table.getn(userdata) .. " users")

end

-- Returns 16 bit random
local function random128()
    local charset = {"0","1","2","3","4","5","6","7","8","9","a","b","c","d","e","f"}
    return tostring(math.random(10000000, 99999999))..tostring(math.random(10000000, 99999999))..tostring(math.random(10000000, 99999999))..tostring(math.random(10000000, 99999999))
end

-- Generates a user
local function generateUser(username, hostnames) -- Hostnames is an array of hostnames

    local user = {}
    user["pin"] = uuid.Generate()
    user["msgpassword"] = random128()
    user["hmacpassword"] = random128()
    user["username"] = username

    for key, value in pairs(hostnames) do
        hosts[value] = username
    end

    hosts[username] = username
    reverseHosts[username] = hostnames
    usernameByPin[user["pin"]] = username

    return user

end

-- Split cleartext data for server
local function splitDataServer(cleartext)

    -- Data table
    data = {}

    -- Split cleartext
    cleartextSplit = split(cleartext)

    -- Put into data table
    data["hmac"] = table.remove(cleartextSplit, 1)
    data["destination"] = table.remove(cleartextSplit, (1))
    data["nextmsgpassword"] = table.remove(cleartextSplit, (1))
    data["nexthashpasswd"] = table.remove(cleartextSplit, (1))
    data["message"] = table.concat(cleartextSplit)

    return data

end

-- Lookup usrname from hostname
local function lookup(name)
    return hosts[name]
end

-- Reverse lookup
local function reverseLookup(username)
    return reverseHosts[username]
end

-- Extracts data from packet
local function handleData(packet)

    -- Split message
    local messagesplit = split(packet)

    -- Filter for spam
    if usernameByPin[messagesplit[1]] ~= nil then

        -- remove pin from message
       table.remove(messagesplit, 1)

        -- Decrypt data
        local decryptedData = decrypt(userdata["messagepassword"], base64.dec(messagesplit), textutils.unserialize(table.remove(messagesplit,1)))

        -- Split data
        local splitErrorHappened, data = pcall(splitDataServer, decryptedData)

        if splitErrorHappened ~= true then
            error("Invalid packet")

        -- Split successful
        else

            -- Check hmac
            local hmacErrorHappened, valid = pcall(data["hmac"], data["message"])

            if hmacErrorHappened ~= true then
                error("Error comparing HMACs!")
            end

            if valid then
                return data

            else
                error("HMAC invalid!")
            end

        end
    end
end

-- Send message
local function send(message, destinationip, sender)

    -- Check for connection
    assert(connected, "Not connected to server!")

    -- Create header
    local pin = userdata["pin"]
    local destination = destinationip
    local nextmsgpasswd = random128()
    local nexthashpasswd = random128()
    local nextpin = uuid.Generate()
    local hmac = sha.hmac(message, userdata[destinationuser]["hmacpassword"])
    local iv = generateIV()

    -- Create encrypted message body
    local messagebody = " " .. base64.enc(encrypt(userdata[destinationuser]["msgpassword"], hmac .. "" .. sender .. " " .. nextpin .. " " .. msgpasswd .. " " .. hashpasswd .. " " .. userdata[destinationuser]["msgpassword"] ..  message))

    -- Concat pin with messagebody
    local message = pin .. " " .. iv .. " " .. messagebody

    -- Transmit to server
    modem.transmit(channel, channel, message)

    -- Save the next passwords and pins
    userdata["messagepassword"] = nextmspasswd
    userdata["hmacpassword"] = nexthashpasswd
    userdata["pin"] = nextpin


end

-- Listen for packets
local function listen()
    while true do
        os.queueEvent(".")
        os.pullEvent(".")
        modem.open(4000)
        local event = {os.pullEvent()} -- event[1] = event --Wait for event
        if event[1] == "modem_message" then -- event[2] = modemside; event[3] = senderchannel; event[4] = replychannel; event[5] = message; event[6] = senderDistance;
            print(event[5])
            local handleDataErrorHappened, data = pcall(handleData, event[5])

            if handleDataErrorHappened ~=  true then

            else

                local usrname = usernameByPin[data["pin"]]

                userdata[usrname]["hmacpassword"] = messageData["nexthashpasswd"]
                userdata[usrname]["msgpassword"] = messageData["nextmsgpasswd"]
                userdata[usrname]["pin"] = messageData["nextpin"]

                lookupErrorHappened, username = pcall(lookup(data["destination"]))

                if (lookupErrorHappened ~= true) then

                elseif (username ~= nil) then
                    send(data["message"], data["destination"], usrname)
                end
            end
        end
    end
end


-- Shell (runs for 1 cycle)
local function usrshell()
    io.write("> ")
    input = read()

    if input == "end" then
        io.write("Enter server password: ")
        saveUserdata(read("*"))
        os.queueEvent("terminate")
    end

    local usr = generateUser(input, {input})
    print("User generated.")
    io.write("Enter server password: ")
    local pass = read("*")
    userdata[input] = usr
    saveUserdata(pass)
end

-- Main thread
function main()

    -- Gets server data
    repeat
        io.write("Enter server password: ")
        password = read("*")
        success, return_value = pcall(getUserdata, password)
        if success ~= true then print("Error! Wrong password") end
    until success == true

    print(textutils.serialize(userdata))

    while true do
         parallel.waitForAny(listen, usrshell)
    end


end

main()