Friday, May 23, 2014

An example of Auth using Tornado + Motor

Tips from by Yutong Zhao on Google group
 https://groups.google.com/forum/#!topic/python-tornado/l8slxrWHq-c
import tornado.testing
import tornado.web
import tornado.httpserver
import motor
import functools
import pymongo


def authenticate(method):
    @tornado.gen.coroutine
    @functools.wraps(method)
    def wrapper(self, *args, **kwargs):
        key = self.request.headers['Authorization']
        query = {'token': key}
        cursor = self.application.db.users.admins
        result = yield cursor.find_one(query, {'_id': 1})
        if not result:
            self.set_status(401)
            return self.write('Unauthorized')
        else:
            return method(self, *args, **kwargs)
    return wrapper


class GetCurrentUser(tornado.web.RequestHandler):
    @tornado.gen.coroutine
    def get(self):
        current_user = yield self.get_current_user()
        print('Current user:', current_user)

    @tornado.gen.coroutine
    def get_current_user(self):
        key = self.request.headers['Authorization']
        query = {'token': key}
        cursor = self.application.db.users.admins
        result = yield cursor.find_one(query, {'_id': 1})
        if result:
            return result['_id']
        else:
            return None


class GetDecoratedUser(tornado.web.RequestHandler):
    @authenticate
    def get(self):
        print('Success')


class TestCase(tornado.testing.AsyncHTTPTestCase):
    def setUp(self):
        db = pymongo.MongoClient()
        db.users.admins.insert({'_id': 'bob', 'token': '12345'})
        super(TestCase, self).setUp()

    def tearDown(self):
        db = pymongo.MongoClient()
        db.users.admins.remove({})
        super(TestCase, self).tearDown()

    def get_app(self):
        app = tornado.web.Application([
            (r'/current_user', GetCurrentUser),
            (r'/decorator', GetDecoratedUser)
        ])
        app.db = motor.MotorClient()
        return app

    def test_get_current_user(self):
        self.fetch('/current_user', headers={'Authorization': '12345'})
        self.fetch('/current_user', headers={'Authorization': '54321'})

    def test_decorator(self):
        self.fetch('/decorator', headers={'Authorization': '12345'})
        self.fetch('/decorator', headers={'Authorization': '54321'})

No comments:

Post a Comment