-- Function to generate a key for encryption and decryption
function generateKey()
    local chars = {}
    for i = 32, 126 do -- ASCII range for printable characters
        table.insert(chars, string.char(i))
    end

    local key = {}
    while #chars > 0 do
        local index = math.random(#chars)
        table.insert(key, table.remove(chars, index))
    end

    return table.concat(key)
end

-- Function to create a map from a key
function createMap(key)
    local map = {}
    for i = 32, 126 do
        map[string.char(i)] = key:sub(i - 31, i - 31)
    end
    return map
end

-- Function to invert a map
function invertMap(map)
    local invMap = {}
    for k, v in pairs(map) do
        invMap[v] = k
    end
    return invMap
end

-- Function to encrypt text using a key
function encrypt(text, key)
    local map = createMap(key)
    local encrypted = {}
    for i = 1, #text do
        local char = text:sub(i, i)
        table.insert(encrypted, map[char] or char)
    end
    return table.concat(encrypted)
end

-- Function to decrypt text using a key
function decrypt(text, key)
    local map = invertMap(createMap(key))
    local decrypted = {}
    for i = 1, #text do
        local char = text:sub(i, i)
        table.insert(decrypted, map[char] or char)
    end
    return table.concat(decrypted)
end

-- Function to save the key to a file
function saveKeyToFile(key)
    local randomNumber = math.random(1, 1000)
    local fileName = "key" .. randomNumber .. ".txt"
    local file = fs.open(fileName, "w")
    file.write(key)
    file.close()
    return fileName
end

-- Function to load the key from a file
function loadKeyFromFile(fileName)
    local file = fs.open(fileName, "r")
    local key = file.readAll()
    file.close()
    return key
end

-- Function to traverse and process files in a directory
function processDirectory(path, processFunction, key)
    local files = fs.list(path)
    for _, file in ipairs(files) do
        local filePath = fs.combine(path, file)
        if fs.isDir(filePath) then
            processDirectory(filePath, processFunction, key)
        else
            local file = fs.open(filePath, "r")
            local content = file.readAll()
            file.close()

            local processedContent = processFunction(content, key)

            file = fs.open(filePath, "w")
            file.write(processedContent)
            file.close()
        end
    end
end

-- Main Program
local tArgs = {...}
if #tArgs < 1 then
    print("Usage: jambler <encrypt|decrypt> [keyfile (for decrypt)]")
    return
end

local mode = tArgs[1]

if mode == "encrypt" then
    local key = generateKey()
    processDirectory("/disk", encrypt, key)
    local keyFileName = saveKeyToFile(key)
    print("All files in /disk/ encrypted.")
    print("Key saved to file: " .. keyFileName)
elseif mode == "decrypt" then
    if #tArgs < 2 then
        print("A key file is required for decryption.")
        return
    end
    local keyFileName = tArgs[2]
    if not fs.exists(keyFileName) then
        print("Key file does not exist.")
        return
    end
    local key = loadKeyFromFile(keyFileName)
    processDirectory("/disk", decrypt, key)
    print("All files in /disk/ decrypted.")
else
    print("Invalid mode. Use 'encrypt' or 'decrypt'.")
end