-- Function to load DNS records from files
local function loadDNS()
    local dns = {}
    for _, file in ipairs(fs.list("/DNS")) do
        local f = fs.open("/DNS/" .. file, "r")
        local ip = f.readAll()
        f.close()
        dns[file] = ip
    end
    return dns
end

-- Function to load stored passwords from file
local function loadPasswords()
    local passwords = {}
    if fs.exists("passwords.txt") then
        local file = fs.open("passwords.txt", "r")
        for line in file.readLine do
            passwords[line] = true
        end
        file.close()
    end
    return passwords
end

-- Function to load stored authorized computer IDs
local function loadAuthorizedIDs()
    local ids = {}
    if fs.exists("authorized_ids.txt") then
        local file = fs.open("authorized_ids.txt", "r")
        for line in file.readLine do
            ids[line] = true
        end
        file.close()
    end
    return ids
end

-- Function to save a new authorized computer ID
local function saveAuthorizedID(id)
    local file = fs.open("authorized_ids.txt", "a")
    file.writeLine(id)
    file.close()
end

-- Function to log DNS connections
local function logConnection(url, ip)
    local logFile = fs.open("dns_log.txt", "a")
    local logMessage = "URL: " .. url .. ", IP: " .. (ip or "Not found")
    logFile.writeLine(logMessage)
    logFile.close()
    print(logMessage)
end

-- Function to check if a computer ID is already authorized
local function isAuthorizedID(id, ids)
    return ids[tostring(id)] ~= nil
end

-- Function to authenticate a computer with a password
local function authenticate(senderID, password)
    local authorizedPasswords = loadPasswords()  -- Load passwords from file
    if authorizedPasswords[password] then
        saveAuthorizedID(senderID)  -- Save computer ID if authentication is successful
        return true
    else
        return false
    end
end

-- Main DNS Server Loop
local authorizedIDs = loadAuthorizedIDs()  -- Load authorized IDs
rednet.open("top")  -- Assuming the modem is on the top
print("DNS Server running...")

while true do
    local senderID, urlOrPassword = rednet.receive()
    local dns = loadDNS()

    if isAuthorizedID(senderID, authorizedIDs) then
        -- If the computer is already authorized, handle DNS request
        local ip = dns[urlOrPassword]
        logConnection(urlOrPassword, ip)
        if ip then
            rednet.send(senderID, ip)
        else
            rednet.send(senderID, "404_NOT_FOUND")
        end
    else
        -- If not authorized, require password
        rednet.send(senderID, "PASSWORD_REQUIRED")

        local event, id, password = os.pullEvent("rednet_message")
        if id == senderID then
            if authenticate(senderID, password) then
                authorizedIDs[tostring(senderID)] = true  -- Update the in-memory authorized IDs list
                rednet.send(senderID, "AUTH_SUCCESS")
            else
                rednet.send(senderID, "AUTH_FAILED")
            end
        end
    end
end