local ajp = require "ajp"
local io = require "io"
local shortport = require "shortport"
local stdnse = require "stdnse"

description = [[
Requests a URI over the Apache JServ Protocol and displays the result
(or stores it in a file). Different AJP methods such as; GET, HEAD,
TRACE, PUT or DELETE may be used.

The Apache JServ Protocol is commonly used by web servers to communicate with
back-end Java application server containers.
]]

---
-- @usage
-- nmap -p 8009 <ip> --script ajp-request
--
-- @output
-- PORT     STATE SERVICE
-- 8009/tcp open  ajp13
-- | ajp-request: 
-- | <!DOCTYPE HTML>
-- | <html>
-- | <head>
-- | <title>JSP Test</title>
-- | 
-- | </head>
-- | <body>
-- | <h2>Hello, World.</h2>
-- | Fri May 04 02:09:40 UTC 2012
-- | </body>
-- |_</html>
--
-- @args method AJP method to be used when requesting the URI (default: GET)
-- @args path the path part of the URI to request
-- @args filename the name of the file where the results should be stored
-- @args username the username to use to access protected resources
-- @args password the password to use to access protected resources
--

author = "Patrik Karlsson"
license = "Same as Nmap--See http://nmap.org/book/man-legal.html"
categories = {"discovery", "safe"}


portrule = shortport.port_or_service(8009, 'ajp13', 'tcp')

local arg_method   = stdnse.get_script_args(SCRIPT_NAME .. ".method") or "GET"
local arg_path     = stdnse.get_script_args(SCRIPT_NAME .. ".path") or "/"
local arg_file     = stdnse.get_script_args(SCRIPT_NAME .. ".filename")
local arg_username = stdnse.get_script_args(SCRIPT_NAME .. ".username")
local arg_password = stdnse.get_script_args(SCRIPT_NAME .. ".password")

local function fail(err) return ("\n  ERROR: %s"):format(err or "") end

action = function(host, port)

	local helper = ajp.Helper:new(host, port)
	if ( not(helper:connect()) ) then
		return fail("Failed to connect to AJP server")
	end
	
	local valid_methods = {
		["GET"]    = true,
		["HEAD"]   = true, 
		["TRACE"]  = true,
		["PUT"]    = true,
		["DELETE"] = true,
		["OPTIONS"]= true,
	}
	
	local method = arg_method:upper()
	if ( not(valid_methods[method]) ) then
		return fail(("Method not supported: %s"):format(arg_method))
	end
	
	local options = { auth = { username = arg_username, password = arg_password } }
	local status, response = helper:request(arg_method, arg_path, nil, nil, options)
	if ( not(status) ) then
		return fail("Failed to retrieve response for request")
	end
	helper:close()

	if ( response ) then
		local output = response['status-line'] .. "\n" .. 
						stdnse.strjoin("\n", response.rawheaders) .. 
						(response.body and "\n\n" .. response.body or "")
		if ( arg_file ) then
			local f = io.open(arg_file, "w")
			if ( not(f) ) then
				return fail(("Failed to open file %s for writing"):format(arg_file))
			end
			f:write(output)
			f:close()
			return ("Response was written to file: %s"):format(arg_file)
		else
			return "\n" .. output
		end
	end
end

