dnmgmt/bin/whoisd

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# This file is part of dnmgmt, a number resource management system
# Licensed under GNU General Public License v3 or later
# Written by Sebastian Lohff (seba@someserver.de)

import argparse
import socketserver
import os
import sys
import django
import logging
import select

os.environ.setdefault("DJANGO_SETTINGS_MODULE", "dnmgmt.settings")
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(BASE_DIR)
django.setup()

from whoisdb.helpers import findInDatabase, getWhoisObjectFields

log = logging.getLogger("whoisd")

__VERSION__ = '0.1'

class WhoisHandler(socketserver.BaseRequestHandler):
    header = "% This is the DARKNET database query service.\n" \
            "% The objects should be in something like RPSL format.\n" \
            "%\n" \
            "% The DARKNET database is subject to terms and conditions.\n" \
            "% Mostly these are \"be nice\" and \"don't knowingly break things\".\n" \
            "\n"

    def handle(self):
        self.request.sendall(self.header.encode())
        line = self.request.recv(1024)
        line = line.split(b"\n", 2)[0].strip()
        print("Request object is %s" % line)
        log.info("Request by %s for %s" % (self.request.getpeername()[0], line))
        self.request.send(b"\n")

        objs = findInDatabase(line.decode())
        if len(objs) > 0:
            self.request.sendall(("%% %d result%s\n" % (len(objs), "" if len(objs)==1 else "s")).encode())
            for obj in objs:
                self.sendObject(obj)
        else:
            self.request.sendall(b"%% NOT FOUND\n\n")

    def sendObject(self, obj):
        result = [
            "",
            "%% Object %s (%s)" % (obj, obj.getClassName()),
            ""
        ]

        for field, value in getWhoisObjectFields(obj, False):
            fieldName = field.lower().replace(" ", "-") + ":"
            if not value:
                value = ""

            if hasattr(value, "through"):
                for v in value.all():
                    result.append("%-16s %s" % (fieldName, v))
            else:
                result.append("%-16s %s" % (fieldName, value))
        result.extend(["", ""])

        self.request.sendall("\n".join(result).encode())

class ThreadedTCPServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
    allow_reuse_address = True


def _parser():
    parser = argparse.ArgumentParser(
                 #prog='foo',
                 #description='do some awesome foo',
    )

    parser.add_argument("-p", "--port", default=43, type=int, help="whoisd port")
    #parser.add_argument("-v", "--verbose", default=False, action="store_true", help="Be more verbose")

    parser.add_argument("--version", action="version", version="%(prog)s " + __VERSION__)

    return parser

def main():
    parser = _parser()
    args = parser.parse_args()

    server = ThreadedTCPServer(('', args.port), WhoisHandler)

    while True:
        try:
            server.serve_forever()
        except select.error as e:
            log.exception(e)
        except KeyboardInterrupt:
            log.info("^c hit, quitting.")
            break

    server.server_close()


if __name__ == '__main__':
    main()