VeriCode 2 - Easy code signing for ComputerCraft

vericode2.lua

--- VeriCode 2 - Easy code signing for ComputerCraft
-- By JackMacWindows
--
---@module vericode
--
-- Code signing uses encryption and hashes to easily verify a) that the sender of
-- the code is trusted, and b) that the code hasn't been changed mid-transfer.
-- VeriCode applies this concept to Lua code sent over Rednet to add a layer of
-- security to Rednet. Just plainly receiving code from whoever sends it is
-- dangerous, and invites the possibility of getting malware (in fact, I've made
-- a virus that spreads through this method). Adding code signing ensures that
-- any code received is safe and trusted.
--
-- Requires ccryptolib (https://github.com/migeyel/ccryptolib).

--[[ Basic usage:
1. Generate keypair files with vericode.generateKeypair
2. Copy the .key.pub file (NOT the standard .key file!!!) to each client that
   needs to receive signed code
3. Require the API & load the key (.key on server, .key.pub on clients) - on the
   server, make sure to store the key returned from loadKey as you'll need it to send
4. Call vericode.send to send a Lua script to a client computer
5. Call vericode.receive on the client to listen for code from the server (note
   that it returns after receiving a function, so call it in an infinite loop if
   you want it to always accept code)

Example code:

-- On server:
local vericode = require "vericode"
if not fs.exists("mykey.key") then
    vericode.generateKeypair("mykey.key")
    print("Please copy mykey.key.pub to the client computer.")
    return
end
local key = vericode.loadKey("mykey.key")
vericode.send(otherComputerID, "turtle.forward()", key, "turtleInstructions")

-- On client:
local vericode = require "vericode"
vericode.loadKey("mykey.key.pub")
while true do vericode.receive(true, "turtleInstructions") end

--]]

-- MIT License
--
-- Copyright (c) 2024 JackMacWindows
--
-- Permission is hereby granted, free of charge, to any person obtaining a copy
-- of this software and associated documentation files (the "Software"), to deal
-- in the Software without restriction, including without limitation the rights
-- to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-- copies of the Software, and to permit persons to whom the Software is
-- furnished to do so, subject to the following conditions:
--
-- The above copyright notice and this permission notice shall be included in all
-- copies or substantial portions of the Software.
--
-- THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-- IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-- FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-- AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-- LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-- OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-- SOFTWARE.

local function minver(version)
    local res
    if _CC_VERSION then res = version <= _CC_VERSION
    elseif not _HOST then res = version <= os.version():gsub("CraftOS ", "")
    else res = version <= _HOST:match("ComputerCraft ([0-9%.]+)") end
    assert(res, "This program requires ComputerCraft " .. version .. " or later.")
end

minver "1.91.0" -- string.pack, string.unpack

local expect = require "cc.expect"
local ed25519 = require "ccryptolib.ed25519"
local random = require "ccryptolib.random"

local vericode = {}
local keyStore = {}

local b64str = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"

local function base64encode(str)
    local retval = ""
    for s in str:gmatch "..." do
        local n = s:byte(1) * 65536 + s:byte(2) * 256 + s:byte(3)
        local a, b, c, d = bit32.extract(n, 18, 6), bit32.extract(n, 12, 6), bit32.extract(n, 6, 6), bit32.extract(n, 0, 6)
        retval = retval .. b64str:sub(a+1, a+1) .. b64str:sub(b+1, b+1) .. b64str:sub(c+1, c+1) .. b64str:sub(d+1, d+1)
    end
    if #str % 3 == 1 then
        local n = str:byte(-1)
        local a, b = bit32.rshift(n, 2), bit32.lshift(bit32.band(n, 3), 4)
        retval = retval .. b64str:sub(a+1, a+1) .. b64str:sub(b+1, b+1) .. "=="
    elseif #str % 3 == 2 then
        local n = str:byte(-2) * 256 + str:byte(-1)
        local a, b, c, d = bit32.extract(n, 10, 6), bit32.extract(n, 4, 6), bit32.lshift(bit32.extract(n, 0, 4), 2)
        retval = retval .. b64str:sub(a+1, a+1) .. b64str:sub(b+1, b+1) .. b64str:sub(c+1, c+1) .. "="
    end
    return retval
end

local function base64decode(str)
    local retval = ""
    for s in str:gmatch "...." do
        if s:sub(3, 4) == '==' then
            retval = retval .. string.char(bit32.bor(bit32.lshift(b64str:find(s:sub(1, 1)) - 1, 2), bit32.rshift(b64str:find(s:sub(2, 2)) - 1, 4)))
        elseif s:sub(4, 4) == '=' then
            local n = (b64str:find(s:sub(1, 1))-1) * 4096 + (b64str:find(s:sub(2, 2))-1) * 64 + (b64str:find(s:sub(3, 3))-1)
            retval = retval .. string.char(bit32.extract(n, 10, 8)) .. string.char(bit32.extract(n, 2, 8))
        else
            local n = (b64str:find(s:sub(1, 1))-1) * 262144 + (b64str:find(s:sub(2, 2))-1) * 4096 + (b64str:find(s:sub(3, 3))-1) * 64 + (b64str:find(s:sub(4, 4))-1)
            retval = retval .. string.char(bit32.extract(n, 16, 8)) .. string.char(bit32.extract(n, 8, 8)) .. string.char(bit32.extract(n, 0, 8))
        end
    end
    return retval
end

vericode.base64 = {encode = base64encode, decode = base64decode}
vericode.random = random
vericode.ed25519 = ed25519

if not random.isInit() then random.initWithTiming() end

--- Generates a keypair for code signing.
-- Outputs a pub/priv keypair at `path`, and a public-only key (for receivers) at `path`.pub.
-- The generated key will be added to the store.
---@param path string The path to the file to generate.
---@return string pub The new public key.
---@return string priv The new private key. (Not required for this API, but might be useful otherwise.)
function vericode.generateKeypair(path)
    expect(1, path, "string")
    local priv = random.random(32)
    local pub = ed25519.publicKey(priv)
    pub, priv = base64encode(string.char(table.unpack(pub))), base64encode(string.char(table.unpack(priv)))
    local file, err = fs.open(path, "w")
    if not file then error("Could not open certificate file: " .. err, 2) end
    file.write(textutils.serialize({
        public = pub,
        private = priv
    }))
    file.close()
    file, err = fs.open(path .. ".pub", "w")
    if not file then error("Could not open public certificate file: " .. err, 2) end
    file.write(textutils.serialize({
        public = pub
    }))
    file.close()
    keyStore[pub] = {
        public = pub,
        private = priv
    }
    return pub, priv
end

--- Loads a key from disk. This can be a full keypair, or only a public key.
---@param path string The path to the key.
---@return string key The loaded public key.
function vericode.loadKey(path)
    expect(1, path, "string")
    local file, err = fs.open(path, "r")
    if not file then error("Could not open certificate file: " .. err, 2) end
    local t = textutils.unserialize(file.readAll())
    file.close()
    if type(t) ~= "table" or t.public == nil then error("Invalid certificate file", 2) end
    keyStore[t.public] = t
    return t.public
end

--- Loads all keys in the specified folder.
---@param path string The path to the folder to check
function vericode.loadKeys(path)
    expect(1, path, "string")
    if not fs.isDir(path) then error(path .. ": Not a directory", 2) end
    for _, p in ipairs(fs.list(path)) do if p:match "%.key%.?p?u?b?$" then vericode.loadKey(fs.combine(path, p)) end end
end

--- Adds a public (and private if provided) key to the key store.
---@param pub string The public key to add.
---@param priv string|nil The private key to add, if desired.
function vericode.addKey(pub, priv)
    expect(1, pub, "string")
    expect(2, priv, "string", "nil")
    keyStore[pub] = {
        public = pub,
        private = priv
    }
end

--- Signs a public key file using a master private key. This can be used to sign code using a key which the receiver does not know.
---@param path string The path to the public key to sign
---@param priv string The private key to sign with (which may the corresponding public key if loaded)
---@return string pub The newly signed public key loaded from the file
function vericode.signKey(path, priv)
    expect(1, path, "string")
    expect(2, priv, "string", "nil")
    local pub, priv
    if keyStore[key] then
        if not keyStore[key].private then error("No private key associated with selected public key", 2) end
        pub, priv = key, keyStore[key].private
    else
        for _,v in pairs(keyStore) do
            if v.public == key then
                if not v.private then error("No private key associated with selected public key", 2) end
                pub, priv = key, v.private
                break
            elseif v.private == key then
                pub, priv = v.public, key
                break
            end
        end
    end
    if not pub or not priv then error("Could not find private key", 2) end
    local file, err = fs.open(path, "r")
    if not file then error("Could not open certificate file: " .. err, 2) end
    local t = textutils.unserialize(file.readAll())
    file.close()
    if type(t) ~= "table" or t.public == nil then error("Invalid certificate file", 2) end
    t.signer = pub
    t.signature = ed25519.sign(priv, pub, t.public)
    file, err = fs.open(path, "w")
    if not file then error("Could not open certificate file: " .. err, 2) end
    file.write(textutils.serialize(t))
    file.close()
    keyStore[t.public] = t
    return t.public
end

---@class VeriCode A chunk of signed code. Don't use the members of this table directly!
---@field code string The code delivered in the payload
---@field key string The public key used to sign the code
---@field signature string The signature of the code using the key specified
---@field keySigner string|nil The signer of the key, if required
---@field keySignature string|nil The signature of the key, if required

--- Signs a Lua chunk into a table.
---@param code string The Lua code to compile.
---@param key string The public or private key to use. The private key associated with this key must exist in the key store.
---@return VeriCode chunk A signed Lua chunk. This chunk can be loaded with `vericode.load`.
function vericode.dump(code, key)
    expect(1, code, "string")
    expect(2, key, "string")
    local pub, priv, kpub, ksig
    local t = keyStore[key]
    if t then
        if not t.private then error("No private key associated with selected public key", 2) end
        pub, priv, kpub, ksig = key, t.private, t.signer, t.signature
    else
        for _,v in pairs(keyStore) do
            if v.public == key then
                if not v.private then error("No private key associated with selected public key", 2) end
                pub, priv, kpub, ksig = key, v.private, v.signer, v.signature
                break
            elseif v.private == key then
                pub, priv, kpub, ksig = v.public, key, v.signer, v.signature
                break
            end
        end
    end
    if not pub or not priv then error("Could not find private key", 2) end
    ---@type VeriCode
    return {
        code = code,
        key = pub,
        signature = ed25519.sign(priv, pub, code),
        keySigner = kpub,
        keySignature = ksig
    }
end

--- Loads and verifies a previously signed code chunk.
-- The public key associated with the chunk must be present in the key store.
---@param code VeriCode The code chunk to load.
---@param name string|nil The name of the chunk.
---@param _mode nil Ignored (for compatibility).
---@param env table|nil The environment to give the chunk.
---@return function|nil fn The returned function, or nil on error.
---@return nil|string err If an error occurred, the error message.
function vericode.load(code, name, _mode, env)
    expect(1, code, "table")
    expect.field(code, "code", "string")
    expect.field(code, "key", "string")
    expect.field(code, "signature", "string")
    expect(2, name, "string", "nil")
    expect(4, env, "table", "nil")
    if not keyStore[code.key] and not (type(code.keySigner) == "string" and type(code.keySignature) == "string" and keyStore[code.keySigner] and ed25519.verify(code.keySigner, code.key, code.keySignature)) then
        return nil, "Unrecognized key: " .. base64encode(code.key)
    end
    if not ed25519.verify(code.key, code.code, code.signature) then return nil, "Invalid code signature" end
    return load(code.code, name, "t", env)
end

--- Sends a signed code chunk over Rednet.
---@param recipient number The ID of the recipient.
---@param code string The code chunk to send.
---@param key string The key to use to sign the chunk.
---@param protocol string|nil The protocol to set, if desired.
---@return boolean ok Whether the message was sent.
function vericode.send(recipient, code, key, protocol)
    expect(1, recipient, "number")
    expect(2, code, "string")
    expect(3, key, "string")
    expect(4, protocol, "string", "nil")
    return rednet.send(recipient, vericode.dump(code, key), protocol)
end

--- Broadcasts a signed code chunk over Rednet.
---@param code string The code chunk to send.
---@param key string The key to use to sign the chunk.
---@param protocol string|nil The protocol to set, if desired.
---@return boolean ok Whether the message was sent.
function vericode.broadcast(code, key, protocol)
    expect(1, code, "string")
    expect(2, key, "string")
    expect(3, protocol, "string", "nil")
    return rednet.broadcast(vericode.dump(code, key), protocol)
end

--- Waits to receive a signed code chunk, and either returns the loaded function or the results from calling it.
---@param run boolean|nil Whether to run the code, or just return the function.
---@param filter string|nil The name of the protocol to listen for (nil for any).
---@param timeout number|nil The maximum amount of time to wait.
---@param name string|nil The name to give the loaded chunk (defaults to "=VeriCode chunk").
---@param env table|nil The environment to give the function.
---@return any res Either the loaded function, or the results from the function, or nil if the timeout was passed.
function vericode.receive(run, filter, timeout, name, env)
    expect(1, run, "boolean", "nil")
    expect(2, filter, "string", "nil")
    expect(3, timeout, "number", "nil")
    expect(4, name, "string", "nil")
    expect(5, env, "table", "nil")
    local res = {n = 0}
    local function receive()
        while true do
            local _, message = rednet.receive(filter)
            if type(message) == "table" and type(message.code) == "string" and type(message.key) == "string" and type(message.signature) == "string" then
                local fn = vericode.load(message, name or "=VeriCode chunk", nil, env)
                if fn then
                    if run then res = table.pack(fn())
                    else res = {fn, n = 1} end
                    return table.unpack(res, 1, res.n)
                end
            end
        end
    end
    if timeout then
        parallel.waitForAny(receive, function() sleep(timeout) end)
        return table.unpack(res, 1, res.n)
    else return receive() end
end

if ... then vericode.generateKeypair(...) end

return vericode