-- public domain 20080404 lua@ztact.com

-- todo: quick (default) header generation

-- todo: nxdomain, error handling

-- todo: cache results of encodeName

-- reference: http://tools.ietf.org/html/rfc1035

-- reference: http://tools.ietf.org/html/rfc1876 (LOC)

require 'socket'

local ztact = require 'util.ztact'

local coroutine, io, math, socket, string, table =

      coroutine, io, math, socket, string, table

local get, set = ztact.get, ztact.set

-------------------------------------------------- module dns

module ('dns')

local dns = _M;

-- dns type & class codes ------------------------------ dns type & class codes

local append = table.insert

local function highbyte (i)    -- - - - - - - - - - - - - - - - - - -  highbyte

  return (i-(i%0x100))/0x100

  end

local function augment (t)    -- - - - - - - - - - - - - - - - - - - -  augment

  local a = {}

  for i,s in pairs (t) do  a[i] = s  a[s] = s  a[string.lower (s)] = s  end

  return a

  end

local function encode (t)    -- - - - - - - - - - - - - - - - - - - - -  encode

  local code = {}

  for i,s in pairs (t) do

    local word = string.char (highbyte (i), i %0x100)

    code[i] = word

    code[s] = word

    code[string.lower (s)] = word

    end

  return code

  end

dns.types = {

  'A', 'NS', 'MD', 'MF', 'CNAME', 'SOA', 'MB', 'MG', 'MR', 'NULL', 'WKS',

  'PTR', 'HINFO', 'MINFO', 'MX', 'TXT',

  [ 28] = 'AAAA', [ 29] = 'LOC',   [ 33] = 'SRV',

  [252] = 'AXFR', [253] = 'MAILB', [254] = 'MAILA', [255] = '*' }

dns.classes = { 'IN', 'CS', 'CH', 'HS', [255] = '*' }

dns.type      = augment (dns.types)

dns.class     = augment (dns.classes)

dns.typecode  = encode  (dns.types)

dns.classcode = encode  (dns.classes)

local function standardize (qname, qtype, qclass)    -- - - - - - - standardize

  if string.byte (qname, -1) ~= 0x2E then  qname = qname..'.'  end

  qname = string.lower (qname)

  return qname, dns.type[qtype or 'A'], dns.class[qclass or 'IN']

  end

local function prune (rrs, time, soft)    -- - - - - - - - - - - - - - -  prune

  time = time or socket.gettime ()

  for i,rr in pairs (rrs) do

    if rr.tod then

      -- rr.tod = rr.tod - 50    -- accelerated decripitude

      rr.ttl = math.floor (rr.tod - time)

      if rr.ttl <= 0 then  rrs[i] = nil  end

    elseif soft == 'soft' then    -- What is this?  I forget!

      assert (rr.ttl == 0)

      rrs[i] = nil

      end  end  end

-- metatables & co. ------------------------------------------ metatables & co.

local resolver = {}

resolver.__index = resolver

local SRV_tostring

local rr_metatable = {}    -- - - - - - - - - - - - - - - - - - -  rr_metatable

function rr_metatable.__tostring (rr)

  local s0 = string.format (

    '%2s %-5s %6i %-28s', rr.class, rr.type, rr.ttl, rr.name )

  local s1 = ''

  if rr.type == 'A' then  s1 = ' '..rr.a

  elseif rr.type == 'MX' then

    s1 = string.format (' %2i %s', rr.pref, rr.mx)

  elseif rr.type == 'CNAME' then  s1 = ' '..rr.cname

  elseif rr.type == 'LOC'   then  s1 = ' '..resolver.LOC_tostring (rr)

  elseif rr.type == 'NS'    then  s1 = ' '..rr.ns

  elseif rr.type == 'SRV'   then  s1 = ' '..SRV_tostring (rr)

  elseif rr.type == 'TXT'   then  s1 = ' '..rr.txt

  else  s1 = ' <UNKNOWN RDATA TYPE>'  end

  return s0..s1

  end

local rrs_metatable = {}    -- - - - - - - - - - - - - - - - - -  rrs_metatable

function rrs_metatable.__tostring (rrs)

  for i,rr in pairs (rrs) do  append (t, tostring (rr)..'\n')  end

  return table.concat (t)

  end

local cache_metatable = {}    -- - - - - - - - - - - - - - - -  cache_metatable

function cache_metatable.__tostring (cache)

  local time = socket.gettime ()

  local t = {}

  for class,types in pairs (cache) do

    for type,names in pairs (types) do

      for name,rrs in pairs (names) do

        prune (rrs, time)

        append (t, tostring (rrs))  end  end  end

  return table.concat (t)

  end

function resolver:new ()    -- - - - - - - - - - - - - - - - - - - - - resolver

  local r = { active = {}, cache = {}, unsorted = {} }

  setmetatable (r, resolver)

  setmetatable (r.cache, cache_metatable)

  setmetatable (r.unsorted, { __mode = 'kv' })

  return r

  end

-- packet layer -------------------------------------------------- packet layer

function dns.random (...)    -- - - - - - - - - - - - - - - - - - -  dns.random

  math.randomseed (10000*socket.gettime ())

  dns.random = math.random

  return dns.random (...)

  end

local function encodeHeader (o)    -- - - - - - - - - - - - - - -  encodeHeader

  o = o or {}

  o.id = o.id or		-- 16b	(random) id

    dns.random (0, 0xffff)

  o.rd = o.rd or 1		--  1b  1 recursion desired

  o.tc = o.tc or 0		--  1b	1 truncated response

  o.aa = o.aa or 0		--  1b	1 authoritative response

  o.opcode = o.opcode or 0	--  4b	0 query

				--      1 inverse query

				--	2 server status request

				--	3-15 reserved

  o.qr = o.qr or 0		--  1b	0 query, 1 response

  o.rcode = o.rcode or 0	--  4b  0 no error

				--	1 format error

				--	2 server failure

				--	3 name error

				--	4 not implemented

				--	5 refused

				--	6-15 reserved

  o.z  = o.z  or 0		--  3b  0 resvered

  o.ra = o.ra or 0		--  1b  1 recursion available

  o.qdcount = o.qdcount or 1	-- 16b	number of question RRs

  o.ancount = o.ancount or 0	-- 16b	number of answers RRs

  o.nscount = o.nscount or 0	-- 16b	number of nameservers RRs

  o.arcount = o.arcount or 0	-- 16b  number of additional RRs

  -- string.char() rounds, so prevent roundup with -0.4999

  local header = string.char (

    highbyte (o.id),  o.id %0x100,

    o.rd + 2*o.tc + 4*o.aa + 8*o.opcode + 128*o.qr,

    o.rcode + 16*o.z + 128*o.ra,

    highbyte (o.qdcount),  o.qdcount %0x100,

    highbyte (o.ancount),  o.ancount %0x100,

    highbyte (o.nscount),  o.nscount %0x100,

    highbyte (o.arcount),  o.arcount %0x100 )

  return header, o.id

  end

local function encodeName (name)    -- - - - - - - - - - - - - - - - encodeName

  local t = {}

  for part in string.gmatch (name, '[^.]+') do

    append (t, string.char (string.len (part)))

    append (t, part)

    end

  append (t, string.char (0))

  return table.concat (t)

  end

local function encodeQuestion (qname, qtype, qclass)    -- - - - encodeQuestion

  qname  = encodeName (qname)

  qtype  = dns.typecode[qtype or 'a']

  qclass = dns.classcode[qclass or 'in']

  return qname..qtype..qclass;

  end

function resolver:byte (len)    -- - - - - - - - - - - - - - - - - - - - - byte

  len = len or 1

  local offset = self.offset

  local last = offset + len - 1

  if last > #self.packet then

    error (string.format ('out of bounds: %i>%i', last, #self.packet))  end

  self.offset = offset + len

  return string.byte (self.packet, offset, last)

  end

function resolver:word ()    -- - - - - - - - - - - - - - - - - - - - - -  word

  local b1, b2 = self:byte (2)

  return 0x100*b1 + b2

  end

function resolver:dword ()    -- - - - - - - - - - - - - - - - - - - - -  dword

  local b1, b2, b3, b4 = self:byte (4)

  -- print ('dword', b1, b2, b3, b4)

  return 0x1000000*b1 + 0x10000*b2 + 0x100*b3 + b4

  end

function resolver:sub (len)    -- - - - - - - - - - - - - - - - - - - - - - sub

  len = len or 1

  local s = string.sub (self.packet, self.offset, self.offset + len - 1)

  self.offset = self.offset + len

  return s

  end

function resolver:header (force)    -- - - - - - - - - - - - - - - - - - header

  local id = self:word ()

  -- print (string.format (':header  id  %x', id))

  if not self.active[id] and not force then  return nil  end

  local h = { id = id }

  local b1, b2 = self:byte (2)

  h.rd      = b1 %2

  h.tc      = b1 /2%2

  h.aa      = b1 /4%2

  h.opcode  = b1 /8%16

  h.qr      = b1 /128

  h.rcode   = b2 %16

  h.z       = b2 /16%8

  h.ra      = b2 /128

  h.qdcount = self:word ()

  h.ancount = self:word ()

  h.nscount = self:word ()

  h.arcount = self:word ()

  for k,v in pairs (h) do  h[k] = v-v%1  end

  return h

  end

function resolver:name ()    -- - - - - - - - - - - - - - - - - - - - - -  name

  local remember, pointers = nil, 0

  local len = self:byte ()

  local n = {}

  while len > 0 do

    if len >= 0xc0 then    -- name is "compressed"

      pointers = pointers + 1

      if pointers >= 20 then  error ('dns error: 20 pointers')  end

      local offset = ((len-0xc0)*0x100) + self:byte ()

      remember = remember or self.offset

      self.offset = offset + 1    -- +1 for lua

    else    -- name is not compressed

      append (n, self:sub (len)..'.')

      end

    len = self:byte ()

    end

  self.offset = remember or self.offset

  return table.concat (n)

  end

function resolver:question ()    -- - - - - - - - - - - - - - - - - -  question

  local q = {}

  q.name  = self:name ()

  q.type  = dns.type[self:word ()]

  q.class = dns.type[self:word ()]

  return q

  end

function resolver:A (rr)    -- - - - - - - - - - - - - - - - - - - - - - - -  A

  local b1, b2, b3, b4 = self:byte (4)

  rr.a = string.format ('%i.%i.%i.%i', b1, b2, b3, b4)

  end

function resolver:CNAME (rr)    -- - - - - - - - - - - - - - - - - - - -  CNAME

  rr.cname = self:name ()

  end

function resolver:MX (rr)    -- - - - - - - - - - - - - - - - - - - - - - -  MX

  rr.pref = self:word ()

  rr.mx   = self:name ()

  end

function resolver:LOC_nibble_power ()    -- - - - - - - - - -  LOC_nibble_power

  local b = self:byte ()

  -- print ('nibbles', ((b-(b%0x10))/0x10), (b%0x10))

  return ((b-(b%0x10))/0x10) * (10^(b%0x10))

  end

function resolver:LOC (rr)    -- - - - - - - - - - - - - - - - - - - - - -  LOC

  rr.version = self:byte ()

  if rr.version == 0 then

    rr.loc           = rr.loc or {}

    rr.loc.size      = self:LOC_nibble_power ()

    rr.loc.horiz_pre = self:LOC_nibble_power ()

    rr.loc.vert_pre  = self:LOC_nibble_power ()

    rr.loc.latitude  = self:dword ()

    rr.loc.longitude = self:dword ()

    rr.loc.altitude  = self:dword ()

    end  end

local function LOC_tostring_degrees (f, pos, neg)    -- - - - - - - - - - - - -

  f = f - 0x80000000

  if f < 0 then  pos = neg  f = -f  end

  local deg, min, msec

  msec = f%60000

  f    = (f-msec)/60000

  min  = f%60

  deg = (f-min)/60

  return string.format ('%3d %2d %2.3f %s', deg, min, msec/1000, pos)

  end

function resolver.LOC_tostring (rr)    -- - - - - - - - - - - - -  LOC_tostring

  local t = {}

  --[[

  for k,name in pairs { 'size', 'horiz_pre', 'vert_pre',

			'latitude', 'longitude', 'altitude' } do

    append (t, string.format ('%4s%-10s: %12.0f\n', '', name, rr.loc[name]))

    end

  --]]

  append ( t, string.format (

    '%s    %s    %.2fm %.2fm %.2fm %.2fm',

    LOC_tostring_degrees (rr.loc.latitude, 'N', 'S'),