#!/usr/bin/ruby
#
# DNS Balance --- 動的負荷分散を行なう DNS サーバ
#
# By: YOKOTA Hiroshi <yokota@netlab.is.tsukuba.ac.jp>

# $Id: dns_balance.rb,v 1.23 2003/01/31 21:50:34 elca Exp $

# DNS Balance の存在するパス名
if ENV["ROOT"] == nil
  print("ROOT environment required")
  exit(111)
else
  Prefix = ENV["ROOT"]
end
$LOAD_PATH.unshift(Prefix)

require 'socket'
require 'thread'
require 'getopts'

require 'datatype.rb'
require 'multilog.rb'
require 'log_writer.rb'
require 'util.rb'
require 'as_search.rb'

require 'namespace.rb'
require 'addrdb.rb'

#####################################################################
# ユーザ定義例外
class DnsNotImplementedError < StandardError ; end
class DnsTruncatedError      < StandardError ; end
class DnsNoQueryError        < StandardError ; end
class DnsNoMoreResourceError < StandardError ; end

###################################################################
# 関数

def usage()
  print "Usage: ruby dns_balance.rb [-lh] [-i ipaddr]\n"
  print "       -i ipaddr  listen IP address (default:0.0.0.0)\n"
  print "       -h         help message\n"
  print "       --as       Enable AS namespace\n"
  exit(111)
end

# DNS パケットから質問内容と質問のタイプと質問のクラスを取り出す
def parse_packet(packet)
  (number, flags, num_q, ans_rr, ort_rr, add_rr, str) =  packet.unpack("a2 a2 a2 a2 a2 a2 a*")

  if num_q != "\0\1"
    ML.log("Q must be 1")
    return [nil, nil, nil]
  end

  qlen = str.split("\0")[0].length + 1

  (q, q_type, q_class) = str.unpack("a#{qlen} a2 a2")

  return [q, q_type, q_class]
end

# クライアントのIPアドレスを返す
def parse_client_addr(str)
  (family, port, ipaddr, pad) = str.unpack("a2 a2 a4 a*")

  return ipaddr
end

# クライアントのIPアドレスによって返答内容を変える事が出来る
# 名前空間を選択。当てはまる物がなければ "default" になる
def select_namespace(addrstr, name)

  netaddrs = [ipstr_to_str(addrstr)] + ip_masklist(ipstr_to_str(addrstr))

  # custom namespace
  netaddrs.each {
    |i|
    if $namespace_db[i]                  != nil &&
	$addr_db[$namespace_db[i]]       != nil &&
	$addr_db[$namespace_db[i]][name] != nil
      return $namespace_db[i]
    end
  }

  # address number namespace
  netaddrs.each {
    |i|
    if $addr_db[i]        != nil &&
	$addr_db[i][name] != nil
      return i
    end
  }

  # AS namespace
  if $OPT_as &&
      # RFC1918 / プライベートアドレスはどこの AS にも属していない
      ipstr_to_str(ipstr_mask(addrstr,  8)) != "10.0.0.0"      &&
      ipstr_to_str(ipstr_mask(addrstr, 12)) != "172.16.0.0"    &&
      ipstr_to_str(ipstr_mask(addrstr, 16)) != "192.168.0.0"   &&
      ipstr_to_str(ipstr_mask(addrstr, 21)) != "204.152.184.0" &&
      ipstr_to_str(addrstr                ) != "127.0.0.1"

    as = as_search(ipstr_to_str(addrstr))
    if as                  != nil &&
	$addr_db[as]       != nil &&
	$addr_db[as][name] != nil
      return as
    end
  end

  return "default"
end

# 重みつき変数のための表を作る
def make_rand_array(namespace, name)
  rnd_max = 0
  rnd_slesh = []

  $addr_db[namespace][name].each {
    |i|
    rnd_max += (10000 - min(10000, i[1])) # badness の最大値は 10000
    rnd_slesh.push(rnd_max)
  }

  return [rnd_max, rnd_slesh]
end

# 重みつき乱数で選択
def select_rand_array(namespace, name, size)
  (rnd_max, rnd_slesh) = make_rand_array(namespace, name)

  if rnd_max == 0  # 全てのホストの Badness が 10000 だった
    return []
  end

  arr = []
  (0...size).each {
    |i|
    rnd = rand(rnd_max)
    (0...rnd_slesh.size).each {
      |j|
      if rnd <= rnd_slesh[j]
	arr.push(j)
	break
      end
    }
  }

  return arr
end

# パケットの正当性チェック
def check_packet(q, q_type, q_class)
  # ゾーン転送は無し
  if q_type == DnsType::AXFR
    ML.log("AXFR: " + q.dump + ":" + q_type.dump + ":" + q_class.dump)
    raise DnsNotImplementedError
  end

  # IP(UDP) のみ受け付け
  if !(q_class == DnsClass::INET || q_class == DnsClass::ANY)
    ML.log("noIP: " + q.dump + ":" + q_type.dump + ":" + q_class.dump)
    raise DnsNoQueryError
  end

  # 使用不可な文字がある
  if (q =~ /[()<>@,;:\\\"\.\[\]]/) != nil
    ML.log("char: " + q.dump + ":" + q_type.dump + ":" + q_class.dump)
    raise DnsNoQueryError
  end

  # A/ANY レコードのみ受け付け
  if q_type != DnsType::A && q_type != DnsType::ANY
    ML.log("noA: " + q.dump + ":" + q_type.dump + ":" + q_class.dump)
    raise DnsNoQueryError
  end

end

######################################################################
# main

srand()
getopts("lh", "i:0.0.0.0", "d:", "as")

usage() if $OPT_h

ML = MultiLog.new
ML.open

ML.log("dir: " + Prefix)
ML.log("start")


#
# アドレスデータベースの動的更新
#
Thread.start do
  loop do
    if test(?r, Prefix + "/addr") || test(?r, "./addr")
      begin
	load("addr")

	ML.log("reload")
      rescue NameError,SyntaxError,Exception
	ML.log("reload failed")
      end
    end

    #if test(?r, Prefix + "/addr-once") || test(?r, "./addr-once")
    #  Thread.exit 
    #end

    #p $addr_db
    sleep(5*60) # 5 分毎に更新
  end
end

gs = Socket.open(Socket::AF_INET, Socket::SOCK_DGRAM, 0)
gs.bind([Socket::AF_INET, 53, str_to_ipstr($OPT_i), ""].pack("s n a4 a8"))

#
# メインループ
#
loop do
  (packet, client_addr) = gs.recvfrom(512)
  Thread.start do
    $SAFE = 2
    begin
      client = parse_client_addr(client_addr)
      (q, q_type, q_class) = parse_packet(packet)
      check_packet(q, q_type, q_class) # -> NoQuery, NotImpl

      name      = dnsstr_to_str(q).downcase
      namespace = select_namespace(client, name)

      if $addr_db[namespace] == nil || $addr_db[namespace][name] == nil
	raise DnsNoQueryError
      end

      if $addr_db[namespace][name].size > 1  # -> NameError -> NoQuery
	size = 8
      else
	size = 1
      end
      a_array = select_rand_array(namespace, name, size).uniq

      if a_array.size == 0
	raise DnsNoMoreResourceError
      end

      # 返答生成
      r = packet[0,12] + q + q_type + q_class
      r[2] |= 0x84   # answer & authenticated
      r[3] = 0  # no error

      ans_addrs = []
      a_array.each {
	|i|
	addr = $addr_db[namespace][name][i][0]
	ans_addrs.push(addr.pack("CCCC"))   # ログ出力用

	# TTL 選択。近い所がある事が分かっているなら TTL を長くする
	if (a_array.size == 1)
	  ttl = "\0\0\x0e\x10" # 1時間
	else
	  ttl = "\0\0\0\x3c"   # 60秒
	end

	# 返答生成。 オフセットは 0x000c
	r += "\xc0\x0c" + DnsType::A + DnsClass::INET + ttl + "\0\4" + addr.pack("CCCC")
      }

      # 返答の数をセット
      r[6..7] = [a_array.size].pack("n")

      # 長過ぎたら削る
      if r.length > 512
	raise DnsTruncatedError
      end

      status = "ok"

    rescue DnsNotImplementedError
      r = packet[0,12] + q + q_type + q_class
      r[2] |= 0x80  # answer
      r[2] &= ~0x04 # not authenticated
      r[3] = 0
      r[3] |= 0x04  # not implemented error
      status = "NotImpl"

    rescue DnsTruncatedError
      # 長過ぎる時は削ってフラグを立てる
      r = r[0,512]
      r[2] |= 0x02
      status = "Truncated"

    rescue DnsNoQueryError,DnsNoMoreResourceError,NameError,StandardError
      r = packet[0,12] + q + q_type + q_class
      r[2] |= 0x84  # answer & authenticated
      r[3] = 0
      r[3] |= 0x03  # name error

      status = "NoQuery"

    rescue
      # ここには来ないはず
      r = packet[0,12] + q + q_type + q_class
      r[2] |= 0x80  # answer
      r[2] &= ~0x04 # not authenticated
      r[3] = 0
      r[3] |= 0x05  # query refused error
      status = "other"
    end

    #print packet.dump, "\n"
    #print r.dump, "\n"
    #p q

    gs.send(r, 0, client_addr)

    logger(ML, client, status, name, namespace, ans_addrs)

  end
end

# end
