diff options
Diffstat (limited to 'unmess/api.py')
| -rw-r--r-- | unmess/api.py | 322 |
1 files changed, 322 insertions, 0 deletions
diff --git a/unmess/api.py b/unmess/api.py new file mode 100644 index 0000000..91872f5 --- /dev/null +++ b/unmess/api.py @@ -0,0 +1,322 @@ +import datetime +import functools +import time + +from bson import ObjectId +from flask import Blueprint, jsonify, current_app, request, abort, url_for, g +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 + +# Setup PyMongo +mongo = PyMongo(current_app, tz_aware=True) +mongo.db.nodes.create_index([('fields.value', pymongo.TEXT), ('name', pymongo.TEXT)], name='fields.value_text_name_text') +mongo.db.nodes.create_index([('parent_id', pymongo.ASCENDING)], name='parent_id') +mongo.db.nodes.create_index([('user_id', pymongo.ASCENDING)], name='user_id') + +app = Blueprint('api', __name__) + + +# Error handling +@app.errorhandler(400) +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 + + +@app.errorhandler(500) +def on_internal_error(e): + return jsonify({'message': e.description or 'Invalid request'}), 500 + + +@app.errorhandler(ValidationError) +def on_validation_error(e): + return jsonify({'message': 'Validation error', 'fields': e.messages}), 400 + + +def auth_required(f): + @functools.wraps(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) + g.user = user + return f(*args, **kwargs) + return wrapper + + +def format_node(node): + parents = dict((parent['_id'], parent) for parent in node.pop('parents')) + current = node + # Recursively assign each node their respective parent + while current: + # NOTE: Python assigns from left to right + current['parent'] = current = parents.get(current.get('parent_id')) + + children = dict((child['_id'], child) for child in node.pop('children')) + children_list = node['children'] = [] + for child in children.values(): + # Direct children should be assigned to the top-most list + if child['parent_id'] == node['_id']: + children_list.append(child) + continue + # Otherwise assign to their respective parent + parent = children[child['parent_id']] + parent.setdefault('children', []).append(child) + + return node + +def format_nodes(nodes): + return (format_node(node) for node in nodes) + +# Routes +@app.route('/nodes') +@auth_required +def root_nodes(): + result = mongo.db.nodes.aggregate([ + {'$match': {'parent_id': None, 'user_id': g.user['_id']}}, + { + '$graphLookup': { + 'from': 'nodes', + 'startWith': '$parent_id', + 'connectFromField': 'parent_id', + 'connectToField': '_id', + 'as': 'parents', + }, + }, + { + '$graphLookup': { + 'from': 'nodes', + 'startWith': '$_id', + 'connectFromField': '_id', + 'connectToField': 'parent_id', + 'as': 'children', + 'maxDepth': 1, + }, + }, + ]) + schema = NodeSchema(many=True) + data = schema.dump(format_nodes(result)) + return jsonify(data) + + +@app.route('/nodes', methods=['POST']) +@auth_required +def add_node(): + data = request.json + if data is None or not isinstance(data, dict): + abort(400, 'Payload must be a JSON object') + schema = NodeSchema() + node = schema.load(data) + node['created_at'] = pytz.utc.localize(datetime.datetime.utcnow()) + node['user_id'] = g.user['_id'] + result = mongo.db.nodes.insert_one(node) + if not result.acknowledged: + abort(500, 'Write operation not acknowledged') + node_id = result.inserted_id + return jsonify({'id': str(node_id), 'url': url_for('.node', node_id=node_id)}), 201 + + +@app.route('/nodes/<ObjectId:node_id>') +@auth_required +def node(node_id): + result = mongo.db.nodes.aggregate([ + {'$match': {'_id': node_id, 'user_id': g.user['_id']}}, + { + '$graphLookup': { + 'from': 'nodes', + 'startWith': '$parent_id', + 'connectFromField': 'parent_id', + 'connectToField': '_id', + 'as': 'parents', + }, + }, + { + '$graphLookup': { + 'from': 'nodes', + 'startWith': '$_id', + 'connectFromField': '_id', + 'connectToField': 'parent_id', + 'as': 'children', + 'maxDepth': 1, + }, + }, + ]) + try: + node = result.next() + except StopIteration: + abort(404, 'No node found') + + node = format_node(node) + + schema = NodeSchema() + return jsonify(schema.dump(node)) + + +@app.route('/nodes/<ObjectId:node_id>', methods=['PUT']) +@auth_required +def update_node(node_id): + data = request.json + if data is None or not isinstance(data, dict): + abort(400, 'Payload must be a JSON object') + schema = NodeSchema() + node = schema.load(data) + node['updated_at'] = pytz.utc.localize(datetime.datetime.utcnow()) + result = mongo.db.nodes.update_one({'_id': node_id, 'user_id': g.user['_id']}, {'$set': node}) + if not result.acknowledged: + abort(500, 'Write operation not acknowledged') + return '', 204 + + +@app.route('/nodes/<ObjectId:node_id>', methods=['DELETE']) +@auth_required +def delete_node(node_id): + result = mongo.db.nodes.delete_one({'_id': node_id, 'user_id': g.user['_id']}) + if result.deleted_count == 0: + abort(404, 'No node found') + return jsonify({}), 204 + + +@app.route('/search', methods=['POST']) +@auth_required +def find_nodes(): + if 'q' not in request.form: + abort(400, 'Missing q argument') + schema = NodeSchema(many=True) + data = schema.dump(mongo.db.nodes.find({'$text': {'$search': request.form['q']}, 'user_id': g.user['_id']})) + 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(): + 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' |