dbwhois.py

#!/usr/bin/python

import psycopg2, os, sys
import json
import pprint

whois_table=os.environ.get("WHOIS_TABLE", "whois")
dsn=os.environ.get("WHOIS_DSN", "")

check_table_sql="select 1 from %s where 1=0"%(whois_table,)
create_sql="create table %s (id bigserial primary key, query text, result jsonb, query_time timestamptz default now())" % (whois_table,)
query_sql="select result from %s where query=%%s"%(whois_table, )
insert_sql="insert into %s (query, result) values (%%s, %%s)" % (whois_table, )


class MyJSONEncoder(json.JSONEncoder):
    dt=__import__("datetime")
    def default(self, obj):
        if isinstance(obj, (self.dt.datetime, self.dt.date)): return obj.isoformat()
        if isinstance(obj, buffer): return str(obj).encode("string_escape")
        raise NotImplementedError("Don't know how to make JSON out of %r" % (type(obj).__name__,))

def make_query(db, query):
    cursor=db.cursor()

    try:
        cursor.execute(check_table_sql)
    except psycopg2.ProgrammingError as e:
        cursor.execute("abort")
        cursor.execute(create_sql)
        cursor.execute("commit")

    cursor.execute(query_sql, (query,))
    ret = cursor.fetchone()
    if ret:
        ret = ret[0]
    else:
        try:
            import pythonwhois
        except ImportError:
            print >>sys.stderr, "Could not import whois, please run something like: pip install pythonwhois"
            raise SystemExit(2)
        ret = pythonwhois.get_whois(query)
        cursor.execute(insert_sql, (query, json.dumps(ret, cls=MyJSONEncoder)))
        cursor.execute("commit")
    return ret

if __name__ == '__main__':
    try:
        db=psycopg2.connect(dsn)
    except psycopg2.OperationalError as e:
        print >>sys.stderr, "Error: cannot connect to database.\nPlease set PG* or WHOIS_DSN environment variables appropriately"
        raise SystemExit(1)
    if len(sys.argv)<2:
        print >>sys.stderr, "Usage: %s {<domain>|<ip>}"%(os.path.basename(sys.argv[0]))
        raise SystemExit(1)
    ret=make_query(db, sys.argv[1])
    pprint.pprint(ret)