Files
inventory-app/inventory-app/decorators.py

106 lines
2.9 KiB
Python

import secrets
import string
from datetime import datetime
from functools import wraps
from bson.objectid import ObjectId
from flask import abort, g, make_response, redirect, request, session
from pymongo import MongoClient
import const
db = MongoClient(const.MONGO_URI).get_default_database()
def generate_password(length):
letters = string.ascii_letters + string.digits + string.punctuation
pw = secrets.choice(string.ascii_lowercase)
pw += secrets.choice(string.ascii_uppercase)
pw += secrets.choice(string.digits)
pw += secrets.choice(string.punctuation)
pw += "".join(secrets.choice(letters) for i in range(length - 4))
return "".join(secrets.SystemRandom().sample(pw, length))
def check_login():
token = db.token.find_one({"cookie": request.cookies.get("LOGINLINKCOOKIE", "")})
if not token:
cookie = generate_password(50)
res = db.token.update_one({
"token": request.args.get("token"),
"cookie": {"$exists": False}
}, {
"$set": {
"used": datetime.utcnow(),
"cookie": cookie
}
})
if res.matched_count >= 1:
return (1, cookie)
else:
return (2, cookie)
db.token.update_one({
"_id": token["_id"]
}, {
"$set": {
"last_seen": datetime.utcnow(),
"user_agent": request.headers.get("User-Agent"),
"remote_addr": request.remote_addr
}
})
g.user = db.member.find_one({"_id": ObjectId(token["member_id"])})
g.token = token
return (0, None);
def has_login():
r, _ = check_login()
return r == 0
def login_redirect():
session['target_path'] = request.path
return redirect("/login")
def login_required(f):
@wraps(f)
def decorated_function(*args, **kwargs):
r, cookie = check_login()
if r == 1:
resp = make_response()
resp.set_cookie("LOGINLINKCOOKIE", cookie)
resp.headers["location"] = request.path
return resp, 302
elif r == 2:
return login_redirect()
return f(*args, **kwargs)
return decorated_function
def rget(d, attr):
if not d:
return None
if "." in attr:
pre, post = attr.split(".", 1)
return rget(d.get(pre), post)
else:
return d.get(attr)
def required(path, msg=None, status=403):
def ff(f):
@wraps(f)
def decorated_function(*args, **kwargs):
if not rget(g.user, path):
abort(status, description=msg)
return f(*args, **kwargs)
return decorated_function
return ff
def board_member_required(f):
@wraps(f)
def decorated_function(*args, **kwargs):
_f = required("access.board", "K-SPACE MTÜ board member required")
return _f(f)(*args, **kwargs)
return decorated_function