# Simple DNS server by J.D. Zamfirescu
# based in part on a simple DNS server in python by Francisco Santos

import socket
import sys
import re
import os.path

def PRINT(str):
  print str
  sys.stdout.flush()

class Matcher:
  def __init__(self, path):
    self.path = path
    
  def getAddress(self, domain):
    for i in self.path:
      if i.matches(domain):
        return i.handle(domain)
    return None

  def __str__(self):
    return "Matcher with patterns:\n"+"\n".join([str(r) for r in self.path])

class RegexpToIP:
  def __init__(self, regexp, ip):
    self.regexp = regexp
    self.re = re.compile(regexp)
    self.ip = ip
    
  def matches(self, domain):
    return self.re.match(domain) is not None
    
  def handle(self, domain):
    return self.ip
  
  def __str__(self):
    return "Map for %s -> %s" % (self.regexp, self.ip)

class DNSQuery:
  def __init__(self, data, matcher):
    self.m = matcher
    self.data=data
    self.domain=''

    type = (ord(data[2]) >> 3) & 15   # Opcode bits
    if type == 0:                     # Standard query
      ini=12
      lon=ord(data[ini])
      while lon != 0:
        self.domain+=data[ini+1:ini+lon+1]+'.'
        ini+=lon+1
        lon=ord(data[ini])        

  def response(self):
    packet=''
    ip = ''
    if self.domain:
      ip = self.m.getAddress(self.domain)
      if ip is not None:
        packet+=self.data[:2] + "\x85\x80"
        packet+=self.data[4:6] + self.data[4:6] + '\x00\x00\x00\x00'   # Questions and Answers Counts
        packet+=self.data[12:]                                         # Original Domain Name Question
        packet+='\xc0\x0c'                                             # Pointer to domain name
        packet+='\x00\x01\x00\x01\x00\x00\x00\x3c\x00\x04'             # Response type, ttl and resource data length -> 4 bytes
        packet+=str.join('',map(lambda x: chr(int(x)), ip.split('.'))) # 4bytes of IP
      else:
        packet+=self.data[:2] + "\x81\x82"
        packet+= self.data[4:6] + '\x00\x00' + '\x00\x00\x00\x00'
        packet+= self.data[12:]
    return (packet, ip)

def usage():
  PRINT(sys.argv[0], " <regexp> <ip> [<regexp> <ip> ...]")
  PRINT("invalid: ", " ".join(sys.argv))

def parseArgs(args):
  path = []
  if len(args) % 2 != 0:
    usage()
    sys.exit(0)
  pairs = [(args[2*i], args[2*i+1]) for i in range(len(args)/2)]
  for (r, i) in pairs:
    path.append(RegexpToIP(r, i))
  return Matcher(path)


if __name__ == '__main__':
  if os.path.isfile("/etc/dns-listing"):
    m = parseArgs(file("/etc/dns-listing").read().split())
    last_update = os.path.getmtime("/etc/dns-listing")
  else:
    m = parseArgs(sys.argv[1:])
    last_update = 0
  PRINT('pyminifakeDNS:: dom.query. 60 IN A <ip>')
  PRINT(str(m))
  
  udps = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
  udps.bind(('',53))
  
  try:
    while 1:
      data, addr = udps.recvfrom(1024)
      # check data file:
      if os.path.isfile("/etc/dns-listing") and (os.path.getmtime("/etc/dns-listing") > last_update):
        m = parseArgs(file("/etc/dns-listing").read().split())  
        last_update = os.path.getmtime("/etc/dns-listing")
        PRINT('reparsing')
        PRINT(str(m))
      p=DNSQuery(data, m)
      res = p.response()
      udps.sendto(res[0], addr)
      PRINT('response: %s -> %s' % (p.domain, res[1]))
  except KeyboardInterrupt:
    PRINT('Finishing')
    udps.close()