OAuth2 Implementation

Question

I am trying to implement an OAuth2 service. I am not planning to use it anytime soon (like will write unit tests). This is my fun project I do to practice. I have used the python_jwt library to create and validate signatures.

OAuth Request handler

"""
 Copyright (C) 2016 Maruf Maniruzzaman
 Website: https://github.com/kuasha/cosmos
 Author: Maruf Maniruzzaman
 License :: OSI Approved :: MIT License

 This service handler follows the style found in Azure (https://msdn.microsoft.com/en-us/library/azure/dn645542.aspx)
 Also for more information look at The OAuth 2.0 Authorization Framework (http://tools.ietf.org/html/rfc6749)
"""

import ast
import base64
import json

import datetime
import uuid

import tornado.web
from tornado import gen
import logging

from cosmos.rbac.object import COSMOS_USERS_OBJECT_NAME
from cosmos.auth.oauth2 import get_token
from cosmos.rbac.object import SYSTEM_USER

try:
    import urlparse  # py2
    from urllib import urlencode
    import urllib as urllib_parse
except ImportError:
    import urllib.parse as urlparse  # py3
    from urllib.parse import urlencode
    import urllib.parse as urllib_parse

from cosmos.service.requesthandler import RequestHandler

OAUTH2_REQUESTS_OBJECT_NAME = "cosmos.auth.oauth2.requests"
OAUTH2_TOKENS_OBJECT_NAME = "cosmos.auth.oauth2.tokens"
OAUTH2_CODE_STATUS_OBJECT_NAME = "cosmos.auth.oauth2.codestatus"
OAUTH2_MAX_CODE_RESPONSE_LENGTH = 200

class OAuth2RequestException(Exception):
    pass

#TODO: security check for entire class

class OAuth2ServiceHandler(RequestHandler):

    @gen.coroutine
    def get(self, tenant_id, function):
        try:
            req_id = None
            auth_request = None
            serve_request = self.get_argument("serve_request", None)
            if not (serve_request):
                auth_request = self._collect_auth_request_parameters(tenant_id, function)

                req_id = yield self.insert_auth_request(auth_request)

                if not req_id:
                    logging.critical("Could not save OAuth2 request to database.")
                    raise tornado.web.HTTPError(500, "Server error")

                logging.debug("Request saved in db {0}".format(str(req_id)))
            else:
                req_id = self.get_argument("serve_request")
                auth_request = yield self.load_auth_request(req_id)
                if not auth_request:
                    self.write("Could not find auth request. Correlation id {}".format(req_id))
                    self.finish()
                    return

            user = self.get_current_user()

            if not user:
                url = self.request.uri
                parts = list(urlparse.urlparse(url))
                query = {"serve_request": str(req_id)}
                parts[4] = urlencode(query)

                redirect_url = urlparse.urlunparse(parts)

                self.initiate_login(redirect_url)

            if function == "authorize":
                yield self._do_authorize(user, auth_request)
                return
            elif function == "token":
                yield self._do_token(user, auth_request)
                return
            else:
                raise tornado.web.HTTPError(404, "Not found")
        except OAuth2RequestException as re:
            self.clear()
            self.set_status(400)
            self.finish("<html><body>400 Bad Request<br/>{}</body></html>".format(str(re)))

    @gen.coroutine
    def _do_authorize(self, user, params):
        response_type = params.get("response_type", None)
        if not "code" == response_type:
            raise OAuth2RequestException("Response type must be code")

        client_id = params.get("client_id", None)
        redirect_uri = params.get("redirect_uri", None)
        state = params.get("state", None)
        resource = params.get("resource", None)

        oauth2_public_key_handler = self.settings.get("oauth2_public_key_handler")

        code_attributes = ["user_id", "client_id", "resource", "iat", "code_status_id"]

        response = {}
        response.update({"user_id": user["_id"]})
        response["client_id"] = client_id
        response["resource"] = resource
        response["iat"] = int(datetime.datetime.utcnow().timestamp())

        code_status_id = yield self.insert_code_status(response)
        if not code_status_id:
            logging.critical("Could not save [OAuth2 Code] to database.")
            raise tornado.web.HTTPError(500, "Server error")

        response["code_status_id"] = str(code_status_id)

        code_response = { k: response.get(k) for k in code_attributes }

        str_response = json.dumps(code_response)
        bytes_resp = str_response.encode()

        if(len(bytes_resp) > OAUTH2_MAX_CODE_RESPONSE_LENGTH):
            raise tornado.web.HTTPError(414, "Request-URI Too Long")

        enc_response = oauth2_public_key_handler.encrypt(bytes_resp, 32)

        b64_enc_response_bytes = base64.urlsafe_b64encode(enc_response[0])
        b64_enc_response = b64_enc_response_bytes.decode()

        url = redirect_uri
        session_state = str(user.get("session_id", str(uuid.uuid4())))

        params = {response_type: b64_enc_response, "session_state": session_state}
        if state:
            params["state"] = state

        parts = list(urlparse.urlparse(url))
        query = dict(urlparse.parse_qsl(parts[4]))
        query.update(params)
        parts[4] = urlencode(query)

        redirect_url = urlparse.urlunparse(parts)

        self.redirect(redirect_url)

    @gen.coroutine
    def _do_token(self, requesting_user, params):
        code = params.get("code", None)
        grant_type = params.get("grant_type", None)
        redirect_uri = params.get("redirect_uri", None)
        resource = params.get("resource", None)
        state = params.get("state", None)
        client_id = params.get("client_id", None)
        client_secret = params.get("client_secret", None)
        refresh_token = params.get("refresh_token", None)
        tid = params.get("tenant_id", None)
        token_type = "bearer"

        # if grant_type == code then code must not be empty
        # if grant_type == refresh_token then refresh_token must not be empty
        if grant_type == "code":
            if not code:
                raise OAuth2RequestException("Code must not be empty when grant_type is code")
        elif grant_type == "refresh_token":
            if not refresh_token:
                raise OAuth2RequestException("Code must not be empty when grant_type is refresh_token")
        else:
            raise OAuth2RequestException("Value of grant_type must be either code or refresh_token")

        request_host = params.get("request_host")
        oauth2_token_issuer = self.settings.get("oauth2_token_issuer", request_host)
        oauth2_token_expiry_seconds = self.settings.get("oauth2_token_expiry_seconds")
        exp = datetime.timedelta(seconds=oauth2_token_expiry_seconds)

        enc_code = base64.urlsafe_b64decode(code.encode())
        oauth2_private_key_handler = self.settings.get("oauth2_private_key_handler")
        code_json_bytes = oauth2_private_key_handler.decrypt(enc_code)
        code_json = code_json_bytes.decode()
        code_dict = json.loads(code_json)

        user_id = code_dict.get("user_id")
        code_iat = code_dict.get("iat")

        if not user_id or not code_iat:
            raise OAuth2RequestException("Invalid code")

        user = yield self.load_user(user_id)
        if not user:
            raise OAuth2RequestException("Invalid code")

        #TODO: allow use of secret - otherwise its not secure

        #Make sure the code was not used earlier - otherwise its not secure
        code_status_id = code_dict.get("code_status_id")
        if not code_status_id:
            raise OAuth2RequestException("Invalid code")

        code_status = yield self.load_code_status(code_status_id)
        if not code_status:
            raise OAuth2RequestException("Invalid code. Possibly server error.")

        if code_status.get("used_at"):
            code_status["duplicate_attempt"] = True
            yield self.update_code_status(code_status)
            raise OAuth2RequestException("Code already used")

        code_status["used_at"] = int(datetime.datetime.utcnow().timestamp())
        saved = yield self.update_code_status(code_status)
        if not saved:
            logging.error("Could not save code status Id={}".format(code_status_id))

        resource = code_dict.get("resource")
        oauth2_private_key_pem = self.get_private_key_pem()

        response = get_token(aud=client_id,
                                     exp=exp,
                                     family_name=user.get("family_name"),
                                     given_name=user.get("given_name"),
                                     iat=str(int(datetime.datetime.utcnow().timestamp())),
                                     iss=oauth2_token_issuer,
                                     nbf=str(int(datetime.datetime.utcnow().timestamp())),
                                     oid=str(user.get("_id")),
                                     sub=str(user.get("_id")),
                                     tid=tid,
                                     unique_name=user.get("username"),
                                     upn=user.get("username"),
                                     service_private_pem=oauth2_private_key_pem)
        url = redirect_uri
        session_state = user.get("session_id", str(uuid.uuid4()))
        response_type = "access_token"
        params = {response_type: response, "session_state": session_state, "token_type": token_type, "resource": resource}

        if state:
            params["state"] = state

        parts = list(urlparse.urlparse(url))
        query = dict(urlparse.parse_qsl(parts[4]))
        query.update(params)
        parts[4] = urlencode(query)

        redirect_url = urlparse.urlunparse(parts)

        self.redirect(redirect_url)

    def _collect_auth_request_parameters(self, tenant_id, function):
        params = {k: self.get_argument(k) for k in self.request.arguments}
        params["request_protocol"] = self.request.protocol
        params["request_host"] = self.request.host
        params["request_uri"] = self.request.uri
        params["tenant_id"] = tenant_id
        params["function"] = function

        return params

    @gen.coroutine
    def insert_auth_request(self, params):
        obj_serv = self.settings['object_service']
        promise = obj_serv.insert(SYSTEM_USER, OAUTH2_REQUESTS_OBJECT_NAME, params)
        req_id = yield promise
        return req_id

    @gen.coroutine
    def load_auth_request(self, req_id):
        obj_serv = self.settings['object_service']
        cursor = obj_serv.load(SYSTEM_USER, OAUTH2_REQUESTS_OBJECT_NAME, req_id, [])
        request = yield cursor
        return request

    @gen.coroutine
    def insert_code_status(self, response):
        obj_serv = self.settings['object_service']
        promise = obj_serv.insert(SYSTEM_USER, OAUTH2_CODE_STATUS_OBJECT_NAME, response)
        code_status_id = yield promise
        return code_status_id

    @gen.coroutine
    def update_code_status(self, code_status):
        obj_serv = self.settings['object_service']
        promise = obj_serv.save(SYSTEM_USER, OAUTH2_CODE_STATUS_OBJECT_NAME, code_status)
        code_status_id = yield promise
        return code_status_id

    @gen.coroutine
    def load_code_status(self, code_status_id):
        obj_serv = self.settings['object_service']
        promise = obj_serv.load(SYSTEM_USER, OAUTH2_CODE_STATUS_OBJECT_NAME, code_status_id, [])
        code_status = yield promise
        return code_status

    @gen.coroutine
    def load_user(self, user_id):
        obj_serv = self.settings['object_service']
        cursor = obj_serv.load(SYSTEM_USER, COSMOS_USERS_OBJECT_NAME, user_id, [])
        user = yield cursor
        return user

    def get_private_key_pem(self):
        return self.settings.get("oauth2_private_key_pem")

Token generation and validation code

def get_token(**kwargs):
    service_private_pem = kwargs.get("service_private_pem")
    if not service_private_pem:
        raise ValueError("service_private_pem is not defined")

    aud = kwargs.get("aud", None)
    exp = kwargs.get("exp", None)
    family_name = kwargs.get("family_name", None)
    given_name = kwargs.get("given_name", None)

    token_payload = {
        "aud": aud,
        "exp": exp,
        "family_name": family_name,
        "given_name": given_name,
        "ver": "1.0"
    }

    iss = kwargs.get("iss", None)
    if iss:
        token_payload["iss"] = iss

    nbf = kwargs.get("nbf", None)
    if nbf:
        token_payload["nbf"] = nbf

    oid = kwargs.get("oid", None)
    if oid:
        token_payload["oid"] = oid

    sub = kwargs.get("sub", None)
    if sub:
        token_payload["sub"] = sub

    tid = kwargs.get("tid", None)
    if tid:
        token_payload["tid"] = tid

    unique_name = kwargs.get("unique_name", None)
    if unique_name:
        token_payload["unique_name"] = unique_name

    upn = kwargs.get("upn", None)
    if upn:
        token_payload["upn"] = upn

    priv_key = RSA.importKey(service_private_pem)
    token = jwt.generate_jwt(token_payload, priv_key, 'RS256', exp)

    return token


def verify_token(token, public_key_pem, alg_list):
    try:
        header, claims = jwt.verify_jwt(token, public_key_pem, alg_list)
        return header, claims
    except Exception as ex:
        raise ValidationError(ex)

Welcome to Code Review! "Draft" is OK, as long can claim that it "works". — 200_success, Commented Apr 6, 2016 at 16:21
Its a piece of code that deals with security - so very hard to tell people it is stable for production without deep tests and some evidence in production that it works. Its new code, other than that it I can say it works. — Maruf Maniruzzaman, Commented Apr 15, 2016 at 6:28
I have rolled back the last edit. Please see What to do when someone answers. — Mast, Commented May 12, 2016 at 13:29

ferada · Accepted Answer · 2016-04-15 23:08:25Z

Second part first:

In verify_token the return values are assigned - only to be immediately returned. I get that it might be nice in a debugger(?) but there's really no reason for this - just return the values. I'd possibly create a function that wraps other functions in such an exception handler, but that's probably overkill here.
dict.get already defaults to returning None - I don't see a particular reason why you'd specify that everywhere. Defaults are really useful to remove that clutter from your code.
I'd also inline a lot of those variables to make the methods shorter (i.e. aud, exp and so on), or create a similar loop as described in the next point.
The repeated pattern with x = kwargs.get("x"); if x: ...["x"] = x should also be refactored to remove repetition, e.g. by using a loop instead:
```
for name in ["iss", "nbf", "oid", "sub", "tid", "unique_name", "upn"]:
    x = kwargs.get(name)
    if x:
        token_payload[name] = x
```
I'd also reorder code fragments that belong together. At the moment the reader is jumping from service_private_pem to the initialisation of token_payload to only go back to the first variable in the second to last line, while token_payload is used in the last line. I'd move the RSA.importKey after the first block, then handle the token_payload and all associated things after, that way the order of operation is clearly distinguished.

That leaves the following (could probably still be improved with more factoring of code, but hey, looks compact enough to me):

def get_token(**kwargs):
    service_private_pem = kwargs.get("service_private_pem")
    if not service_private_pem:
        raise ValueError("service_private_pem is not defined")
    priv_key = RSA.importKey(service_private_pem)

    token_payload = {"ver": "1.0"}

    for name in ["aud", "exp", "family_name", "given_name"]:
        token_payload[name] = kwargs.get(name)

    for name in ["iss", "nbf", "oid", "sub", "tid", "unique_name", "upn"]:
        x = kwargs.get(name)
        if x:
            token_payload[name] = x

    return jwt.generate_jwt(token_payload, priv_key, 'RS256', exp)

Now for the first part. Well that's a lot of code. Okay, so small stuff:

Empty returns with no other afterwards don't do anything, so I'd leave them out.
logging already supports arguments using %s and friends - the calls can thus be written like logging.debug("... %s", req_id) etc.

the six coroutine methods look awfully similar. I'm not too familiar with generators or the framework, so this could be wrong:

def _request_helper(self, *args):
    obj_serv = self.settings["object_service"]
    promise = obj_serv.insert(SYSTEM_USER, OAUTH2_REQUESTS_OBJECT_NAME, *args)
    req_id = yield promise
    return req_id

This helper would then reused from the individual methods (though at that point you could as well use that method directly?):

@gen.coroutine
def insert_auth_request(self, params):
    return self._request_helper(params)

In _collect_auth_request_parameters the update of the dictionary could be a bit more compact:

def _collect_auth_request_parameters(self, tenant_id, function):
    params = {k: self.get_argument(k) for k in self.request.arguments}
    params.update({
        "request_protocol": self.request.protocol,
        "request_host": self.request.host,
        "request_uri": self.request.uri,
        "tenant_id": tenant_id,
        "function": function
    })
    return params

Having these giant try/except blocks is just increasing the indentation. I'd suggest that the inner part is either compacted so much as to fit into one screen, or moved into its own method.
Similarly having the return cases early can remove the indentation from the else branch, which makes the flow also a bit more clear. Initialising variables to None is also mostly not that useful.
_do_authorize is just convoluted and can be cleaned up a lot by removing all the code moving values around between the dictionaries. And again, sprinkling assignments out over the method like this makes for a great treasure hunt, but doesn't help understanding the code too much. Reassigning params also tripped me up during rewriting.
_do_token is similarly just too much - perhaps no wonder considering it looks mostly like copy & paste from the method before. IMO that's the biggest mistake as it requires a lot of attention when trying to deduplicate it again. It's better to spend the effort creating a "DRY" implementation from the start.
The exceptions are mostly similar and can be moved into methods.

Actually I got confused when trying to clean up _do_authorize and _do_token, so I'll give up on that for now. I can imagine that other readers might have a similar problem. I strongly suggest to aggressively factor out common functions and creating more methods, helper functions, whatever to help keep the line count and number of variables to read down.

Stack Exchange Network

OAuth2 Implementation

OAuth Request handler

Token generation and validation code

1 Answer 1

Your Answer

Not the answer you're looking for? Browse other questions tagged
python
oauth
jwt
or ask your own question.

Hot Network Questions

OAuth2 Implementation

OAuth Request handler

Token generation and validation code

1 Answer 1

Your Answer

Sign up or log in

Post as a guest

Not the answer you're looking for? Browse other questions tagged pythonoauthjwt or ask your own question.

Related

Hot Network Questions

Not the answer you're looking for? Browse other questions tagged
python
oauth
jwt
or ask your own question.