-- Copyright © by Jeff Foley 2017-2023. All rights reserved.
-- Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
-- SPDX-License-Identifier: Apache-2.0

local json = require("json")

name = "PentestTools"
type = "api"

function start()
    set_rate_limit(1)
end

function check()
    local c
    local cfg = datasrc_config()
    if (cfg ~= nil) then
        c = cfg.credentials
    end

    if (c ~= nil and c.key ~= nil and c.key ~= "") then
        return true
    end
    return false
end

function vertical(ctx, domain)
    local c
    local cfg = datasrc_config()
    if (cfg ~= nil) then
        c = cfg.credentials
    end

    if (c == nil or c.key == nil or c.key == "") then
        return
    end

    local id = start_scan(domain, c.key)
    if (id == "") then
        return
    end

    while(true) do
        local status = get_scan_status(id, c.key)
        if (status == "failed") then
            return
        elseif (status == "finished") then
            break
        end

        for _=1,5 do check_rate_limit() end
    end

    local output = get_output(id, c.key)
    if (output ~= "") then
        for _, r in pairs(output) do
            new_name(ctx, r[1])
            new_addr(ctx, r[2], r[1])
        end
    end
end

function start_scan(domain, key)
    local body, err = json.encode({
        ['op']="start_scan",
        ['tool_id']=20,
        ['target']=domain,
        ['tool_params'] = {
            ['web_details']="off",
            ['do_bing_search']="off",
        },
    })
    if (err ~= nil and err ~= "") then
        return ""
    end

    local resp, err = request(ctx, {
        ['url']=build_url(key),
        ['method']="POST",
        ['header']={['Content-Type']="application/json"},
        ['body']=body,
    })
    if (err ~= nil and err ~= "") then
        log(ctx, "start_scan request to service failed: " .. err)
        return ""
    elseif (resp.status_code < 200 or resp.status_code >= 400) then
        log(ctx, "start_scan request to service returned with status: " .. resp.status)
        return ""
    end

    d = json.decode(resp.body)
    if (d == nil) then
        log(ctx, "failed to decode the JSON start_scan response")
        return ""
    elseif (d.op_status == nil or d.op_status ~= "success") then
        return ""
    end
    return d.scan_id
end

function get_scan_status(id, key)
    local body, err = json.encode({
        ['op']="get_scan_status",
        ['scan_id']=id,
    })
    if (err ~= nil and err ~= "") then
        return "failed"
    end

    local resp, err = request(ctx, {
        ['url']=build_url(key),
        ['method']="POST",
        ['header']={['Content-Type']="application/json"},
        ['body']=body,
    })
    if (err ~= nil and err ~= "") then
        log(ctx, "get_scan_status request to service failed: " .. err)
        return "failed"
    elseif (resp.status_code < 200 or resp.status_code >= 400) then
        log(ctx, "get_scan_status request to service returned with status: " .. resp.status)
        return "failed"
    end

    d = json.decode(resp.body)
    if (d == nil) then
        log(ctx, "failed to decode the JSON get_scan_status response")
        return "failed"
    elseif (d.op_status == nil or d.op_status ~= "success") then
        return "failed"
    elseif (d.scan_status ~= nil and (d.scan_status == "waiting" or d.scan_status == "running")) then
        return "progress"
    else
        return "finished"
    end
end

function get_output(id, key)
    local body, err = json.encode({
        ['op']="get_output",
        ['scan_id']=id,
        ['output_format']="json",
    })
    if (err ~= nil and err ~= "") then
        return ""
    end

    local resp, err = request(ctx, {
        ['url']=build_url(key),
        ['method']="POST",
        ['header']={['Content-Type']="application/json"},
        ['body']=body,
    })
    if (err ~= nil and err ~= "") then
        log(ctx, "get_output request to service failed: " .. err)
        return ""
    elseif (resp.status_code < 200 or resp.status_code >= 400) then
        log(ctx, "get_output request to service returned with status: " .. resp.status)
        return ""
    end

    d = json.decode(resp.body)
    if (d == nil) then
        log(ctx, "failed to decode the JSON get_output response")
        return ""
    elseif (d.op_status == nil or d.op_status ~= "success" or 
        d.output_json == nil or #(d['output_json'].output_data) == 0) then
        return ""
    end
    return d['output_json'][1].output_data
end

function build_url(key)
    return "https://pentest-tools.com/api?key=" .. key
end