local ipset =   require "ipset_message_pb"
local ipset_socket = "unix:/var/run/imunify360/libiplists-daemon.sock"

local function write_size(little_endian, format, ...)
    local res = ''
    local values = {...}
    for i=1,#format do
        local size = tonumber(format:sub(i,i))
        local value = values[i]
        local str = ""
        for j=1,size do
            str = str .. string.char(value % 256)
            value = math.floor(value / 256)
        end
        if not little_endian then
            str = string.reverse(str)
        end
        res = res .. str
    end
    return res
end


local function read_size(big_endian, format, ...)
    values = {...}
    local total = 0
    for i=1,#format do
        local size = tonumber(format:sub(i,i))
        local value = values[i]
        if not big_endian then
            value = value:reverse()
        end
        local res = 0
        for j=1,size do
            res = res * 256 + string.byte(value:sub(j,j))
            print(res)
        end
        total = total + res
    end
    return total
end


local function composeMessage(ip)
    local msg = ipset.Request()

    msg.ip = ngx.var.remote_addr
    msg.ua = "-"
    msg.url = "-"
    msg.uid = "-"
    msg.captcha = "-"
    msg.host = "-"
    msg.proxyname = ngx.var.remote_proxy
    msg.port = ngx.var.server_port
    msg.realip = ip
    if ngx.var.domain_whitelisted == "1" then
        msg.wldomain = true
    end

    local serialized = msg:SerializeToString()
    return write_size(false, "2", string.len(serialized)) .. serialized
end

local function decomposeMessage(pb_msg)
    local status_map = {
        [0]  = "NOT_FOUND",
        [10] = "WHITE",
        [11] = "REMOTE_PROXY",
        [12] = "WHITE_STATIC",
        [20] = "BLACK",
        [21] = "COUNTRY_BLACK",
        [22] = "BLACK_LOCAL",
        [30] = "GRAY",
        [40] = "SPLASH"
    }
    local msg = ipset.Response()
    msg:ParseFromString(pb_msg)

    if msg == nil then
        return {["ip"] = nil, ["verdict"] = "ERROR"}
    end

    local verdict = status_map[msg.action]
    if verdict == nil and msg.status == 1 then
        return {["ip"] = msg.ip, ["verdict"] = "ERROR"}
    end

    return {["ip"] = msg.ip, ["verdict"] = verdict}
end

local function query_peer(ip)

    local sock = ngx.socket.tcp()
    sock:settimeout(1000)
    local ok, err = sock:connect(ipset_socket)

    if not ok then
        ngx.log(ngx.ERR, "Could not connect to ipset socket: ", err)
        return
    end

    local sent, err = sock:send(composeMessage(ip))

    if not sent then
        ngx.log(ngx.ERR, "Could not send to ipset socket: ", err)
        return
    end

    local raw_size, err = sock:receive(2)
    if not raw_size then
        ngx.log(ngx.ERR, "Could not read size from ipset socket: ", err)
        return
    end

    local expected_size = read_size(true, "2", raw_size)

    local data, err = sock:receive(expected_size)
    if not data then
        ngx.log(ngx.ERR, "Could not read data from ipset socket: ", err)
        return
    end

    sock:close()

    return decomposeMessage(data)
end

local function check()

    local remote_addr = ngx.var.remote_addr

    if remote_addr == "127.0.0.1" or remote_addr == "::1" then
        ngx.var.webshield_ip_status = "NOT_FOUND"
        return
    end

    local remote_address = ngx.var.wsuserip or remote_addr

    local response = query_peer(remote_address)
    if response == nil then
        ngx.log(ngx.WARN, "Got empty response from IPSET daemon")
        ngx.var.webshield_ip_status = "NOT_FOUND"
        return
    end

    ngx.var.webshield_ip_status = response.verdict
end

if ngx.var.access_check_enabled == "0" then
  return check()
end