DogLocker

-- Helper function to generate a seed-based pseudo-random number generator
function createPRNG(seed)
    local state = seed
    return function()
        state = (state * 16807) % 2147483647  -- A simple PRNG algorithm (multiplicative congruential)
        return state
    end
end

-- XOR encryption function
function xorEncrypt(data, key)
    local encryptedData = {}
    local prng = createPRNG(key)  -- PRNG with the key
    for i = 1, #data do
        local keyByte = prng() % 256  -- Get a random byte for XORing
        encryptedData[i] = string.char(bit32.bxor(string.byte(data, i), keyByte))
    end
    return table.concat(encryptedData)
end

-- XOR decryption function (symmetric)
function xorDecrypt(encryptedData, key)
    return xorEncrypt(encryptedData, key)  -- Decryption is the same as encryption with XOR
end

-- Encrypt the entire disk directory (all files and subdirectories)
function encryptDiskDirectory(diskPath, key)
    for _, fileName in pairs(fs.list(diskPath)) do
        local fullPath = diskPath .. "/" .. fileName
        if fs.isDir(fullPath) then
            encryptDiskDirectory(fullPath, key)  -- Recursively encrypt subdirectories
        else
            local file = fs.open(fullPath, "r")
            local data = file.readAll()
            file.close()

            -- Encrypt the file data and overwrite it
            local encryptedFile = xorEncrypt(data, key)
            local encryptedFilePath = fullPath .. ".enc"
            local encFile = fs.open(encryptedFilePath, "w")
            encFile.write(encryptedFile)
            encFile.close()

            -- Optionally, delete the original file after encryption
            fs.delete(fullPath)
        end
    end
end

-- Decrypt the entire disk directory (all files and subdirectories)
function decryptDiskDirectory(diskPath, key)
    for _, fileName in pairs(fs.list(diskPath)) do
        local fullPath = diskPath .. "/" .. fileName
        if fs.isDir(fullPath) then
            decryptDiskDirectory(fullPath, key)  -- Recursively decrypt subdirectories
        else
            local file = fs.open(fullPath, "r")
            local data = file.readAll()
            file.close()

            -- Decrypt the file data and overwrite it
            local decryptedFile = xorDecrypt(data, key)
            local decryptedFilePath = fullPath:sub(1, -5)  -- Remove the ".enc" suffix
            local decFile = fs.open(decryptedFilePath, "w")
            decFile.write(decryptedFile)
            decFile.close()

            -- Optionally, delete the encrypted file after decryption
            fs.delete(fullPath)
        end
    end
end

-- Load disk encryption key and data from the DogLocker subdirectory
function loadDiskData(diskPath)
    local dirPath = "/disk/DogLocker/" .. diskPath:sub(6)  -- Remove the "/disk" part of the path to use the name

    if not fs.exists(dirPath) then
        return nil  -- Disk directory not found
    end

    -- Load the data from files in the subdirectory
    local labelFile = fs.open(dirPath .. "/label.txt", "r")
    local passwordFile = fs.open(dirPath .. "/password.txt", "r")

    local label = labelFile.readLine()
    local encryptedPassword = passwordFile.readLine()  -- Get the encrypted password
    labelFile.close()
    passwordFile.close()

    -- Generate key from the disk path length (it will match what was used in encryption)
    local key = diskPath:len()

    return { label = label, encryptedPassword = encryptedPassword, key = key }
end

-- Store disk encryption data (Disk data, password, label) in the DogLocker subdirectory
function storeDiskData(diskPath, label, password)
    local dirPath = "/disk/DogLocker/" .. diskPath:sub(6)  -- Remove the "/disk" part of the path to use the name

    if not fs.exists(dirPath) then
        fs.makeDir(dirPath)
    end

    -- Encrypt the password before storing it
    local encryptedPassword = xorEncrypt(password, diskPath:len())

    -- Save the label and encrypted password
    local labelFile = fs.open(dirPath .. "/label.txt", "w")
    local passwordFile = fs.open(dirPath .. "/password.txt", "w")

    labelFile.writeLine(label)
    passwordFile.writeLine(encryptedPassword)

    labelFile.close()
    passwordFile.close()
end

-- Encrypt and store the entire disk directory
function encryptDiskData(diskPath)
    local diskData = loadDiskData(diskPath)
    if not diskData then
        error("Disk data not found.")
    end

    encryptDiskDirectory(diskPath, diskData.key)
end

-- Decrypt the entire disk directory
function decryptDiskData(diskPath, password)
    local diskData = loadDiskData(diskPath)
    if not diskData then
        error("Disk data not found.")
    end

    -- Decrypt the password with the disk path length as the key
    local decryptedPassword = xorDecrypt(diskData.encryptedPassword, diskPath:len())

    -- Validate password input when unlocking
    if decryptedPassword ~= password then
        error("Incorrect password.")
    end

    decryptDiskDirectory(diskPath, diskData.key)
end

-- List connected drives
function listConnectedDrives()
    local drives = {}
    local driveCount = 0
    for _, disk in pairs(fs.list("/disk/")) do
        if fs.isDir("/disk/" .. disk) then
            driveCount = driveCount + 1
            table.insert(drives, "/disk/" .. disk)
        end
    end
    return drives
end

-- UI for protecting and unlocking drives
function displayUI()
    print("Welcome to DogLocker!")
    print("1. Protect (Encrypt) Drive")
    print("2. Unlock (Decrypt) Drive")
    print("3. Exit")

    local choice = read()

    if choice == "1" then
        protectDriveUI()
    elseif choice == "2" then
        unlockDriveUI()
    elseif choice == "3" then
        print("Exiting...")
        return
    else
        print("Invalid choice. Please try again.")
        displayUI()
    end
end

-- UI for protecting a drive (encrypting data)
function protectDriveUI()
    -- Ask the user for the exact path to encrypt
    print("Enter the full path of the drive to protect (e.g., /disk/, /disk2/):")
    local diskPath = read()

    if not fs.exists(diskPath) or not fs.isDir(diskPath) then
        print("Invalid disk path. Please try again.")
        return
    end

    print("You selected: " .. diskPath)

    print("Enter label for the disk:")
    local label = read()

    print("Enter password to protect the disk:")
    local password = read("#")  -- Hide password input

    -- Store the disk's information and encrypt its data
    storeDiskData(diskPath, label, password)
    encryptDiskData(diskPath)

    print("Drive protected successfully.")
    displayUI()
end

-- UI for unlocking a drive (decrypting data)
function unlockDriveUI()
    -- Ask the user for the exact path to unlock
    print("Enter the full path of the drive to unlock (e.g., /disk/, /disk2/):")
    local diskPath = read()

    if not fs.exists(diskPath) or not fs.isDir(diskPath) then
        print("Invalid disk path. Please try again.")
        return
    end

    print("You selected: " .. diskPath)

    print("Enter password to unlock the disk:")
    local password = read("#")  -- Hide password input

    -- Decrypt the disk's data
    local status, err = pcall(function() decryptDiskData(diskPath, password) end)

    if status then
        print("Drive unlocked successfully.")
    else
        print("Error: " .. err)
    end

    displayUI()
end

-- Start the UI
displayUI()