Deep Freeze

-- Full-Disk Encryption

local startupStub = [[
    local oldOSPull = os.pullEvent
    os.pullEvent = os.pullEventRaw
    os.loadAPI("arcfour")

    local function getAllFilesInDir(dir, includeDirs)
        dir = dir or ""
        local files = {}
        for i,v in ipairs(fs.list(dir)) do
            if fs.isDir(dir.."/"..v) then
                if string.sub(v, 1,4) ~= "disk" and v ~= "rom" then
                    if includeDirs then
                        table.insert(files, v)
                    end
                    for i2, v2 in ipairs(getAllFilesInDir(dir.."/"..v, includeDirs)) do
                        table.insert(files, v.."/"..v2)
                    end
                end
            else
                table.insert(files, v)
            end
        end
        return files
    end

    local function unpackFiles(files)
        for i,v in pairs(files) do
            local fileName = fs.getName(i)
            local nameBegin = string.find(i, fileName)-1
            if nameBegin then
                if not fs.exists(string.sub(i, 1, nameBegin)) then
                    fs.makeDir(string.sub(i, 1, nameBegin))
                end
                local handle = fs.open(i, "w")
                handle.write(v)
                handle.close()
            end
        end
    end

    if fs.exists "HDD" then
        print "This computer's files have been encrypted."
        print "Please enter your password."
        local encrypt_key = read()
        if #encrypt_key == 0 then
            print("Invalid input.")
            os.sleep(2)
            os.reboot()
        end
        local handle = fs.open("HDD", "rb")
        local bytes = {}
        while true do
            local byte = handle.read()
            if byte then
                table.insert(bytes, byte)
            else
                break
            end
        end
        handle.close()
        local ciphertext = string.char(unpack(bytes))
        local files = arcfour.rc4_cipher(encrypt_key, ciphertext)
        files = textutils.unserialize(files)
        if type(files) == "table" then
            unpackFiles(files)
        else
            print "Incorrect key/invalid file!"
            os.sleep(2)
            os.shutdown()
        end
    end

    local oldFSOpen = fs.open
    local oldFSMove = fs.delete
    local oldFSCopy = fs.copy
    local oldFSDelete = fs.delete

    local oldShutdown = os.shutdown
    os.shutdown = function()
        files = getAllFilesInDir("", true)
        for i,v in ipairs(files) do
            if fs.exists(v) and v ~= "startup" and v ~= "HDD" and v ~= "arcfour" then
                oldFSDelete(v)
            end
        end
        oldShutdown()
    end

    local oldReboot = os.reboot
    os.reboot = function()
        files = getAllFilesInDir("", true)
        for i,v in ipairs(files) do
            if fs.exists(v) and v ~= "startup" and v ~= "HDD" and v ~= "arcfour" then
                oldFSDelete(v)
            end
        end
        oldReboot()
    end

    fs.open = function(file, mode)
        if file ~= "startup" and file ~= "arcfour" and file ~= "HDD" then
            return oldFSOpen(file, mode)
        end
    end

    fs.delete = function(file)
        if file ~= "startup" and file ~= "arcfour" and file ~= "HDD" then
            return oldFSDelete(file)
        end
    end

    fs.move = function(from, to)
        if (to ~= "startup" and to ~= "arcfour" and to ~= "HDD") and (from ~= "startup" and from ~= "arcfour" and from ~= "HDD") then
            return oldFSMove(from, to)
        end
    end

    fs.copy = function(from, to)
        if (to ~= "startup" and to ~= "arcfour" and to ~= "HDD") and (from ~= "startup" and from ~= "arcfour" and from ~= "HDD") then
            return oldFSCopy(from, to)
        end
    end

    if fs.exists("startup_2") then
        os.pullEvent = oldOSPull
        shell.run("startup_2")
    else
        shell.run("shell")
    end

    os.shutdown()

]]

local arc4API = [[

-- ARCFOUR implementation in pure Lua
-- Copyright 2008 Rob Kendrick <rjek@rjek.com>
-- Distributed under the MIT licence
--[=[
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated document
ation 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.
]=]--
-- To create a new context;
--  rc4_context = arcfour.new_arc4(optional_key_string)
--
-- To schedule a key
--  rc4_context:schedule(key_string)
--
-- To generate a string of bytes of "random" data from the generator
--  bytes = rc4_context:generate(number_of_bytes)
--
-- To encrypt/decrypt a string
--  ciphertext = rc4_context:cipher(plaintext)
--
-- Example usage;
--  enc_state = arcfour.new "My Encryption Key"
--  enc_state:generate(3072)
--  ciphertext = enc_state:rc4_cipher "Hello, world!"
--
--  dec_state = arcfour.new "My Encryption Key"
--  dec_state:generate(3072)
--  plaintext = dec_state:rc4_cipher(ciphertext)
--
-- Best practise says to discard the first 3072 bytes from the generated
-- stream to avoid leaking information about the key.  Additionally, if using
-- a nonce, you should hash your key and nonce together, rather than
-- concatenating them.

-- Given a binary boolean function b(x,y) defined by a table
-- of four bits { b(0,0), b(0,1), b(1,0), b(1,1) },
-- return a 2D lookup table f[][] where f[x][y] is b() applied
-- bitwise to the lower eight bits of x and y.

local function make_byte_table(bits)
    local f = { }
    for i = 0, 255 do
        f[i] = { }
    end

    f[0][0] = bits[1] * 255

    local m = 1

    for k = 0, 7 do
        for i = 0, m - 1 do
            for j = 0, m - 1 do
                local fij = f[i][j] - bits[1] * m
                f[i  ][j+m] = fij + bits[2] * m
                f[i+m][j  ] = fij + bits[3] * m
                f[i+m][j+m] = fij + bits[4] * m
            end
        end
        m = m * 2
    end

    return f
end

local byte_xor = make_byte_table { 0, 1, 1, 0 }

local function generate(self, count)
    local S, i, j = self.S, self.i, self.j
    local o = { }
    local char = string.char

    for z = 1, count do
        i = (i + 1) % 256
        j = (j + S[i]) % 256
        S[i], S[j] = S[j], S[i]
        o[z] = char(S[(S[i] + S[j]) % 256])
    end

    self.i, self.j = i, j
    return table.concat(o)
end

local function cipher(self, plaintext)
    local pad = generate(self, #plaintext)
    local r = { }
    local byte = string.byte
    local char = string.char

    for i = 1, #plaintext do
        r[i] = char(byte_xor[byte(plaintext, i)][byte(pad, i)])
    end

    return table.concat(r)
end

local function schedule(self, key)
    local S = self.S
    local j, kz = 0, #key
    local byte = string.byte

    for i = 0, 255 do
        j = (j + S[i] + byte(key, (i % kz) + 1)) % 256;
        S[i], S[j] = S[j], S[i]
    end
end

function new_arc4(key)
    local S = { }
    local r = {
        S = S, i = 0, j = 0,
        generate = generate,
        cipher = cipher,
        schedule = schedule
    }

    for i = 0, 255 do
        S[i] = i
    end

    if key then
        r:schedule(key)
    end

    return r
end

function rc4_cipher(key, text)
    local rc4 = new_arc4(key)
    rc4:generate(3072)
    local encrypted = rc4:cipher(text)
    return encrypted
end

]]
local args = {...}

if term.setTextColor == nil or not term.isColor() then
    local oldTermSetTextColor = term.setTextColor
    term.setTextColor = function() -- Pre-1.45 compatibility stub
        return
    end
end


local function getAllFilesInDir(dir, includeDirs)
    dir = dir or ""
    local files = {}
    for i,v in ipairs(fs.list(dir)) do
        if fs.isDir(dir.."/"..v) then
            if string.sub(v, 1,4) ~= "disk" and v ~= "rom" then
                if includeDirs then
                    table.insert(files, v)
                end
                for i2, v2 in ipairs(getAllFilesInDir(dir.."/"..v, includeDirs)) do
                    table.insert(files, v.."/"..v2)
                end
            end
        else
            table.insert(files, v)
        end
    end
    return files
end

local function packFiles(password, includeSelf)
    if fs.exists("startup") then
        fs.move("startup", "startup_2")
    end
    local fileList = getAllFilesInDir("", false)
    local files = {}
    for i,v in ipairs(fileList) do
        if v ~= "arcfour" and v ~= "startup" and (v ~= shell.getRunningProgram() or includeSelf ~= nil) and v ~= "HDD" then
            local handle = fs.open(v, "r")
            local data = handle.readAll()
            handle.close()
            files[v] = data
        end
    end
    local filesStr = textutils.serialize(files)
    local ciphertext = arcfour.rc4_cipher(password, filesStr)
    local cipherbytes = {}
    for i=1, #ciphertext do
        table.insert(cipherbytes, string.byte(ciphertext, i, i))
    end
    if fs.getFreeSpace("/") < #cipherbytes then
        error "Insufficent space"
    end
    local handle = fs.open("HDD", "wb")
    if handle then
        for i=1, #cipherbytes do
            term.clear()
            term.setCursorPos(1,1)
            term.setTextColor(colors.white)
            print("Currently writing to hard drive image. This may take some time.")
            for i=1, term.getSize()-1 do
                write("=")
            end
            print("Byte: "..i.."/"..#cipherbytes)
            local progress = math.ceil(i/#cipherbytes*100)
            print("Progress: "..progress.."%")
            local num_bars = math.floor(progress/2)
            print(" ")
            for i=1, term.getSize()-1 do
                write("=")
            end
            term.setTextColor(colors.lime)
            for i=1, num_bars do
                write("#")
            end
            term.setCursorPos(1, 9)
            term.setTextColor(colors.white)
            for i=1, term.getSize()-1 do
                write("=")
            end
            handle.write(cipherbytes[i])
            os.sleep(0)
        end
        handle.close()
        local startupHandle = fs.open("startup", "w")
        startupHandle.write(startupStub)
        startupHandle.close()
        term.clear()
        term.setCursorPos(1,1)
        write("Writing complete. Restart now? (Y/N) >")
        local yn = read()
        yn = string.upper(yn)
        if yn == "Y" then
            os.reboot()
        end
    else
        error("Couldn't open hard drive image!")
    end
end

if #args == 1 then
    local handle = fs.open("arcfour", "w")
    handle.write(arc4API)
    handle.close()
    os.loadAPI("arcfour")
    if fs.exists "HDD" then
        print("There already appears to be a hard drive image here.")
        print("You cannot encrypt a frozen volume twice.")
        write("Overwrite? (Y/N)> ")
        local yn = read()
        yn = string.upper(yn)
        if yn ~= "Y" then
            error("Cannot freeze volume twice")
        end
    end
    packFiles(args[1], args[2])
else
    print("USAGE: "..fs.getName(shell.getRunningProgram()).." [password] [include encryptor]")
end

local oldTermSetTextColor = term.setTextColor
term.setTextColor = oldTermSetTextColor