term.clear()
term.setCursorPos(1, 1)
-- Find modems
for _, sModem in ipairs(peripheral.getNames()) do
    if peripheral.getType(sModem) == "modem" then rednet.open(sModem) end
end

local knownHosts = settings.get("knownhosts")
if knownHosts == nil then
	settings.set("knownhosts", {})
	settings.save()
end

local isClient = true
local PK = nil
local crypto = require("ccryptolib")
crypto.random.initWithTiming()
local to = arg[1] -- Who we're trying to talk to (computer id)
write("Your Nickname: ")
local from = read()
if to == nil then print("You need to provide a computer ID!") end
local hexPrivate = settings.get("keys.private")
local hexPublic = settings.get("keys.public")

if hexPrivate == nil then
    settings.set("keys.private", crypto.util.toHex(crypto.random.random(32)))
    hexPrivate = settings.get("keys.private")
	settings.save()
end

if hexPublic == nil then
	settings.set("keys.public", crypto.util.toHex(crypto.x25519.publicKey(crypto.util.fromHex(settings.get("keys.private")))))
	hexPublic = settings.get("keys.public")
    settings.save()
end

local privateKey = crypto.util.fromHex(hexPrivate)
local publicKey = crypto.util.fromHex(hexPublic)

local termX, termY = term.getSize()

function string.starts(String, Start)
    return string.sub(String, 1, string.len(Start)) == Start
end

function rx()
    while true do
        sleep(0)
        sender, rawData = rednet.receive("encChat")
        if sender ~= tonumber(to) then return end
        data = textutils.unserialise(crypto.chacha20.crypt(PK, rawData.nonce, rawData.data))
        if not data.test then return end -- Failed decrypt, possible bad key exchange?
        term.setCursorPos(1, termY)
        term.clearLine()
        print(data.from .. ":" .. data.msg)
		term.blit("> ", "30", "ff")
        term.setCursorPos(3, termY)
    end
end

function tx()
    while true do
        sleep(0)
        term.setCursorPos(1, termY)
        term.blit("> ", "30", "ff")
        local msg = read()
        if (string.starts(msg, "/")) then
            -- Command handler
			msg = string.lower(msg)
			if msg == ("/q" or "/quit" or "/exit") then
				error("Goodbye!")
			end
        else
            term.scroll(-1)
            term.clearLine()
            write(from .. ":" .. msg)
            term.scroll(1)
            nonce = crypto.random.random(12)
            msgData = {
                timestamp = os.epoch("utc"),
                msg = msg,
                from = from,
                test = true
            }
            rednet.send(tonumber(to), {
                nonce = nonce,
                data = crypto.chacha20
                    .crypt(PK, nonce, textutils.serialise(msgData))
            }, "encChat")
        end
    end
end

function ping()
    local fails = 0
    if isClient then
        while true do
            sleep(1)
            nonce = crypto.random.random(12)
            msgData = {timestamp = os.epoch("utc"), test = true}
            rednet.send(tonumber(to), {
                nonce = nonce,
                data = crypto.chacha20
                    .crypt(PK, nonce, textutils.serialise(msgData))
            }, "encPing")

            -- Wait for reply
            sender, reply = rednet.receive("encPingReply", 5)
            if fails > 2 then
                error(
                    "No response from other user, They likely closed their client exiting...")
            end
            if sender == nil or data == nil then
                fails = fails + 1
            elseif crypto.chacha20.crypt(PK, reply.nonce, reply.data) ~=
                tostring(msgData.timestamp) then
                error("Invalid reply from ping, exiting for security.")
            else
                fails = 0
            end
        end
    else
        while true do
            sleep(1)
            sender, data = rednet.receive("encPing", 10)
            if fails > 1 then
                error("No response from other user, They likely closed their client exiting...")
            end
            if sender == nil or data == nil then
                fails = fails + 1
            elseif sender == tonumber(to) then
                timestamp = textutils.unserialise(crypto.chacha20.crypt(PK, data.nonce, data.data)).timestamp
                if os.epoch("utc") > timestamp then
                    fails = 0
                    local nonce = crypto.random.random(12)
                    rednet.send(sender, {
                        nonce = nonce,
                        data = crypto.chacha20
                            .crypt(PK, nonce, tostring(timestamp))
                    }, "encPingReply")
                end
            end
        end
    end
end

function exchange()
    print("sending ping to " .. to)
    rednet.send(tonumber(to), publicKey, "keyExClient")
    sender, data = rednet.receive("keyEx", 1)
    if sender ~= tonumber(to) then -- Switch to server
        print("no response from " .. to .. ". listening for ping...")
        isClient = false
        while true do
            sender, data = rednet.receive("keyExClient")
            if sender == tonumber(to) then
                print("ping from " .. sender .. ". sending and generating keys...")
                rednet.send(tonumber(to), publicKey, "keyEx")
				if knownHosts[tonumber(to)] == nil then
					write("remote pubkey is " .. crypto.util.toHex(data) ..". Do you trust it? (Y/N)")
					local trust = string.lower(read())
					if trust == "y" then
						knownHosts[tonumber(to)] = data
						settings.set("knownhosts", knownHosts)
						settings.save()
						PK = crypto.x25519.exchange(privateKey, data)
					else
						error("Exiting...")
					end
				elseif knownHosts[tonumber(to)] ~= data then
					error("==DANGER== Remote pubkey differs from previous interaction!")
				else
					PK = crypto.x25519.exchange(privateKey, data)
				end
                break
            end
        end
    else
        print("response from " .. to)
		if knownHosts[tonumber(to)] == nil then
			write("remote pubkey is " .. crypto.util.toHex(data) ..". Do you trust it? (Y/N)")
			local trust = string.lower(read())
			if trust == "y" then
				knownHosts[tonumber(to)] = data
				settings.set("knownhosts", knownHosts)
				settings.save()
				PK = crypto.x25519.exchange(privateKey, data)
			else
				error("Exiting...")
			end
		elseif knownHosts[tonumber(to)] ~= data then
			error("==DANGER== Remote pubkey differs from previous interaction!")
		else
			PK = crypto.x25519.exchange(privateKey, data)
		end
        PK = crypto.x25519.exchange(privateKey, data)
    end
	term.clear()
	term.setCursorPos(1,1)
    print("==BEGIN ENCRYPTED CONVERSATION==")
end

exchange()

parallel.waitForAll(rx, tx, ping)