local nmap = require "nmap"
local packet = require "packet"
local ipOps = require "ipOps"
local bin = require "bin"
local stdnse = require "stdnse"
local table = require "table"
local math = require "math"
local string = require "string"

description = [[
Queries for the multicast path from a source to a destination host.

This works by sending an IGMP Traceroute Query and listening for IGMP
Traceroute responses. The Traceroute Query is sent to the first hop and
contains information about source, destination and multicast group addresses.
First hop defaults to the multicast All routers address. The default multicast
group address is 0.0.0.0 and the default destination is our own host address. A
source address must be provided. The responses are parsed to get interesting
information about interface addresses, used protocols and error codes.

This is similar to the mtrace utility provided in Cisco IOS.
]]

---
--@args mtrace.fromip Source address from which to traceroute.
--
--@args mtrace.toip Destination address to which to traceroute.
-- Defaults to our host address.
--
--@args mtrace.group Multicast group address for the traceroute.
-- Defaults to <code>0.0.0.0</code> which represents all group addresses.
--
--@args mtrace.firsthop Host to which the query is sent. If not set, the
-- query will be sent to <code>224.0.0.2</code>.
--
--@args mtrace.timeout Time to wait for responses.
-- Defaults to <code>7s</code>.
--
--@usage
-- nmap --script mtrace --script-args 'mtrace.fromip=172.16.45.4'
--
--@output
-- Pre-scan script results:
-- | mtrace: 
-- |   Group 0.0.0.0 from 172.16.45.4 to 172.16.0.1
-- |   Source: 172.16.45.4
-- |     In address: 172.16.34.3
-- |       Out address: 172.16.0.3
-- |       Protocol: PIM
-- |     In address: 172.16.45.4
-- |       Out address: 172.16.34.4
-- |       Protocol: PIM
-- |   Source: 172.16.45.4
-- |     In address: 172.16.13.1
-- |       Out address: 172.16.0.2
-- |       Protocol: PIM / Static
-- |     In address: 172.16.34.3
-- |       Out address: 172.16.13.3
-- |       Protocol: PIM
-- |     In address: 172.16.45.4
-- |       Out address: 172.16.34.4
-- |_      Protocol: PIM

author = "Hani Benhabiles"

license = "Same as Nmap--See http://nmap.org/book/man-legal.html"

categories = {"discovery", "safe", "broadcast"}

-- From: https://tools.ietf.org/id/draft-ietf-idmr-traceroute-ipm-07.txt
PROTO = {
    [0x01] = "DVMRP", 
    [0x02] = "MOSPF", 
    [0x03] = "PIM", 
    [0x04] = "CBT", 
    [0x05] = "PIM / Special table", 
    [0x06] = "PIM / Static", 
    [0x07] = "DVMRP / Static", 
    [0x08] = "PIM / MBGP", 
    [0x09] = "CBT / Special table",
    [0x10] = "CBT / Static", 
    [0x11] = "PIM / state created by Assert processing", 
}

FWD_CODE = {
    [0x00] = "NO_ERROR",
    [0x01] = "WRONG_IF",
    [0x02] = "PRUNE_SENT",
    [0x03] = "PRUNE_RCVD",
    [0x04] = "SCOPED",
    [0x05] = "NO_ROUTE",
    [0x06] = "WRONG_LAST_HOP",
    [0x07] = "NOT_FORWARDING",
    [0x08] = "REACHED_RP",
    [0x09] = "RPF_IF",
    [0x0A] = "NO_MULTICAST",
    [0x0B] = "INFO_HIDDEN",
    [0x81] = "NO_SPACE",
    [0x82] = "OLD_ROUTER",
    [0x83] = "ADMIN_PROHIB",
}

prerule = function()
    if nmap.address_family() ~= 'inet' then
	stdnse.print_verbose("%s is IPv4 only.", SCRIPT_NAME)
	return false
    end
    if not nmap.is_privileged() then
	stdnse.print_verbose("%s not running for lack of privileges.", SCRIPT_NAME)
	return false
    end
    return true
end

--- Generates a raw IGMP Traceroute Query.
--@param fromip Source address.
--@param toip Destination address.
--@param group Multicast group address.
--@param receiver Receiver of the response.
--@return data Raw Traceroute Query.
local traceRaw = function(fromip, toip, group, receiver)
    local data = bin.pack(">C", 0x1f) -- Type: Traceroute Query
    local data = data .. bin.pack(">C", 0x20) -- Hops: 32
    local data = data .. bin.pack(">S", 0x0000) -- Checksum: To be set later
    local data = data .. bin.pack(">I", ipOps.todword(group)) -- Multicast group
    local data = data .. bin.pack(">I", ipOps.todword(fromip)) -- Source 
    local data = data .. bin.pack(">I", ipOps.todword(toip)) -- Destination
    local data = data .. bin.pack(">I", ipOps.todword(receiver)) -- Receiver
    local data = data .. bin.pack(">C", 0x40) -- TTL
    local data = data .. bin.pack(">CS", 0x00, math.random(123456)) -- Query ID

    -- We calculate checksum
    data = data:sub(1,2) .. bin.pack(">S", packet.in_cksum(data)) .. data:sub(5)
    return data
end

--- Sends a raw IGMP Traceroute Query.
--@param interface Network interface to send through.
--@param destination Target host to which the packet is sent.
--@param trace_raw Traceroute raw Query.
local traceSend = function(interface, destination, trace_raw)
    local ip_raw = bin.pack("H", "45c00040ed780000400218bc0a00c8750a00c86b") .. trace_raw
    local trace_packet = packet.Packet:new(ip_raw, ip_raw:len())
    trace_packet:ip_set_bin_src(ipOps.ip_to_str(interface.address))
    trace_packet:ip_set_bin_dst(ipOps.ip_to_str(destination))
    trace_packet:ip_set_len(#trace_packet.buf)
    trace_packet:ip_count_checksum()

    if destination == "224.0.0.2" then
	-- Doesn't affect results as it is ignored but most routers, but RFC
	-- 3171 should be respected.
	trace_packet:ip_set_ttl(1)
    end
    trace_packet:ip_count_checksum()

    local sock = nmap.new_dnet()
    if destination == "224.0.0.2" then
	sock:ethernet_open(interface.device)
	-- Ethernet IPv4 multicast, our ethernet address and packet type IP
	local eth_hdr = bin.pack("HAH", "01 00 5e 00 00 02", interface.mac, "08 00")
	sock:ethernet_send(eth_hdr .. trace_packet.buf)
	sock:ethernet_close()
    else
	sock:ip_open()
	sock:ip_send(trace_packet.buf, destination)
	sock:ip_close()
    end
end

--- Parses an IGMP Traceroute Response and returns it in structured form.
--@param data Raw Traceroute Response.
--@return response Structured Traceroute Response.
local traceParse = function(data)
    local index
    local response = {}

    -- first byte should be IGMP type == 0x1e (Traceroute Response)
    if data:byte(1) ~= 0x1e then return end

    -- Hops
    index, response.hops = bin.unpack(">C", data, 2)

    -- Checksum 
    index, response.checksum = bin.unpack(">S", data, index)

    -- Group
    index, response.group = bin.unpack("<I", data, index)
    response.group = ipOps.fromdword(response.group)

    -- Source address
    index, response.source = bin.unpack("<I", data, index)
    response.source = ipOps.fromdword(response.source)

    -- Destination address
    index, response.destination = bin.unpack("<I", data, index)
    response.receiver = ipOps.fromdword(response.destination)

    -- Response address
    index, response.response = bin.unpack("<I", data, index)
    response.response = ipOps.fromdword(response.response)

    -- Response TTL
    index, response.ttl = bin.unpack(">C", data, index)

    -- Query ID
    index, response.qid = bin.unpack(">C", data, index)
    index, response.qid = response.qid * 2^16 + bin.unpack(">S", data, index)

    local block
    response.blocks = {}
    -- Now, parse data blocks
    while true do
	-- To end parsing and not get stuck in infinite loops.
	if index >= #data then
	    break
	elseif #data - index < 31 then
	    stdnse.print_verbose("%s malformated traceroute response.", SCRIPT_NAME)
	    return
	end

	block = {}
	-- Query Arrival
	index, block.query = bin.unpack(">I", data, index)

	-- In itf address
	index, block.inaddr = bin.unpack("<I", data, index)
	block.inaddr = ipOps.fromdword(block.inaddr)

	-- Out itf address
	index, block.outaddr = bin.unpack("<I", data, index)
	block.outaddr = ipOps.fromdword(block.outaddr)

	-- Previous rtr address
	index, block.prevaddr = bin.unpack("<I", data, index)
	block.prevaddr = ipOps.fromdword(block.prevaddr)

	-- In packets
	index, block.inpkts = bin.unpack(">I", data, index)

	-- Out packets
	index, block.outpkts = bin.unpack(">I", data, index)

	-- S,G pkt count
	index, block.sgpkt = bin.unpack(">I", data, index)

	-- Protocol
	index, block.proto = bin.unpack(">C", data, index)

	-- Forward TTL
	index, block.fwdttl = bin.unpack(">C", data, index)

	-- Options
	index, block.options = bin.unpack(">C", data, index)

	-- Forwarding Code
	index, block.code = bin.unpack(">C", data, index)

	table.insert(response.blocks, block)
    end
    return response
end

-- Listens for IGMP Traceroute responses
--@param interface Network interface to listen on.
--@param timeout Amount of time to listen for in seconds.
--@param responses table to insert responses into.
local traceListener = function(interface, timeout, responses)
    local condvar = nmap.condvar(responses)
    local start = nmap.clock_ms()
    local listener = nmap.new_socket()
    local p, trace_raw, status, l3data, response, _

    -- IGMP packets that are sent to our host
    local filter = 'ip proto 2 and dst host ' .. interface.address
    listener:set_timeout(100)
    listener:pcap_open(interface.device, 1024, true, filter)

    while (nmap.clock_ms() - start) < timeout do
	status, _, _, l3data = listener:pcap_receive()
	if status then
	    p = packet.Packet:new(l3data, #l3data)
	    trace_raw = string.sub(l3data, p.ip_hl*4 + 1)
	    if p then
		-- Check that IGMP Type == 0x1e (Traceroute Response)
		if trace_raw:byte(1) == 0x1e then
		    response = traceParse(trace_raw)
		    if response then 
			response.srcip = p.ip_src
			table.insert(responses, response)
		    end
		end
	    end
	end
    end
    condvar("signal")
end

-- Returns the network interface used to send packets to a target host.
--@param target host to which the interface is used.
--@return interface Network interface used for target host.
local getInterface = function(target)
    -- First, create dummy UDP connection to get interface
    local sock = nmap.new_socket()
    local status, err = sock:connect(target, "12345", "udp")
    if not status then
	stdnse.print_verbose("%s: %s", SCRIPT_NAME, err)
	return
    end
    local status, address, _, _, _ = sock:get_info()
    if not status then
	stdnse.print_verbose("%s: %s", SCRIPT_NAME, err)
	return
    end
    for _, interface in pairs(nmap.list_interfaces()) do
	if interface.address == address then
	    return interface
	end
    end
end


action = function()
    local fromip = stdnse.get_script_args(SCRIPT_NAME .. ".fromip")
    local toip = stdnse.get_script_args(SCRIPT_NAME .. ".toip")
    local group = stdnse.get_script_args(SCRIPT_NAME .. ".group") or "0.0.0.0"
    local firsthop = stdnse.get_script_args(SCRIPT_NAME .. ".firsthop") or "224.0.0.2"
    local timeout = stdnse.parse_timespec(stdnse.get_script_args(SCRIPT_NAME .. ".timeout"))
    local responses = {}
    timeout = (timeout or 7) * 1000

    -- Source address from which to traceroute
    if not fromip then
	stdnse.print_verbose("%s: A source IP must be provided through fromip argument.", SCRIPT_NAME)
	return
    end

    -- Get network interface to use
    local interface = nmap.get_interface() 
    if interface then
	interface = nmap.get_interface_info(interface) 
    else
	interface = getInterface(firsthop)
    end
    if not interface then
	return ("\n ERROR: Couldn't get interface for %s"):format(firsthop)
    end

    -- Destination defaults to our own host
    toip = toip or interface.address

    stdnse.print_debug("%s: Traceroute group %s from %s to %s.", SCRIPT_NAME, group, fromip, toip)
    stdnse.print_debug("%s: will send to %s via %s interface.", SCRIPT_NAME, firsthop, interface.shortname)

    -- Thread that listens for responses
    stdnse.new_thread(traceListener, interface, timeout, responses)

    -- Send request after small wait to let Listener start
    stdnse.sleep(0.1)
    local trace_raw = traceRaw(fromip, toip, group, interface.address)
    traceSend(interface, firsthop, trace_raw)

    local condvar = nmap.condvar(responses)
    condvar("wait")
    if #responses > 0 then
	local outresp 
	local output, outblock = {}
	table.insert(output, ("Group %s from %s to %s"):format(group, fromip, toip))
	for _, response in pairs(responses) do
	    outresp = {}
	    outresp.name = "Source: " .. response.srcip
	    for _, block in pairs(response.blocks) do
		outblock = {}
		outblock.name = "In address: " .. block.inaddr
		table.insert(outblock, "Out address: " .. block.outaddr)
		-- Protocol
		if PROTO[block.proto] then
		    table.insert(outblock, "Protocol: " .. PROTO[block.proto])
		else
		    table.insert(outblock, "Protocol: Unknown")
		end
		-- Error Code, we ignore NO_ERROR which is the normal case.
		if FWD_CODE[block.code] and block.code ~= 0x00 then
		    table.insert(outblock, "Error code: " .. FWD_CODE[block.code])
		elseif block.code ~= 0x00 then
		    table.insert(outblock, "Error code: Unknown")
		end
		table.insert(outresp, outblock)
	    end
	    table.insert(output, outresp)
	end
	return stdnse.format_output(true, output)
    end
end