local nmap = require "nmap"
local stdnse = require "stdnse"
local table = require "table"
local bin = require "bin"
local packet = require "packet"
local ipOps = require "ipOps"
local target = require "target"
local coroutine = require "coroutine"
local string = require "string"

description = [[
Discovers targets that have IGMP Multicast memberships and grabs interesting information.

The scripts works by sending IGMP Membership Query message to the 224.0.0.1 All
Hosts multicast address and listening for IGMP Membership Report messages. The
script then extracts all the interesting information from the report messages
such as the version, group, mode, source addresses (depending on the version).  

The script defaults to sending an IGMPv2 Query but this could be changed to
another version (version 1 or 3) or to sending queries of all three version. If
no interface was specified as a script argument or with the -e option, the
script will proceed to sending queries through all the valid ethernet
interfaces.
]]

---
-- @args broadcast-igmp-discovery.timeout Time to wait for reports in seconds.
-- Defaults to <code>5</code> seconds.
--
-- @args broadcast-igmp-discovery.version IGMP version to use. Could be
-- <code>1</code>, <code>2</code>, <code>3</code> or <code>all</code>. Defaults to <code>2</code>
--
-- @args broadcast-igmp-discovery.interface Network interface to use.
--
--@usage
-- nmap --script broadcast-igmp-discovery
-- nmap --script broadcast-igmp-discovery -e wlan0
-- nmap --script broadcast-igmp-discovery
-- --script-args 'broadcast-igmp-discovery.version=all, broadcast-igmp-discovery.timeout=3'
--
--@output
--Pre-scan script results:
-- | broadcast-igmp-discovery: 
-- |   192.168.2.2
-- |     Interface: tap0
-- |     Version: 3
-- |     Group: 239.1.1.1
-- |       Mode: EXCLUDE
-- |     Group: 239.1.1.2
-- |       Mode: EXCLUDE
-- |     Group: 239.1.1.44
-- |       Mode: INCLUDE
-- |       Sources:
-- |           192.168.31.1
-- |   192.168.1.3
-- |     Interface: wlan0
-- |     Version: 2
-- |     Group: 239.255.255.250
-- |   192.168.1.3
-- |     Interface: wlan0
-- |     Version: 2
-- |     Group: 239.255.255.253
-- |_  Use the newtargets script-arg to add the results as targets
--


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 due to lack of privileges.", SCRIPT_NAME)
	return false
    end
    return true
end

author = "Hani Benhabiles"

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

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

--- Parses a raw igmp packet and return a structred packet.
-- @param data string IGMP Raw packet.
-- @return response table Structured igmp packet.
local igmpParse = function(data)
    local index
    local response = {}
    local group, source
    -- Report type (0x12 == v1, 0x16 == v2, 0x22 == v3)
    index, response.type = bin.unpack(">C", data, index)
    if response.type == 0x12 or response.type == 0x16 then
	-- Max response time
	index, response.maxrt = bin.unpack(">C", data, index)
	-- Checksum
	index, response.checksum = bin.unpack(">S", data, index)
	-- Multicast group
	index, response.group = bin.unpack("<I", data, index)
	response.group = ipOps.fromdword(response.group)
	return response
    elseif response.type == 0x22 and #data >= 12 then
	-- Skip reserved byte
	index = index + 1
	-- Checksum
	index, response.checksum = bin.unpack(">S", data, index)
	-- Skip reserved byte
	index = index + 2
	-- Number of groups
	index, response.ngroups = bin.unpack(">S", data, index)
	response.groups = {}
	for i=1,response.ngroups do
	    group = {}
	    -- Mode is either INCLUDE or EXCLUDE
	    index, group.mode = bin.unpack(">C", data, index)
	    -- Auxiliary data length in the group record (in 32bits units)
	    index, group.auxdlen = bin.unpack(">C", data, index)
	    -- Number of source addresses
	    index, group.nsrc = bin.unpack(">S", data, index)
	    index, group.address = bin.unpack("<I", data, index)
	    group.address = ipOps.fromdword(group.address)
	    group.src = {}
	    if group.nsrc > 0 then
		for i=1,group.nsrc do
		    index, source = bin.unpack("<I", data, index)
		    table.insert(group.src, ipOps.fromdword(source))
		end
	    end
	    -- Skip auxiliary data
	    index = index + group.auxdlen
	    -- Insert group
	    table.insert(response.groups, group)
	end
	return response
    end
end

--- Listens for IGMP Membership reports packets.
-- @param interface Interface to listen on.
-- @param timeout Amount of time to listen for.
-- @param responses table to put valid responses into.
local igmpListener = function(interface, timeout, responses)
    local condvar = nmap.condvar(responses)
    local start = nmap.clock_ms()
    local listener = nmap.new_socket()
    local p, igmp_raw, status, l3data, response, _
    local devices = {}
    listener:set_timeout(100)
    listener:pcap_open(interface.device, 1024, true, 'ip proto 2')

    while (nmap.clock_ms() - start) < timeout do
	status, _, _, l3data = listener:pcap_receive()
	if status then
	    p = packet.Packet:new(l3data, #l3data)
	    igmp_raw = string.sub(l3data, p.ip_hl*4 + 1)
	    if p then
		-- check the first byte before sending to the parser
		-- response 0x12 == Membership Response version 1
		-- response 0x16 == Membership Response version 2
		-- response 0x22 == Membership Response version 3
		local igmptype = igmp_raw:byte(1)
		if igmptype == 0x12 or igmptype == 0x16 or igmptype == 0x22 then
		    response = igmpParse(igmp_raw)
		    if response then
			response.src = p.ip_src
			response.interface = interface.shortname
			-- Many hosts return more than one same response message
			-- this is to not output duplicates
			if not devices[response.src..response.type..(response.group or response.ngroups)] then
			    devices[response.src..response.type..(response.group or response.ngroups)] = true
			    table.insert(responses, response)
			end
		    end
		end
	    end
	end
    end
    condvar("signal")
end

--- Crafts a raw IGMP packet.
-- @param interface Source interface of the packet.
-- @param vesion IGMP version. Could be 1, 2 or 3.
-- @return string Raw IGMP packet.
local igmpRaw = function(interface, version)
    -- Only 1, 2 and 3 are valid IGMP versions
    if version ~= 1 and version ~= 2 and version ~= 3 then
	stdnse.print_debug("IGMP version %s doesn't exist.", version)
	return
    end

    -- Let's craft an IGMP Membership Query
    local igmp_raw = bin.pack(">C", 0x11) -- Membership Query, same for all versions
    if version == 1 then
	igmp_raw = igmp_raw .. bin.pack(">C", 0x00) -- Unused, 0x00 for version 1 only
    else
	igmp_raw = igmp_raw .. bin.pack(">C", 0x16) -- Max response time: 10 Seconds, for version 2 and 3
    end

    igmp_raw = igmp_raw .. bin.pack(">S", 0x00) -- Checksum, calculated later
    igmp_raw = igmp_raw .. bin.pack(">I", 0x00) -- Multicast Address: 0.0.0.0

    if version == 3 then
	-- Reserved = 4 bits (Should be zeroed)
	-- Supress Flag = 1 bit 
	-- QRV (Querier's Robustness Variable) = 3 bits
	-- all are set to 0
	igmp_raw = igmp_raw .. bin.pack(">C", 0x00)
	-- QQIC (Querier's Query Interval Code) in seconds = Set to 0 to get insta replies.
	igmp_raw = igmp_raw .. bin.pack(">C", 0x10)
	-- Number of sources (in the next arrays) = 1 ( Our IP only)
	igmp_raw = igmp_raw .. bin.pack(">S", 0x01)
	-- Source = Our IP address
	igmp_raw = igmp_raw .. bin.pack(">I", ipOps.todword(interface.address))
    end

    igmp_raw = igmp_raw:sub(1,2) .. bin.pack(">S", packet.in_cksum(igmp_raw)) .. igmp_raw:sub(5)

    return igmp_raw
end


local igmpQuery;
--- Sends an IGMP Membership query.
-- @param interface Network interface to send on.
-- @param vesion IGMP version. Could be 1, 2, 3 or all.
igmpQuery = function(interface, version)
    local srcip = interface.address
    local dstip = "224.0.0.1"

    if version == 'all' then
	-- Small pause to let listener begin and not miss reports.
	stdnse.sleep(0.5)
	igmpQuery(interface, 3)
	igmpQuery(interface, 2)
	igmpQuery(interface, 1)
    else
	local igmp_raw = igmpRaw(interface, version)

	local ip_raw = bin.pack("H", "45c00040ed780000010218bc0a00c8750a00c86b") .. igmp_raw
	local igmp_packet = packet.Packet:new(ip_raw, ip_raw:len())
	igmp_packet:ip_set_bin_src(ipOps.ip_to_str(srcip))
	igmp_packet:ip_set_bin_dst(ipOps.ip_to_str(dstip))
	igmp_packet:ip_set_len(#igmp_packet.buf)
	igmp_packet:ip_count_checksum()

	local sock = nmap.new_dnet()
	sock:ethernet_open(interface.device)
	 
	-- Ethernet IPv4 multicast, our ethernet address and type IP
	local eth_hdr = bin.pack("HAH", "01 00 5e 00 00 01", interface.mac, "08 00")
	sock:ethernet_send(eth_hdr .. igmp_packet.buf)
	sock:ethernet_close()
    end
end

-- Function to compare wieght of an IGMP response message.
-- Used to sort elements in responses table.
local respCompare = function(a,b)
    return ipOps.todword(a.src) + a.type + (a.ngroups or ipOps.todword(a.group))
	   < ipOps.todword(b.src) + b.type + (b.ngroups or ipOps.todword(b.group))
end


action = function(host, port)
    local timeout = tonumber(stdnse.get_script_args(SCRIPT_NAME .. ".timeout")) or 7
    local version = stdnse.get_script_args(SCRIPT_NAME .. ".version") or 2
    local interface = stdnse.get_script_args(SCRIPT_NAME .. ".interface")
    timeout = timeout * 1000
    if version ~= 'all' then
	version = tonumber(version)
    end

    local responses, results, interfaces, lthreads = {}, {}, {}, {}
    local result, grouptable, sourcetable

    -- Check the interface
    interface = interface or nmap.get_interface()
    if interface then
	-- Get the interface information
	interface = nmap.get_interface_info(interface)
	if not interface then
	    return ("ERROR: Failed to retreive %s interface information."):format(interface)
	end
	interfaces = {interface}
	stdnse.print_debug("%s: Will use %s interface.", SCRIPT_NAME, interface.shortname)
    else
	local ifacelist = nmap.list_interfaces()
	for _, iface in ipairs(ifacelist) do
	    -- Match all ethernet interfaces
	    if iface.address and iface.link=="ethernet" and 
		iface.address:match("%d+%.%d+%.%d+%.%d+") then

		stdnse.print_debug("%s: Will use %s interface.", SCRIPT_NAME, iface.shortname)
		table.insert(interfaces, iface)
	    end
	end
    end


    -- We should iterate over interfaces
    for _, interface in pairs(interfaces) do
	local co = stdnse.new_thread(igmpListener, interface, timeout, responses)
	igmpQuery(interface, version)
	lthreads[co] = true
    end

    local condvar = nmap.condvar(responses)
    -- Wait for the listening threads to finish
    repeat
	for thread in pairs(lthreads) do
	    if coroutine.status(thread) == "dead" then lthreads[thread] = nil end
	end
	if ( next(lthreads) ) then
		condvar("wait")
	end
    until next(lthreads) == nil;

    -- Output useful info from the responses
    if #responses > 0 then
	-- We should sort our list here.
	-- This is useful to have consistent results for tools such as Ndiff.
	table.sort(responses, respCompare)

	for _, response in pairs(responses) do
	    result = {}
	    result.name = response.src
	    table.insert(result, "Interface: " .. response.interface)
	    -- Add to new targets if newtargets script arg provided
	    if target.ALLOW_NEW_TARGETS then target.add(response.src) end
	    if response.type == 0x12 then
		table.insert(result, "Version: 1")
		table.insert(result, "Multicast group: ".. response.group)
	    elseif response.type == 0x16 then
		table.insert(result, "Version: 2")
		table.insert(result, "Group: ".. response.group)
	    elseif response.type == 0x22 then
		table.insert(result, "Version: 3")
		for _, group in pairs(response.groups) do
		    grouptable = {}
		    grouptable.name = "Group: " .. group.address
		    if group.mode == 0x01 then
			table.insert(grouptable, "Mode: INCLUDE")
		    elseif group.mode == 0x02 then
			table.insert(grouptable, "Mode: EXCLUDE")
		    end
		    if group.nsrc > 0 then
			sourcetable = {}
			sourcetable.name = "Sources:"
			table.insert(sourcetable, group.src)
			table.insert(grouptable, sourcetable) 
		    end
		    table.insert(result, grouptable)
		end
	    end
	    table.insert(results, result)
	end
	if #results>0 and not target.ALLOW_NEW_TARGETS then
	    table.insert(results,"Use the newtargets script-arg to add the results as targets")
	end
	return stdnse.format_output(true, results)
    end
end