#!/usr/bin/lua
package.path = package.path..';'..(arg[0]:gsub('/[^/]+$', ''))..'/?.lua'

-- Copyright Julien Thomas < julthomas @ free.fr >
-- Copyright Zenetys < jthomas @ zenetys.com >
-- Author: Julien Thomas
-- Licence: MIT

lc = require 'libcheck'
lp = require 'libperfdata'
curl = require 'lcurl'

lc.checkname = 'REST'
lc.shortdescr = 'Check a value returned by an HTTP JSON REST API'
lc.progtype = 'rest'

lc.optsdef = {
    { short = 'U', long = 'url', required = true, help = 'Request URL' },
    { short = 'P', long = 'protocol', help = 'Request protocol (http, https), -H required, -U must start with /' },
    { short = 'H', long = 'hostname', help = 'Request host or IP, -P required, -U must start with /' },
    { short = 'r', long = 'header', call = lc.setter_opt_array, help = 'Request header(s)' },
    { short = 'n', long = 'username', help = 'Authentication username' },
    { short = 'p', long = 'password', help = 'Authentication password' },
    { short = 'd', long = 'post-data', help = 'Request data to POST, otherwise GET' },
    { short = 'X', long = 'method', help = 'Force request method, eg: GET when using --post-data' },
    { short = 'L', long = 'follow', call = lc.setter_opt_boolean, help = 'Follow location' },
    { short = 's', long = 'parameter', call = lc.setter_opt_kv, help = 'Replace ${key} by value in --post-data, format: key=value' },
    { short = 'l', long = 'label', call = lc.setter_opt_array, help = 'Metric(s) name' },
    { short = 'a', long = 'addon', help = 'Custom addon script to handle JSON data' },
    { short = 'K', long = 'addon-insecure', call = lc.setter_opt_boolean, help = 'Do not restrict lua env in addon' },
    { short = 'w', long = 'warning', call = lc.setter_opt_array, help = 'Metric(s) warning' },
    { short = 'c', long = 'critical', call = lc.setter_opt_array, help = 'Metric(s) critical' },
    { short = 'u', long = 'uom', call = lc.setter_opt_array, help = 'Metric(s) uom' },
    { short = 'm', long = 'min', call = lc.setter_opt_array, help = 'Metric(s) min' },
    { short = 'M', long = 'max', call = lc.setter_opt_array, help = 'Metric(s) max' },
    { short = 'x', long = 'null-value', call = lc.setter_opt_array, help = 'Metric(s) null value' },
    { short = 't', long = 'timeout', call = lc.setter_opt_number, help = 'cURL timeout in seconds' },
    { short = 'R', long = 'netrc', call = lc.setter_opt_iboolean, help = 'Enable ~/.netrc' },
    { short = 'A', long = 'raw', call = lc.setter_opt_boolean, help = 'Raw data, do not decode JSON, require --addon' },
    -- 0: CURL_NETRC_IGNORED, 1: CURL_NETRC_OPTIONAL
    { short = 'j', long = 'cookies', call = lc.setter_opt_boolean, help = 'Store and reuse cookies' },
    { short = 'N', long = 'check-name', call = function (o,v) lc.checkname = v; return v end,
      help = 'Set output prefix' },
}

lc.init_opts()

-- extra options checks to set request url
has_protocol = (lc.opts.protocol and #lc.opts.protocol > 0)
has_hostname = (lc.opts.hostname and #lc.opts.hostname > 0)
if (has_protocol and not has_hostname) or
   (has_hostname and not has_protocol) then
    lc.die(lc.UNKNOWN, 'Options -P/-H are mutually required')
end
if has_protocol then -- and has_hostname
    if lc.opts.url:sub(1, 1) ~= '/' then
        lc.die(lc.UNKNOWN, 'URL must start with / (slash) when using -P/-H')
    end
    url = lc.opts.protocol..'://'..lc.opts.hostname..lc.opts.url
else
    url = lc.opts.url
end

-- defaults
if not lc.opts.parameter then lc.opts.parameter = {} end
if not lc.opts.label then lc.opts.label = {} end
if not lc.opts.warning then lc.opts.warning = {} end
if not lc.opts.critical then lc.opts.critical = {} end
if not lc.opts.uom then lc.opts.uom = {} end
if not lc.opts.min then lc.opts.min = {} end
if not lc.opts.max then lc.opts.max = {} end
if not lc.opts.null_value then lc.opts.null_value = {} end
if not lc.opts.timeout then lc.opts.timeout = 10 end
if not lc.opts.netrc then lc.opts.netrc = 0 end
if not lc.opts.raw then lc.opts.raw = false end

-- query
if lc.opts.post_data ~= nil then
    if lc.opts.post_data:sub(1, 1) == '@' then
        fd, err = io.open(lc.opts.post_data:sub(2), 'rb')
        if not fd then lc.die(lc.UNKNOWN, 'Cannot open data file: '..err) end
        lc.opts.post_data, err = fd:read('*a')
        fd:close()
        if not lc.opts.post_data then lc.die(lc.UNKNOWN, 'Cannot read data file: '..err) end
    end
    for k,v in pairs(lc.opts.parameter) do
        lc.opts.post_data = lc.opts.post_data:gsub('%${'..k:gsub('%.', '%%.')..'}', v)
    end
end

curlopts = {
    url = url,
    ssl_verifypeer = 0,
    ssl_verifyhost = 0,
    timeout = lc.opts.timeout,
    writefunction = function (x)
        data = (data or '')..x
        return x and x:len() or 0
    end,
    verbose = lc.opts.debug,
    username = lc.opts.username,
    password = lc.opts.password,
    followlocation = lc.opts.follow,
    postfields = lc.opts.post_data,
    [curl.OPT_CUSTOMREQUEST] = lc.opts.method,
    [curl.OPT_FAILONERROR] = true,
    [curl.OPT_NETRC] = lc.opts.netrc,
}
-- some versions of lua-curl complain if httpheader is nil or an empty array
if lc.opts.header and #lc.opts.header > 0 then
    curlopts['httpheader'] = lc.opts.header
end
if lc.opts.cookies then
    lc.init_cache()
    curlopts['cookiefile'] = lc.cachedir..'/cookies'
    curlopts['cookiejar'] = lc.cachedir..'/cookies'
end
lc.dump(curlopts, 'Dump curlopts')
c = curl.easy(curlopts)
success, err = pcall(c.perform, c)
code = c:getinfo(curl.INFO_RESPONSE_CODE)
time = c:getinfo(curl.INFO_TOTAL_TIME)
c:close()

if not success then
    msg = 'cURL failed: '
    if code ~= nil and code > 0 then
        -- on some version of lua / lua-curl code will be print as float,
        -- cleanup with string.format()
        msg = msg..string.format('HTTP status %d, ', code)
    end
    lc.die(lc.UNKNOWN, msg..err:msg())
end

-- json
lc.dump(data, 'Dump data')
if not lc.opts.raw then
    data, err = lc.cjson.decode(data)
    if not data then lc.die(lc.UNKNOWN, 'JSON decode failed: '..err) end
end

-- perfdata
perfdata = {}

-- limited env for addons/eval code
env_limited = {
    -- check related
    data = data, lc = lc, perfdata = perfdata,
    -- lua functions
    ipairs = ipairs, math = math, pairs = pairs, print = print,
    string = string, table = table, type = type, tonumber = tonumber,
    tostring = tostring,
}

for i = 1, #lc.opts.label do
    -- path[:label], path is split[3], label is split[4]
    split = { lc.opts.label[i]:find('^([^:]+):?(.*)') }
    if split[1] ~= nil then
        p = {
            name = split[4]:len() > 0 and split[4] or split[3],
            warning = lc.opts.warning[i],
            critical = lc.opts.critical[i],
            uom = lc.opts.uom[i],
            min = lc.opts.min[i],
            max = lc.opts.max[i],
        }
        fn, err = load('return data.'..split[3], nil, nil, env_limited)
        if fn == nil then
            lc.perr('Failed to load eval code: '..err)
        else
            success, value = pcall(fn)
            if success then
                p.value = tonumber(value)
            else
                lc.perr('Failed to run eval code: '..value)
                if lc.opts.null_value[i] then
                    lc.perr('Using provided null-value: '..lc.opts.null_value[i])
                    p.value = tonumber(lc.opts.null_value[i])
                else
                    p.value = nil
                end
            end
        end
        table.insert(perfdata, p)
    end
end
-- custom processing
if lc.opts.addon ~= nil then
    local addon_file = lc.opts.addon
    if lc.lfs.attributes(addon_file, 'mode') ~= 'file' then
        if addon_file:sub(1,1) ~= '/' and
           lc.lfs.attributes(lc.progdir..'/'..addon_file, 'mode') == 'file' then
            addon_file = lc.progdir..'/'..addon_file
        else
            lc.die_unkn('Addon not found: '..addon_file)
        end
    end
    fn, err = loadfile(addon_file, nil, lc.opts.addon_insecure and _ENV or env_limited)
    if fn == nil then lc.die(lc.UNKNOWN, 'Failed to load addon script: '..err) end
    success, err = pcall(fn)
    if not success then lc.die(lc.UNKNOWN, 'Failed to run addon script: '..err) end
end
-- if no label or perfdata was given, add curl time
if #perfdata == 0 then
    table.insert(perfdata, {
        name = 'time',
        value = time,
        uom = 's',
    })
end
lc.dump(perfdata, 'Dump perfdata')

-- set exit state and output if not already done by addon
if not lc.exit_code then
    lc.exit_code = lp.compute_perfdata(perfdata)
end
if not lc.exit_message then
    lc.exit_message = lp.format_output(perfdata)..'|'..
        lp.format_perfdata(perfdata)
end
