local dns = require "dns"
local ipOps = require "ipOps"
local nmap = require "nmap"
local shortport = require "shortport"
local stdnse = require "stdnse"
local table = require "table"

description = [[
Performs a domain lookup using the edns-client-subnet option which
allows clients to specify the subnet that queries supposedly originate
from.  The script uses this option to supply a number of
geographically distributed locations in an attempt to enumerate as
many different address records as possible. The script also supports
requests using a given subnet.

* http://tools.ietf.org/html/draft-vandergaast-edns-client-subnet-00
]]

---
-- @usage
--   nmap -sU -p 53 --script dns-client-subnet-scan  --script-args \
--     dns-client-subnet-scan.domain=www.example.com, \
--     dns-client-subnet-scan.address=192.168.0.1 \
--     [,dns-client-subnet.nameserver=8.8.8.8] \
--     [,dns-client-subnet.mask=24] <target>
--   nmap --script dns-client-subnet-scan --script-args \
--     dns-client-subnet-scan.domain=www.example.com, \
--     dns-client-subnet-scan.address=192.168.0.1 \
--     dns-client-subnet.nameserver=8.8.8.8, \
--     [,dns-client-subnet.mask=24]
--
-- @output
-- 53/udp open  domain  udp-response
-- | dns-client-subnet-scan: 
-- | www.google.com
-- |   1.2.3.4
-- |   5.6.7.8
-- |   9.10.11.12
-- |   13.14.15.16
-- |   .
-- |   .
-- |_  .
---
-- @args dns-client-subnet.domain The domain to lookup eg. www.example.org
-- @args dns-client-subnet.address The client subnet address to use
-- @args dns-client-subnet.mask [optional] The number of bits to use as subnet mask (default: 24)
-- @args dns-client-subnet.nameserver [optional] nameserver to use.  (default = host.ip)
--

author = "John Bond"
license = "Simplified (2-clause) BSD license--See http://nmap.org/svn/docs/licenses/BSD-simplified"
categories = {"discovery", "safe"}


local argNS = stdnse.get_script_args(SCRIPT_NAME .. '.nameserver')
local argDomain = stdnse.get_script_args(SCRIPT_NAME .. '.domain')
local argMask = stdnse.get_script_args(SCRIPT_NAME .. '.mask') or 24
local argAddr = stdnse.get_script_args(SCRIPT_NAME .. '.address')

prerule = function()
	if ( not(argDomain) or nmap.address_family() ~= "inet" ) then
		return false
	end
	return true
end

portrule = function(host, port)
	if ( nmap.address_family() ~= "inet" ) then
		return false
	else
		return shortport.port_or_service(53, "domain", {"tcp", "udp"})(host, port)
	end
end

local areaIPs = {
	A4 = {ip=47763456, desc="GB,A4,Bath"},
	A5 = {ip=1043402336, desc="GB,A5,Biggleswade"},
	A6 = {ip=1364222182, desc="FR,A6,Chèvremont"},
	A7 = {ip=35357952, desc="GB,A7,Birmingham"},
	A8 = {ip=1050694009, desc="FR,A8,Romainville"},
	A9 = {ip=534257152, desc="FR,A9,Montpellier"},
	AB = {ip=2156920832, desc="CA,AB,Edmonton"},
	AK = {ip=202125312, desc="US,AK,Anchorage"},
	B1 = {ip=1041724648, desc="FR,B1,Robert"},
	B2 = {ip=35138048, desc="GB,B2,Bournemouth"},
	B3 = {ip=33949696, desc="FR,B3,Toulouse"},
	B4 = {ip=1050704998, desc="FR,B4,Lomme"},
	B5 = {ip=35213312, desc="GB,B5,Wembley"},
	B6 = {ip=773106752, desc="FR,B6,Amiens"},
	B7 = {ip=35148800, desc="GB,B7,Bristol"},
	B8 = {ip=786088496, desc="FR,B8,Valbonne"},
	B9 = {ip=33753088, desc="FR,B9,Lyon"},
	BC = {ip=201674096, desc="CA,BC,Victoria"},
	C1 = {ip=522223616, desc="FR,C1,Strasbourg"},
	C2 = {ip=41598976, desc="GB,C2,Halifax"},
	C3 = {ip=534676272, desc="GB,C3,Cambridge"},
	C5 = {ip=1043410032, desc="GB,C5,Runcorn"},
	C6 = {ip=773987544, desc="GB,C6,Saltash"},
	C7 = {ip=35165184, desc="GB,C7,Coventry"},
	C8 = {ip=35248128, desc="GB,C8,Croydon"},
	C9 = {ip=1892301824, desc="PH,C9,Iloilo"},
	D1 = {ip=35414016, desc="GB,D1,Darlington"},
	D2 = {ip=35164672, desc="GB,D2,Derby"},
	D3 = {ip=35301376, desc="GB,D3,Chesterfield"},
	D4 = {ip=1043450424, desc="GB,D4,Barnstaple"},
	D5 = {ip=2036385792, desc="PH,D5,Legaspi"},
	D7 = {ip=41451520, desc="GB,D7,Dudley"},
	D8 = {ip=35279104, desc="GB,D8,Durham"},
	D9 = {ip=460228608, desc="PH,D9,Manila"},
	DC = {ip=68514448, desc="US,DC,Washington"},
	E1 = {ip=1040645056, desc="GB,E1,Beverley"},
	E2 = {ip=35206912, desc="GB,E2,Brighton"},
	E3 = {ip=47822848, desc="GB,E3,Enfield"},
	E4 = {ip=39874560, desc="GB,E4,Colchester"},
	E5 = {ip=35270656, desc="GB,E5,Gateshead"},
	E6 = {ip=1368606720, desc="GB,E6,Coleford"},
	E7 = {ip=1051376056, desc="GB,E7,Woolwich"},
	E8 = {ip=1044737528, desc="GB,E8,Hackney"},
	F1 = {ip=1043451648, desc="GB,F1,Hammersmith"},
	F2 = {ip=35176448, desc="GB,F2,Basingstoke"},
	F4 = {ip=47998976, desc="GB,F4,Harrow"},
	F5 = {ip=1040622704, desc="GB,F5,Hart"},
	F6 = {ip=35230720, desc="GB,F6,Romford"},
	F8 = {ip=35214848, desc="GB,F8,Watford"},
	F9 = {ip=41693184, desc="GB,F9,Uxbridge"},
	G1 = {ip=41437184, desc="GB,G1,Hounslow"},
	G2 = {ip=35188224, desc="GB,G2,Ryde"},
	G3 = {ip=41861120, desc="GB,G3,Islington"},
	G4 = {ip=1040704992, desc="GB,G4,Kensington"},
	G5 = {ip=41506816, desc="GB,G5,Ashford"},
	G6 = {ip=786894336, desc="GB,G6,Hull"},
	G8 = {ip=40112128, desc="GB,G8,Huddersfield"},
	G9 = {ip=1380217968, desc="GB,G9,Knowsley"},
	H1 = {ip=1044731464, desc="GB,H1,Lambeth"},
	H2 = {ip=3512017264, desc="GB,H2,Earby"},
	H3 = {ip=35221504, desc="GB,H3,Leeds"},
	H4 = {ip=35158016, desc="GB,H4,Leicester"},
	H5 = {ip=1043402716, desc="GB,H5,Loughborough"},
	H6 = {ip=41732608, desc="GB,H6,Catford"},
	H7 = {ip=41863168, desc="GB,H7,Lincoln"},
	H8 = {ip=35294976, desc="GB,H8,Liverpool"},
	H9 = {ip=35196928, desc="GB,H9,London"},
	I1 = {ip=35253760, desc="GB,I1,Luton"},
	I2 = {ip=35263488, desc="GB,I2,Manchester"},
	I3 = {ip=47714304, desc="GB,I3,Rochester"},
	I4 = {ip=1298651136, desc="GB,I4,Morden"},
	I5 = {ip=1382961968, desc="GB,I5,Middlesborough"},
	I8 = {ip=1371219061, desc="GB,I8,Stepney"},
	I9 = {ip=35282944, desc="GB,I9,Norwich"},
	IA = {ip=201438272, desc="US,IA,Urbandale"},
	J1 = {ip=523578880, desc="GB,J1,Daventry"},
	J2 = {ip=788492344, desc="GB,J2,Grimsby"},
	J3 = {ip=3282790208, desc="GB,J3,Flixborough"},
	J5 = {ip=41759232, desc="GB,J5,Wallsend"},
	J6 = {ip=1043412268, desc="GB,J6,Alnwick"},
	J7 = {ip=41783296, desc="GB,J7,Harrogate"},
	J8 = {ip=35160064, desc="GB,J8,Nottingham"},
	J9 = {ip=47742976, desc="GB,J9,Newark"},
	JA = {ip=1476096512, desc="RU,JA,Kurilsk"},
	K1 = {ip=48015360, desc="GB,K1,Oldham"},
	K2 = {ip=1043402360, desc="GB,K2,Kidlington"},
	K3 = {ip=39956480, desc="GB,K3,Peterborough"},
	K4 = {ip=41735168, desc="GB,K4,Plymouth"},
	K5 = {ip=775747568, desc="GB,K5,Poole"},
	K6 = {ip=774162844, desc="GB,K6,Portsmouth"},
	K7 = {ip=41746432, desc="GB,K7,Reading"},
	K8 = {ip=35229696, desc="GB,K8,Ilford"},
	L1 = {ip=47773696, desc="GB,L1,Twickenham"},
	L2 = {ip=48103424, desc="GB,L2,Rochdale"},
	L3 = {ip=35304192, desc="GB,L3,Rotherham"},
	L4 = {ip=1043416984, desc="GB,L4,Oakham"},
	L5 = {ip=772988024, desc="GB,L5,Salford"},
	L6 = {ip=35336192, desc="GB,L6,Shrewsbury"},
	L7 = {ip=1043419464, desc="GB,L7,Oldbury"},
	L8 = {ip=39936000, desc="GB,L8,Lytham"},
	L9 = {ip=35304448, desc="GB,L9,Sheffield"},
	M1 = {ip=35384320, desc="GB,M1,Slough"},
	M2 = {ip=41470976, desc="GB,M2,Solihull"},
	M4 = {ip=35139584, desc="GB,M4,Southampton"},
	M5 = {ip=1043402176, desc="GB,M5,Southend-on-sea"},
	M6 = {ip=773986248, desc="GB,M6,Hill"},
	M8 = {ip=1443330688, desc="GB,M8,Camberwell"},
	M9 = {ip=35322880, desc="GB,M9,Stafford"},
	MB = {ip=1076550400, desc="CA,MB,Winnipeg"},
	MI = {ip=201393888, desc="US,MI,Saginaw"},
	N1 = {ip=1318741928, desc="GB,N1,Haydock"},
	N2 = {ip=35266560, desc="GB,N2,Stockport"},
	N3 = {ip=41832448, desc="GB,N3,Stockton-on-tees"},
	N4 = {ip=3231559680, desc="GB,N4,Longport"},
	N5 = {ip=1043424608, desc="GB,N5,Beccles"},
	N6 = {ip=35276800, desc="GB,N6,Sunderland"},
	N7 = {ip=41551872, desc="GB,N7,Tadworth"},
	N8 = {ip=41697280, desc="GB,N8,Sutton"},
	N9 = {ip=35252736, desc="GB,N9,Swindon"},
	NB = {ip=2211053568, desc="CA,NB,Fredericton"},
	ND = {ip=201473536, desc="US,ND,Bismarck"},
	NH = {ip=201772808, desc="US,NH,Laconia"},
	NJ = {ip=201352704, desc="US,NJ,Piscataway"},
	NS = {ip=3226164992, desc="CA,NS,Halifax"},
	NT = {ip=3332472320, desc="CA,NT,Yellowknife"},
	NV = {ip=202261184, desc="US,NV,Henderson"},
	O2 = {ip=40251392, desc="GB,O2,Telford"},
	O3 = {ip=35230208, desc="GB,O3,Grays"},
	O4 = {ip=35318784, desc="GB,O4,Torquay"},
	O5 = {ip=1368498352, desc="GB,O5,Poplar"},
	O6 = {ip=1546138112, desc="GB,O6,Stretford"},
	O7 = {ip=35219456, desc="GB,O7,Wakefield"},
	O8 = {ip=35321856, desc="GB,O8,Walsall"},
	O9 = {ip=1359108248, desc="GB,O9,Walthamstow"},
	ON = {ip=201620304, desc="CA,ON,Ottawa"},
	P1 = {ip=1043431736, desc="GB,P1,Wandsworth"},
	P2 = {ip=35260416, desc="GB,P2,Warrington"},
	P3 = {ip=41766912, desc="GB,P3,Nuneaton"},
	P4 = {ip=41893888, desc="GB,P4,Newbury"},
	P5 = {ip=772987648, desc="GB,P5,Westminster"},
	P7 = {ip=41466624, desc="GB,P7,Wigan"},
	P8 = {ip=48087808, desc="GB,P8,Salisbury"},
	P9 = {ip=41793536, desc="GB,P9,Maidenhead"},
	Q1 = {ip=41457664, desc="GB,Q1,Wallasey"},
	Q2 = {ip=1040739840, desc="GB,Q2,Wokingham"},
	Q3 = {ip=35323392, desc="GB,Q3,Wolverhampton"},
	Q4 = {ip=539624744, desc="GB,Q4,Redditch"},
	Q5 = {ip=1043415688, desc="GB,Q5,Wetherby"},
	Q6 = {ip=1043439984, desc="GB,Q6,Antrim"},
	Q7 = {ip=41811456, desc="GB,Q7,Newtownards"},
	Q8 = {ip=1347208672, desc="GB,Q8,Armagh"},
	Q9 = {ip=1044726432, desc="GB,Q9,Connor"},
	QC = {ip=2210594816, desc="CA,QC,Varennes"},
	R1 = {ip=1482707288, desc="GB,R1,Ballymoney"},
	R3 = {ip=47828992, desc="GB,R3,Belfast"},
	R4 = {ip=1051352576, desc="GB,R4,Eden"},
	R5 = {ip=1056827328, desc="GB,R5,Castlereagh"},
	R6 = {ip=47895040, desc="GB,R6,Coleraine"},
	R7 = {ip=3270400320, desc="GB,R7,Dunmore"},
	R8 = {ip=1367996672, desc="GB,R8,Portadown"},
	R9 = {ip=773985608, desc="GB,R9,Square"},
	RI = {ip=67285760, desc="US,RI,Providence"},
	S1 = {ip=1040409048, desc="GB,S1,Drummond"},
	S2 = {ip=1353842208, desc="GB,S2,Enniskillen"},
	S3 = {ip=1368133632, desc="GB,S3,Larne"},
	S4 = {ip=1446384520, desc="GB,S4,Ardmore"},
	S5 = {ip=1043419184, desc="GB,S5,Lisburn"},
	S6 = {ip=1056826304, desc="GB,S6,Londonderry"},
	S7 = {ip=1359111383, desc="GB,S7,Curran"},
	S8 = {ip=1369435392, desc="GB,S8,Waterfoot"},
	S9 = {ip=1043434592, desc="GB,S9,Newry"},
	T1 = {ip=3242033152, desc="GB,T1,Jordanstown"},
	T2 = {ip=1043402000, desc="GB,T2,Bangor"},
	T3 = {ip=1043429728, desc="GB,T3,Omagh"},
	T4 = {ip=1043429520, desc="GB,T4,Strabane"},
	T5 = {ip=39849984, desc="GB,T5,Aberdeen"},
	T6 = {ip=1043407024, desc="GB,T6,Inverurie"},
	T7 = {ip=47917056, desc="GB,T7,Forfar"},
	T8 = {ip=1051457600, desc="GB,T8,Sandbank"},
	T9 = {ip=1043429424, desc="GB,T9,Melrose"},
	TX = {ip=201673024, desc="US,TX,Mckinney"},
	U1 = {ip=1043400976, desc="GB,U1,Alloa"},
	U2 = {ip=1353815544, desc="GB,U2,Langholm"},
	U3 = {ip=1042190336, desc="GB,U3,Dundee"},
	U4 = {ip=1043428036, desc="GB,U4,Newmilns"},
	U5 = {ip=1051334704, desc="GB,U5,Bishopbriggs"},
	U6 = {ip=1040628912, desc="GB,U6,Musselburgh"},
	U7 = {ip=1056881248, desc="GB,U7,Barrhead"},
	U8 = {ip=35188736, desc="GB,U8,Edinburgh"},
	U9 = {ip=1318744616, desc="GB,U9,Blackstone"},
	V1 = {ip=47947776, desc="GB,V1,Kirkcaldy"},
	V2 = {ip=35190784, desc="GB,V2,Glasgow"},
	V4 = {ip=1043417560, desc="GB,V4,Greenock"},
	V5 = {ip=3570359128, desc="GB,V5,Borthwick"},
	V6 = {ip=1398983520, desc="GB,V6,Findhorn"},
	V7 = {ip=1043452928, desc="GB,V7,Saltcoats"},
	V8 = {ip=523564544, desc="GB,V8,Bothwell"},
	V9 = {ip=1353706504, desc="GB,V9,Redland"},
	VT = {ip=201355264, desc="US,VT,Brattleboro"},
	W1 = {ip=1042195200, desc="GB,W1,Perth"},
	W2 = {ip=1043412560, desc="GB,W2,Paisley"},
	W4 = {ip=1056825616, desc="GB,W4,Dundonald"},
	W5 = {ip=1040411544, desc="GB,W5,Douglas"},
	W6 = {ip=41547776, desc="GB,W6,Stirling"},
	W7 = {ip=1443523584, desc="GB,W7,Bearsden"},
	W8 = {ip=534572928, desc="GB,W8,Cross"},
	W9 = {ip=1042221056, desc="GB,W9,Livingston"},
	WA = {ip=201806720, desc="US,WA,Issaquah"},
	WY = {ip=135495936, desc="US,WY,Casper"},
	X1 = {ip=1043425760, desc="GB,X1,Valley"},
	X2 = {ip=773988152, desc="GB,X2,Victoria"},
	X3 = {ip=35149824, desc="GB,X3,Bridgend"},
	X4 = {ip=1043402272, desc="GB,X4,Blackwood"},
	X5 = {ip=39946240, desc="GB,X5,Cardiff"},
	X6 = {ip=1043435700, desc="GB,X6,Aberystwyth"},
	X7 = {ip=1043408760, desc="GB,X7,Llanelli"},
	X8 = {ip=1368926208, desc="GB,X8,Abergele"},
	X9 = {ip=1043411032, desc="GB,X9,Rhyl"},
	Y1 = {ip=1043407256, desc="GB,Y1,Holywell"},
	Y2 = {ip=1043401576, desc="GB,Y2,Caernarfon"},
	Y4 = {ip=1043428692, desc="GB,Y4,Cwmbran"},
	Y5 = {ip=3265794544, desc="GB,Y5,Cwmafan"},
	Y6 = {ip=35153920, desc="GB,Y6,Newport"},
	Y7 = {ip=1353763984, desc="GB,Y7,Haverfordwest"},
	Y8 = {ip=1043430344, desc="GB,Y8,Welshpool"},
	Z1 = {ip=40116224, desc="GB,Z1,Swansea"},
	Z2 = {ip=40189952, desc="GB,Z2,Pontypool"},
	Z3 = {ip=35147776, desc="GB,Z3,Barry"},
	Z4 = {ip=40321024, desc="GB,Z4,Wrexham"}
}

local get_addresses = function(address, mask, domain, nameserver) 

	-- translate the IP's in the areaIPs to strings, as this is what the
	-- DNS library expects
	if ( "number" == type(address) ) then
		address = ipOps.fromdword(address)
		local a, b, c, d = address:match("(%d+)%.(%d+)%.(%d+)%.(%d+)")
		address = ("%d.%d.%d.%d"):format(d,c,b,a)
	end

	local subnet = { family = nmap.address_family(), address = address, mask = mask }
	local status, resp = dns.query(domain, {host = nameserver, retAll=true, subnet=subnet})
	if ( not(status) ) then
		return
	end
	if ( "table" ~= type(resp) ) then resp = { resp } end
	return resp
end

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

action = function(host, port)	

	if ( not(argDomain) ) then
		return fail(SCRIPT_NAME .. ".domain was not specified")
	end
	
	local nameserver = argNS or (host and host.ip)
	-- as the nameserver argument overrides the host.ip, the prerule should
	-- already have done our work, so abort
	if ( argNS and host ) then
		return
	-- if we have no nameserver argument and no host, we dont have sufficient
	-- information to continue, abort
	elseif ( not(argNS) and not(host) ) then
		return		
	end
	
	local addrs = argAddr or areaIPs
	if ( "string" == type(addrs) ) then	addrs = {{ ip = addrs }} end

	local lookup, result = {}, { name = argDomain }
	for _,ip in pairs(addrs) do
		for _, addr in ipairs( get_addresses (ip.ip, argMask, argDomain, nameserver) ) do
			lookup[addr] = true
		end
	end
	for addr in pairs(lookup) do table.insert(result, addr) end
	table.sort(result)
	return stdnse.format_output(true, result)
end
