nando/nandod

#!/usr/bin/env python

import SocketServer,threading
import sqlite3
import time,os,sys

ALIVE=60*30
TABLE_HEAD="Host|IP|Age(d h:m:s)|Description\n"

def setup_options():
    ''' Setup the command line options '''
    from argparse import ArgumentParser

    parser=ArgumentParser(description="Alive notifier.")
    parser.add_argument("--host",action='store', dest='HOST',default='0.0.0.0',type=str,
                      help="Bind to address: %(default)s")
    parser.add_argument("--port", action='store', dest='PORT', default=13370,type=int,
                      help="Bind to port: %(default)s")
    parser.add_argument("--html-port", action='store', dest='WEBPORT', default=13370,type=int,
                      help="Bind www server to port. 0 to disable, default: %(default)s")
    parser.add_argument("--db",action='store', dest='DB',default='/tmp/nando.sqlite',type=str,
                      help="Sqlite file for database: %(default)s")
    parser.add_argument("--quiet","-q",action='store_true', dest='quiet',default=False,
                      help="Quiet operation: %(default)s")
    options=parser.parse_args()
    return options


class UDPHandler(SocketServer.BaseRequestHandler):
    def handle(self):
        data = self.request[0].strip()
        socket = self.request[1]
        if not opts.quiet:
            print "{0} wrote: {1}".format(self.client_address[0],data)
        if data=="list":
            reply=DB.list_all()
        elif data=="alive":
            reply=DB.list_alive()
        elif data=="lost":
            reply=DB.list_lost()
        else:
            DB.update(data)
            reply="OK"
        socket.sendto(reply, self.client_address)
    
class TCPHandler(SocketServer.BaseRequestHandler):
    def handle(self):
        if not opts.quiet:
            print "TCP request: {0}".format(self.client_address[0])
        self.request.send('HTTP/1.0 200 OK\r\n')
        self.request.send("Content-Type: text/html\r\n\r\n")
        self.request.sendall(HTMLDB.HTML_list())
        self.request.close()

class DataBase:
    def __init__(self,DB):
        self.DBfile=DB
        self.conn=None
        self.db=None
        if not os.path.exists(self.DBfile):
            self.createDB()
        self.conn_init()
    
    def conn_init(self):
        self.conn=sqlite3.connect(self.DBfile)
        self.db=self.conn.cursor()
        self.conn.text_factory=str

    def conn_end(self):
        self.conn.commit()
        #self.conn.close()

    def createDB(self):
        self.conn_init()
        self.db.execute('CREATE TABLE alive (id TEXT PRIMARY KEY,\
                    ip TEXT,\
                    desc TEXT,\
                    date INTEGER)')
        self.conn_end()
        
    def update(self,id):
        try:
            host,ip,desc=id.split("|",3)
            desc=''.join([ c for c in desc if c not in ['\n','\r','|'] ])
            if len(ip.split("."))!=4:
                return 1
        except:
            return
        self.db=self.conn.cursor()
        if desc is "":
            self.db.execute("INSERT OR REPLACE INTO alive(id,ip,date) \
               VALUES(?,?,?)",(host,ip,int(time.time())))
        else:
            self.db.execute("INSERT OR REPLACE INTO alive(id,ip,date,desc) \
               VALUES(?,?,?,?)",(host,ip,int(time.time()),desc))
        self.conn_end()
    
    def list_all(self):
        self.db=self.conn.cursor()
        self.db.execute("SELECT id,ip,date,desc FROM alive ORDER BY id")
        alive=[]
        lost=[]
        for row in self.db:
            age=int(time.time())-row[2]
            if age<ALIVE:
                alive.append("{0}|{1}|{2}|{3}".format(row[0],row[1], humanize_time(age),row[3]))
            else:
                lost.append( "{0}|{1}|{2}|{3}".format(row[0],row[1], humanize_time(age),row[3]))
        return TABLE_HEAD+"\n".join(alive)+"\n---|---|---|---\n"+"\n".join(lost)
    
    def list_alive(self):
        self.db=self.conn.cursor()
        self.db.execute("SELECT id,ip,date,desc FROM alive ORDER BY date DESC")
        msg=[]
        for row in self.db:
            age=int(time.time())-row[2]
            if age<ALIVE:
                msg.append("{0}|{1}|{2}|{3}".format(row[0],row[1], humanize_time(age),row[3])) 
        return TABLE_HEAD+"\n".join(msg)
        
    def list_lost(self):
        self.db=self.conn.cursor()
        self.db.execute("SELECT id,ip,date,desc FROM alive ORDER BY date")
        msg=[]
        for row in self.db:
            age=int(time.time())-row[2]
            if age>=ALIVE:
                msg.append("{0}|{1}|{2}|{3}".format(row[0],row[1], humanize_time(age),row[3])) 
        return TABLE_HEAD+"\n".join(msg)

    def HTML_list(self):
        msg=[]
        msg.append('''<html>
<head>
<style>
body {
    font-family: monospace;
}
table {
    background-color: gray;
}
td, th {
    text-align: left;
    background-color: lightgray;
    padding: 0.5ex;
}
.right {
   text-align: right;
}
</style>
</head>
<body><h1>Alive</h1><table>''')
        self.db=self.conn.cursor()
        self.db.execute("SELECT id,ip,date,desc FROM alive ORDER BY id")
        msg.append('<tr><th>{0}</th><th>{1}</th><th class="right">{2}</th><th>{3}</th></tr>'.format('Host','IP','Age[d h:m:s]','Description'))
        for row in self.db:
            age=int(time.time())-row[2]
            if age <ALIVE:
                msg.append('<tr><td>{0}</td><td>{1}</td><td class="right">{2}</td><td>{3}</td></tr>'.format(row[0],row[1],humanize_time(age),row[3]))
        msg.append('</table><h1>Lost</h1><table>')
        msg.append('<tr><th>{0}</th><th>{1}</th><th class="right">{2}</th><th>{3}</th></tr>'.format('Host','IP','Age[d h:m:s]','Description'))
        self.db.execute("SELECT id,ip,date,desc FROM alive ORDER BY id")
        for row in self.db:
            age=int(time.time())-row[2]
            if age >=ALIVE:
                msg.append('<tr><td>{0}</td><td>{1}</td><td class="right">{2}</td><td>{3}</td></tr>'.format(row[0],row[1],humanize_time(age),row[3]))
        msg.append('</table></body></html>')

        return "\n\r".join(msg)

def UDPserve():
    global DB
    global server
    DB=DataBase(opts.DB)
    server = SocketServer.UDPServer((opts.HOST, opts.PORT), UDPHandler)
    server.serve_forever()
def TCPserve():
    global HTMLDB
    global HTMLserver
    time.sleep(2)
    HTMLDB=DataBase(opts.DB)
    HTMLserver = SocketServer.TCPServer((opts.HOST, opts.PORT), TCPHandler)
    HTMLserver.serve_forever()

def humanize_time(secs):
    mins, secs = divmod(secs, 60)
    hours, mins = divmod(mins, 60)
    if hours < 24:
        return '%02d:%02d:%02d' % (hours, mins, secs)
    else:
        days, hours = divmod(hours, 24)
        return '%dd %02d:%02d:%02d' % (days, hours, mins, secs)

if __name__ == "__main__":
    opts=setup_options()
    if not opts.quiet:
        print("Starting NandoD UDP:{0}:{1}, TCP:{0}:{2}".format(
           opts.HOST,
           opts.PORT,
           opts.WEBPORT
        ))
    UDP=threading.Thread(target=UDPserve)
    UDP.daemon=True
    UDP.start()
    if opts.WEBPORT > 0:
        TCP=threading.Thread(target=TCPserve)
        TCP.daemon=True
        TCP.start()
    while True:
        try:
            time.sleep(1)
        except KeyboardInterrupt:
            print("Exiting..")
            server.shutdown()
            if opts.WEBPORT > 0:
                HTMLserver.shutdown()
            sys.exit(0)