rbmd_web_panel/rbmd.py

#!/usr/bin/env python

import bcrypt
import json
from kazoo.client import KazooClient
from kazoo.recipe.watchers import DataWatch
import logging
import requests
import sqlite3
import tornado.escape
import tornado.ioloop
import tornado.options
import tornado.web
import tornado.websocket
import threading
import os.path
import uuid


from tornado.options import define, options

define("port", default=8000, help="run on the given port", type=int)
define("sqlite_db", default=os.path.join(os.path.dirname(__file__), "auth.db"), help="User DB")

class Application(tornado.web.Application):
    def __init__(self):
        handlers = [
            (r"/", MainHandler),
            (r"/socket", SocketHandler),
            (r"/mount", MountHandler),
            (r"/unmount", UnmountHandler),
            (r"/resolve", ResolveHandler),
            (r"/status", StatusHandler),
            (r"/user", User),
            (r"/login", Auth),
            
        ]
        settings = dict(
            cookie_secret="=&r854^9nk7ys49@m7a5eu(g&jn8pytk6f%@quumabt*x5e*)i",
            template_path=os.path.join(os.path.dirname(__file__), "templates/rbmd"),
            static_path=os.path.join(os.path.dirname(__file__), "static/rbmd"),
            xsrf_cookies=True,
            websocket_ping_interval = 59,
            login_url="/login",
        )
        super(Application, self).__init__(handlers, **settings)


class MainHandler(tornado.web.RequestHandler):
    @tornado.web.authenticated
    def get(self):
        user_id = self.get_secure_cookie("user")
        ws = config("ws")
        try: metrics = json.loads(action('metrics', 'get'))
        except ValueError: metrics = {}
        my_status = zk_fetch("/rbmd/log/health") #action('status', 'get')
        dct = {'ws': ws, 'metrics': metrics, 'status': my_status}
        self.render("index.html", **dct)

    def get_current_user(self):
        return self.get_secure_cookie("user")


class Auth(tornado.web.RequestHandler):
    def get(self):
        self.render("login.html", error=None)        

    def post(self):
        con = sqlite3.connect(options.sqlite_db)
        cur = con.cursor()
        cur.execute("SELECT * FROM users WHERE name = :name", {"name": self.get_argument("username")})
        #cur.execute("SELECT * FROM users WHERE name = ?", (self.get_argument("username"),))
        user = cur.fetchone()
        cur.close()
        con.close()
        if not user:
            self.render("login.html", error="user not found")
            return
        hashed_password = bcrypt.hashpw(tornado.escape.utf8(self.get_argument("password")), tornado.escape.utf8(user[2]))
        if hashed_password == user[2]:
            self.set_secure_cookie("user", str(user[0]))
            self.redirect("/")
        else:
            self.render("login.html", error="incorrect password")


#create new user
class User(tornado.web.RequestHandler):
    def get(self):
        con = sqlite3.connect(options.sqlite_db)
        cur = con.cursor()
        data = self.request.arguments
        hashed_password = bcrypt.hashpw(tornado.escape.utf8(self.get_argument("password")), bcrypt.gensalt())
        cur.execute(
            "INSERT INTO users (name, password) VALUES (?, ?)",
             (self.get_argument("name"), hashed_password)
        )
        con.commit()
        if hashed_password == bcrypt.hashpw(tornado.escape.utf8(self.get_argument("password")), hashed_password):
            logging.info('OK')
        cur.close()
        con.close()
        self.write('done')


class MountHandler(tornado.web.RequestHandler):
    def post(self):
        if self.get_secure_cookie("user"):
            data = {k: v[0] for k, v in self.request.arguments.items()}
            res = action('mount', 'post', json.dumps(data))


class UnmountHandler(tornado.web.RequestHandler):
    def get(self):
        #data = self.request.arguments
        data = {k: v[0] for k, v in self.request.arguments.items()}
        res = action('umount', 'post', json.dumps(data))


class ResolveHandler(tornado.web.RequestHandler):
    #my_status = action('status', 'get')
    def get(self):
        #data = self.request.arguments
        data = {k: v[0] for k, v in self.request.arguments.items()}
        res = action('resolve', 'post', json.dumps(data))

class StatusHandler(tornado.web.RequestHandler):
    pass

class SocketHandler(tornado.websocket.WebSocketHandler):
    waiters = set()
    quorum = dict()

    def get_compression_options(self):
        # Non-None enables compression with default options.
        return {}

    def check_origin(self, origin):
       return True

    def open(self):
        SocketHandler.waiters.add(self)
        SocketHandler.send_updates(SocketHandler.quorum)
    
    def on_close(self):
        SocketHandler.waiters.remove(self)

    @classmethod
    def send_updates(cls, data, *stat):
        SocketHandler.quorum = data
        for waiter in cls.waiters:
            try:
                waiter.write_message(SocketHandler.quorum)
            except:
                logging.error("Error sending message", exc_info=True)

    def on_message(self, message):     
        logging.info(logging.info(message))

def config(point):
    with open('conf.json') as conf:
        path = json.load(conf)[point]
    return path

def action(name, method, data=None):
    with open('conf.json') as conf:
        url = json.load(conf)["api"] + '/' + name
        logging.info(url)
        logging.info(data)
        if method == 'get':
            try:
                res = requests.get(url).content
            except:
                res = 'connection can\'t be established'
        elif method == 'post':
            try:
                res = requests.post(url, data).content
            except:
                res = 'connection can\'t be established'
    return res


def zk_handler():
    logging.basicConfig()
    zk = KazooClient(hosts=config("zookeeper"))
    zk.start()
    t1 = threading.Thread(target=DataWatch, args=(zk, "/rbmd/log/quorum"), kwargs=dict(func=SocketHandler.send_updates))
    t1.setDaemon(True)
    t1.start()



def zk_fetch(path):
    zk = KazooClient(hosts=config("zookeeper"))
    zk.start()
    data = zk.get(path)
    zk.stop()


def main():
    tornado.options.parse_command_line()
    zk_handler()
    app = Application()
    app.listen(options.port)
    tornado.ioloop.IOLoop.current().start()


if __name__ == "__main__":
    main()