path: root/inventory/api.py

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'