local program_name = "BlueSSH"
local program_version = "1.0"
local program_string = program_name .. "_" .. program_version

local component = require("component")
local internet = component.internet

local function lazy_require(name)
    local this = {}
    local module
    setmetatable(this, {
        __index = function(_, key)
            module = module or require(name)
            return module[key]
        end,
        __call = function(_, ...)
            module = module or require(name)
            return module(...)
        end
    })
    return this
end

local x25519 = lazy_require("x25519")
local sha2 = lazy_require("sha2")

local function parse_ssh_version(str)
    local proto, software, comments = str:match("^SSH%-(%d+%.%d+)%-(%S+)%s?(.*)\r\n$")
    if not proto then return nil, "invalid SSH version string" end
    return proto, software, comments or nil
end

local function encode_ssh_version(proto, software, comments)
    return "SSH-" .. proto .. "-" .. software .. (comments and (" " .. comments) or "") .. "\r\n"
end

local random

if not component.isAvailable("data") or component.data.random == nil then
    terminal.write("ssh: WARNING: THIS PROGRAM IS NOT CRYPTOGRAPHICALLY SAFE WITHOUT A TIER 2 DATA CARD, DO NOT CONNECT TO A REAL SERVER WITH THIS\n")
    random = function(n)
        local parts = {}
        for i = 1, n do parts[i] = math.random(0, 255) end
        return table.concat(parts)
    end
else
    random = component.data.random
end

local SSH_MSG_DISCONNECT = 1
local SSH_MSG_IGNORE = 2
local SSH_MSG_UNIMPLEMENTED = 3
local SSH_MSG_DEBUG = 4
local SSH_MSG_SERVICE_REQUEST = 5
local SSH_MSG_SERVICE_ACCEPT = 6
local SSH_MSG_KEXINIT = 20
local SSH_MSG_NEWKEYS = 21

-- RFC 4253 7.1: first entry in client list that also appears in server_list wins
local function negotiate_one(client_list, server_list)
    local server_set = {}
    for _, name in ipairs(server_list) do server_set[name] = true end
    for _, name in ipairs(client_list) do
        if server_set[name] then return name end
    end
    return nil
end

local function negotiate(client_algs, server_kex)
    local chosen = {}
    for field, _ in pairs(client_algs) do
        chosen[field] = negotiate_one(client_algs[field] or {}, server_kex[field] or {})
        local is_language_field = field:find("^languages_")
        if not chosen[field] and not is_language_field then
            return nil, "no agreement on " .. field
        end
    end
    return chosen
end

local function xor_with_byte(s, b)
    local out = {}
    for i = 1, #s do
        out[i] = string.char(string.byte(s, i) ~ b)
    end
    return table.concat(out)
end

local function hmac(hash_fn, block_size, key, message)
    if #key > block_size then key = hash_fn(key) end
    if #key < block_size then key = key .. string.rep('\0', block_size - #key) end
    return hash_fn(xor_with_byte(key, 0x5c) .. hash_fn(xor_with_byte(key, 0x36) .. message))
end

local mac_algorithms_priority = {
    "hmac-sha2-256",
    "hmac-sha2-512",
    --"hmac-sha1"
}

local mac_algorithms = {
    ["hmac-sha2-256"] = function(key, message) return hmac(sha2.sha256, 64, key, message) end,
    ["hmac-sha2-512"] = function(key, message) return hmac(sha2.sha512, 128, key, message) end,
    --["hmac-sha1"] = function(key, message) return hmac(sha1, 64, key, message) end,
}

local encryption_algorithms_priority = {
    "chacha20-poly1305@openssh.com",
-- aes128-gcm@openssh.com
-- aes256-gcm@openssh.com
-- aes128-ctr
-- aes192-ctr
-- aes256-ctr
}
local encryption_algorithms = {
    ["chacha20-poly1305@openssh.com"] = nil
}

local server_host_key_algorithms_priority = {
    --[[
    "rsa-sha2-512",
    "rsa-sha2-256",
    "ecdsa-sha2-nistp256",
    "ssh-ed25519"
    ]]
}

local server_host_key_algorithms = {

}

local function socket()
    local this = {}
    local _sock
    local _host
    local _port

    function this.connect(host, port)
        _host = host; _port = port
        _sock = internet.connect(host, port)

        local ok = false
        for _ = 1, 50 do
            local err
            ok, err = _sock.finishConnect()
            if ok == nil then
                return nil, err
            end
            if ok == true then
                return true
            end
            coroutine.yield()
        end
        return nil, "timed out"
    end

    local _read_buf = ""

    function this.read(n)
        while #_read_buf < n do
            local chunk, err = _sock.read(n - #_read_buf)
            if chunk == nil then return nil, err end
            if chunk == "" then
                coroutine.yield()
            else
                _read_buf = _read_buf .. chunk
            end
        end
        local to_return = _read_buf:sub(1, n);
        _read_buf = _read_buf:sub(n + 1, -1)
        return to_return
    end

    function this.read_until_crlf(n)
        local chunk, err = _sock.read(n)
        if not chunk then return nil, err end
        _read_buf = _read_buf .. chunk

        local crlf_pos = _read_buf:find("\r\n", 1, true)
        if crlf_pos == nil then return nil end

        local to_return = _read_buf:sub(1, crlf_pos + 1)
        _read_buf = _read_buf:sub(crlf_pos + 2, -1)

        return to_return
    end

    function this.read_uint8()
        local x, err = this.read(1)
        if x == nil then return nil, err end
        return string.unpack(">I1", x)
    end

    function this.read_uint32()
        local x, err = this.read(4)
        if x == nil then return nil, err end
        return string.unpack(">I4", x)
    end

    function this.read_packet(mac_length)
        mac_length = mac_length or 0

        local packet_length, err = this.read_uint32()
        if not packet_length then return nil, err end

        local padding_length, err = this.read_uint8()
        if not padding_length then return nil, err end

        local payload, err = this.read(packet_length - padding_length - 1)
        if not payload then return nil, err end

        local random_padding, err = this.read(padding_length)
        if not random_padding then return nil, err end

        local mac, err = this.read(mac_length)
        if not mac then return nil, err end

        return payload
    end

    function this.write_packet(payload, block_size)
        block_size = block_size or 8
        local min_padding = 4

        local padding_length = block_size - ((5 + #payload) % block_size)
        if padding_length < min_padding then
            padding_length = padding_length + block_size
        end

        local padding = random(padding_length)

        local packet_length = 1 + #payload + padding_length
        local packet = string.pack(">I4", packet_length)
            .. string.char(padding_length)
            .. payload
            .. padding

        return this.write(packet)
    end

    local function parse_namelist(payload, pos)
        local len = string.unpack(">I4", payload, pos)
        local s = payload:sub(pos + 4, pos + 4 + len - 1)
        local list = {}
        if #s > 0 then
            for name in s:gmatch("[^,]+") do
                list[#list + 1] = name
            end
        end
        return list, pos + 4 + len
    end

    local function encode_namelist(list)
        return string.pack(">s4", table.concat(list or {}, ","))
    end

    function this.parse_kexinit(payload)
        local packet_type = payload:byte(1)
        if packet_type ~= SSH_MSG_KEXINIT then return nil, "unknown server side error" end

        local pos = 2
        local kex = {}
        kex.cookie = payload:sub(pos, pos + 15)
        pos = pos + 16

        kex.kex_algorithms, pos = parse_namelist(payload, pos)
        kex.server_host_key_algorithms, pos = parse_namelist(payload, pos)
        kex.encryption_algorithms_client_to_server, pos = parse_namelist(payload, pos)
        kex.encryption_algorithms_server_to_client, pos = parse_namelist(payload, pos)
        kex.mac_algorithms_client_to_server, pos = parse_namelist(payload, pos)
        kex.mac_algorithms_server_to_client, pos = parse_namelist(payload, pos)
        kex.compression_algorithms_client_to_server, pos = parse_namelist(payload, pos)
        kex.compression_algorithms_server_to_client, pos = parse_namelist(payload, pos)
        kex.languages_client_to_server, pos = parse_namelist(payload, pos)
        kex.languages_server_to_client, pos = parse_namelist(payload, pos)

        kex.first_kex_packet_follows = payload:byte(pos) ~= 0
        pos = pos + 1

        -- next 4 bytes are reserved (uint32 == 0)
        return kex
    end

    function this.build_kexinit(algs)
        return string.char(SSH_MSG_KEXINIT) ..
            random(16) ..
            encode_namelist(algs.kex_algorithms) ..
            encode_namelist(algs.server_host_key_algorithms) ..
            encode_namelist(algs.encryption_algorithms_client_to_server) ..
            encode_namelist(algs.encryption_algorithms_server_to_client) ..
            encode_namelist(algs.mac_algorithms_client_to_server) ..
            encode_namelist(algs.mac_algorithms_server_to_client) ..
            encode_namelist(algs.compression_algorithms_client_to_server) ..
            encode_namelist(algs.compression_algorithms_server_to_client) ..
            encode_namelist(algs.languages_client_to_server) ..
            encode_namelist(algs.languages_server_to_client) ..
            '\0' .. -- first_kex_packet_follows
            '\0\0\0\0' -- reserved
    end

    function this.write(x)
        local written = 0
        while written < #x do
            local bytes, err = _sock.write(x:sub(written + 1, -1))
            if not bytes then return nil, err end
            written = written + bytes
        end
        return written
    end

    return this
end

local sock = socket()
sock.connect("localhost", 22)
-- RFC 4253: The maximum length of the string is 255 characters
local server_version_raw = sock.read_until_crlf(255)
local server_protoversion, server_softwareversion, server_comments = parse_ssh_version(server_version_raw)
if server_protoversion ~= "2.0" then
    return
end

local client_version_line = encode_ssh_version("2.0", program_string)
sock.write(client_version_line)

local V_C = client_version_line:gsub("\r\n$", "")
local V_S = server_version_raw:gsub("\r\n$", "")

local I_S, kerr = sock.read_packet()
if not I_S then
    print("failed to read server KEXINIT: " .. tostring(kerr))
    return
end
local server_kex, perr = sock.parse_kexinit(I_S)
if not server_kex then
    print("failed to parse server KEXINIT: " .. tostring(perr))
    return
end

local my_algs = {
    kex_algorithms = { "curve25519-sha256" },
    server_host_key_algorithms = server_host_key_algorithms,
    encryption_algorithms_client_to_server = encryption_algorithms_priority,
    encryption_algorithms_server_to_client = encryption_algorithms_priority,
    mac_algorithms_client_to_server = mac_algorithms_priority,
    mac_algorithms_server_to_client = mac_algorithms_priority,
    compression_algorithms_client_to_server = { "none" },
    compression_algorithms_server_to_client = { "none" },
    languages_client_to_server = {},
    languages_server_to_client = {},
}

local I_C = sock.build_kexinit(my_algs)
sock.write_packet(I_C)

local chosen, nerr = negotiate(my_algs, server_kex)
if not chosen then
    print("algorithm negotiation failed: " .. tostring(nerr))
    return
end