1 files changed, 134 insertions, 0 deletions

diff --git a/inventory/api.py b/inventory/api.py
index 0fde577..fa09a68 100644
--- a/inventory/api.py
+++ b/inventory/api.py
@@ -1,10 +1,16 @@
 import datetime
+import time
 
+from bson import ObjectId
 from flask import Blueprint, jsonify, current_app, request, abort, url_for
 from flask_pymongo import PyMongo
+import itsdangerous
+import jwt
 from marshmallow import ValidationError
 import pymongo
 import pytz
+import requests
+import yarl
 
 from .schema import NodeSchema
 
@@ -22,6 +28,11 @@ def on_invalid_request(e):
     return jsonify({'message': e.description or 'Invalid request'}), 400
 
 
+@app.errorhandler(403)
+def on_forbidden(e):
+    return jsonify({'message': e.description or 'Forbidden'}), 403
+
+
 @app.errorhandler(404)
 def on_not_found(e):
     return jsonify({'message': e.description or 'Not found'}), 404
@@ -37,6 +48,27 @@ def on_validation_error(e):
     return jsonify({'message': 'Validation error', 'fields': e.messages}), 400
 
 
+def auth_required(f):
+    def wrapper(*args, **kwargs):
+        authorization = request.headers.get('Authorization')
+        if not authorization:
+            abort(403)
+        auth_type, access_token = authorization.split(None, 1)
+        if auth_type != 'Bearer':
+            abort(403)
+        try:
+            token = jwt.decode(access_token, key=current_app.config['JWT_PUBLIC_KEY'],
+                    audience=current_app.config['OAUTH_CLIENT_ID'])
+        except jwt.InvalidTokenError:
+            abort(403)
+        user_id = ObjectId(token['sub'])
+        user = mongo.db.users.find_one({'_id': user_id})
+        if not user:
+            abort(403)
+        return f(user, *args, **kwargs)
+    return wrapper
+
+
 # Routes
 @app.route('/nodes')
 def root_nodes():
@@ -142,6 +174,108 @@ def find_nodes():
     return jsonify(data)
 
 
+@app.route('/auth/request', methods=['POST'])
+def auth_request():
+    signer = itsdangerous.URLSafeSerializer(current_app.config['SECRET_KEY'])
+    state = signer.dumps(int(time.time()))
+    redirect_uri = request.form.get('redirect_uri')
+    url = (yarl.URL(current_app.config['OAUTH_URL']) / 'authorize').with_query(
+        response_type='code',
+        client_id=current_app.config['OAUTH_CLIENT_ID'],
+        redirect_uri=str(redirect_uri),
+        state=state,
+    )
+    return jsonify(url=str(url))
+
+
+def set_refresh_token_cookie(response, refresh_token='', clear=False):
+    response.set_cookie(key='refresh_token', value=refresh_token, path=url_for('.auth_refresh'),
+            max_age=0 if clear else (86400 * 30), samesite='strict', httponly=True, secure=request.scheme == 'https')
+
+
+@app.route('/auth/response', methods=['POST'])
+def auth_response():
+    code = request.form.get('code')
+    state = request.form.get('state')
+    signer = itsdangerous.URLSafeSerializer(current_app.config['SECRET_KEY'])
+    try:
+        signer.loads(state)
+    except itsdangerous.exc.BadSignature:
+        abort(500, 'Invalid state')
+    redirect_uri = request.form.get('redirect_uri')
+    url = yarl.URL(current_app.config['OAUTH_URL']) / 'token'
+    response = requests.post(str(url), data={
+        'grant_type': 'authorization_code',
+        'code': code,
+        'client_id': current_app.config['OAUTH_CLIENT_ID'],
+        'client_secret': current_app.config['OAUTH_CLIENT_SECRET'],
+        'redirect_uri': redirect_uri,
+    })
+    token = response.json()
+    error = token.get('error')
+    if error:
+        abort(500, error)
+
+    token_data = jwt.decode(token['access_token'], key=current_app.config['JWT_PUBLIC_KEY'],
+            audience=current_app.config['OAUTH_CLIENT_ID'])
+    # We're assuming sub is an ObjectId (this is true for jab)
+    user_id = ObjectId(token_data['sub'])
+    user = mongo.db.users.find_one({'_id': user_id})
+    if not user:
+        url = yarl.URL(current_app.config['OAUTH_URL']).parent / 'api' / 'user'
+        response = requests.get(str(url), headers={
+            'Authorization': 'Bearer {}'.format(token['access_token']),
+        })
+        userdata = response.json()
+        user = {
+            '_id': user_id,
+            'username': userdata['username'],
+        }
+        result = mongo.db.users.insert_one(user)
+        if not result.acknowledged:
+            abort(500, 'Write not acknowledged')
+
+    r = jsonify(access_token=token['access_token'])
+    set_refresh_token_cookie(r, token['refresh_token'])
+    return r
+
+
+@app.route('/auth/refresh', methods=['POST'])
+def auth_refresh():
+    refresh_token = request.cookies.get('refresh_token')
+    if not refresh_token:
+        abort(403)
+    url = yarl.URL(current_app.config['OAUTH_URL']) / 'token'
+    response = requests.post(str(url), data={
+        'grant_type': 'refresh_token',
+        'client_id': current_app.config['OAUTH_CLIENT_ID'],
+        'client_secret': current_app.config['OAUTH_CLIENT_SECRET'],
+        'refresh_token': refresh_token,
+    })
+    if not response.ok:
+        abort(403)
+    data = response.json()
+    r = jsonify(access_token=data['access_token'])
+    set_refresh_token_cookie(r, data['refresh_token'])
+    return r
+
+
+@app.route('/auth/logout', methods=['POST'])
+def auth_logout():
+    r = jsonify()
+    set_refresh_token_cookie(r, clear=True)
+    return r
+
+
+@app.route('/user')
+@auth_required
+def user_info(user):
+    url = yarl.URL(current_app.config['OAUTH_URL']).parent / 'api' / 'user'
+    response = requests.get(str(url), headers={'Authorization': request.headers.get('Authorization')})
+    data = response.json()
+    return jsonify(username=data['username'])
+
+
 @app.route('/')
 def index():
     return 'api'