#
# -*- coding: utf-8 -*-
#
# Whois module for the pywhoisd server
 
import SocketServer, re, math, syslog
from ipcalc import IP, Network
 
def FormatOutput(field, value):
    """
    Format a field, value pair for the output
    from the server.
    """
    output = (field + ":").ljust(24)
    output += value + "\n"
    return " "*4 + output
 
def FormatMessage(message):
    """
    Prepare a message to be written from the
    server.
    """
    output = ""
    for line in message.split("\n"):
        line = line.strip()
        output += """\n  :: """ + line
    return output + "\n"
 
 
class Logger():
    """
    Logger class that allow pywhoisd to log
    everything on syslog.
    """
    def __init__(self, verbose = False):
        self.verbose = verbose
 
    def Log(self, message):
        """Log a message to syslog"""
        syslog.syslog ("[pywhoisd] " + message)
        if self.verbose:
            print message
 
    def __call__(self, message):
        self.Log (message)
 
class Config():
    """
    Config class manage configuration contained
    in the config file and acts like a dictionary
    that provides data to other classes
    """
    def __init__(self, config_file):
        """
        Load config file and save data in memory
        """
        self.config = {}
        f = open(config_file, 'r')
        content = f.read()
        f.close()
        for (key, value) in re.findall(r"([^\s]+)\s*=\s*([^;]+)\s*;", content):
            self.config[key] = value
 
    def __getitem__(self, key):
        if self.config.has_key(key):
            return self.config[key]
        else:
            return None
 
    def has_key(self, key):
        return self.config.has_key(key)
 
 
class WhoisRequestHandler(SocketServer.StreamRequestHandler):
    """
    This handler answers to the whois request that
    arrives to the WhoisServer.
    """
 
    def handle(self):
        # Come prima cosa salutiamo
        hello = self.server.config['welcome_banner']
        if hello:
            response = FormatMessage(hello)
        else:
            response = ""
 
        # Obtain the whois request
        request = self.rfile.readline().strip()
 
        if re.search(r"\d+\.\d+\.\d+\.\d+", request.strip()):
            response += self.server.ResolveAddress(request.split("/")[0])
        else:
            response += self.server.ResolveDomain(request)
        response += FormatMessage("Bye")
 
        self.request.send (response)
        return
 
 
class WhoisServer(SocketServer.ThreadingTCPServer):
    """
    The real server. We allow reusing of address
    so we can stop and start in a moment.
    """
    allow_reuse_address = True
    def __init__(self, (host, port), logger, config):
 
        SocketServer.ThreadingTCPServer.__init__(self, (host, port),
                                                WhoisRequestHandler)
        self.logger = logger
        self.config = config
        self.ip_database = {}
        self.domain_database = {}
        self.LoadDatabase(config['database_file'])
 
    def LoadDatabase(self, database_file):
        """
        Populate server database with the data in the database
        file
        """
        self.logger.Log ("Populating database with entry in %s" % database_file)
        try:
            f = open(database_file, 'r')
            database_data = f.read()
        except IOError:
            self.logger.Log ("Error while loading config file. exiting")
 
        nets = re.findall (r"net\s+(\d+\.\d+\.\d+\.\d+[/\d]*)[\s|\n]+\{([^\}]+)\};",
                           database_data)
        for net in nets:
            self.logger.Log ("Loading net %s" % net[0])
            fields = re.findall(r"(\w+)\s*=\s*([^;]+);", net[1])
            name, owner, ns = None, None, None
            data = {}
            for field, value in fields:
                if field.lower() == 'name':
                    name = value.strip()
                elif field.lower() == 'owner':
                    owner = value.strip()
                elif field.lower() == 'ns':
                    ns = value.strip()
                else:
                    data[field] = value
            self.ip_database[net[0]] = Net(name = name, owner = owner, ns = ns,
                                           data = data)
 
        domains = re.findall(r"domain\s+([^\s]+)[\s|\n]+\{([^\}]+)\};",
                             database_data)
        for domain in domains:
            self.logger.Log ("Loading domain %s" % domain[0])
            fields = re.findall(r"(\w+)\s*=\s*([^;]+);", domain[1])
            name, owner, ns = None, None, None
            data = {}
            for field, value in fields:
                if field.lower() == 'name':
                    name = value.strip()
                elif field.lower() == 'owner':
                    owner = value.strip()
                elif field.lower() == 'ns':
                    ns = value.strip()
                else:
                    data[field] = value
            self.domain_database[domain[0]] = Domain(name = name,
                                                     owner = owner,
                                                     ns = ns,
                                                     data = data)
 
    def ResolveDomain(self, domain):
        """
        Find domain data and return them ready for output
        for the server
        """
        domain = domain.strip(".")
        output = ""
        domains = self.GetParentDomains(domain)
        for d in domains:
            output += FormatMessage("GNet whoisd answering for domain %s" % d)
            output += FormatOutput ("Domain", d)
            domain = self.domain_database[d]
            if domain.name is not None:
                output += FormatOutput ("Name", domain.name)
            if domain.owner is not None:
                output += FormatOutput ("Owner", domain.owner)
            if domain.ns is not None:
                output += FormatOutput ("Nameserver", domain.ns)
            if len(domain.data.items()) > 0:
                output += FormatMessage("Additional data")
            for k, v in domain.data.items():
                output += FormatOutput (k.capitalize(),v)
        if len(domains) == 0:
            return FormatMessage("No domain found for %s" % domain)
        return output
 
 
    def GetParentDomains(self, domain):
        """
        Return a list of domains that contain the domain
        passed as argument
        """
        interesting_domains = []
        for domain_name, domain_data in self.domain_database.items():
            if domain.endswith(domain_name):
                interesting_domains.append (domain_name)
        return interesting_domains
 
    def ResolveAddress(self, ip):
        """
        Resolve an ip: it searches all the data that match it
        (i.e. all the subnets that contain it) and return
        output text ready for the server response.
        """
        output = ""
        nets = self.GetParentNetworks(ip)
        for net in nets:
            output += FormatMessage("GNet whoisd answering for ip %s" % ip)
            output += FormatOutput ("Range", net)
            net = self.ip_database[net]
            if net.name is not None:
                output += FormatOutput ("Name", net.name)
            if net.owner is not None:
                output += FormatOutput ("Owner", net.owner)
            if net.ns is not None:
                output += FormatOutput ("Nameserver", net.ns)
            if len(net.data.items()) > 0:
                output += FormatMessage("Additional data")
                for k, v in net.data.items():
                    output += FormatOutput (k.capitalize(),v)
        if len(nets) == 0:
            output = FormatMessage("Address %s not found\n" % ip)
        return output
 
    def GetParentNetworks(self, ip):
        """
        Get network in which the given ip is
        """
        interesting_networks = []
        for (address, net) in self.ip_database.items():
            try:
                if IP(ip) in Network(address):
                    interesting_networks.append (address)
            except:
                continue
        return interesting_networks
 
class Net():
    """
    Net is a class representing the network
    contained in the local database and the
    net obtained from other whois server.
    It act as a C struct containing only the
    data and no method.
    """
    def __init__(self, name, owner, ns, data):
        self.name = name
        self.owner = owner
        self.ns = ns
        self.data = data
 
class Domain():
    """
    The same of Net, but with the domain names.
    """
    def __init__(self, name, owner, ns, data):
        self.name, self.owner, self.ns = name, owner, ns
        self.data = data