from flask import current_app as app
from flask_sqlalchemy import SQLAlchemy
from sqlalchemy.sql import text
from werkzeug.security import generate_password_hash, check_password_hash
import jwt
from application.util import decrypt, encrypt
import logging
db = SQLAlchemy()
class User(db.Model):
id = db.Column(db.Integer, primary_key=True)
username = db.Column(db.String(50), unique=True, nullable=False)
password = db.Column(db.String(80), nullable=False)
class File(db.Model):
id = db.Column(db.Integer, primary_key=True)
user_id = db.Column(db.Integer, nullable=False)
title = db.Column(db.String(256), nullable=False)
filename = db.Column(db.String(256), nullable=False)
filepath = db.Column(db.String(256), nullable=False)
def login_user_db(username, password):
user = User.query.filter_by(username=username).first()
if not user:
return False
if check_password_hash(user.password, password):
return jwt.encode({'username': user.username, 'id':user.id}, app.config['SECRET_KEY'], algorithm="HS256")
return False
def register_user_db(username, password):
check_user = User.query.filter_by(username=username.lower()).first()
if not check_user:
hashed_password = generate_password_hash(password)
new_user = User(username=username.lower(), password=hashed_password)
db.session.add(new_user)
db.session.commit()
return True
return False
def list_user_files_db(user_id):
files = File.query.filter_by(user_id=user_id).all()
results = []
for file in files:
file.filepath = decrypt(file.filepath).decode()
file.filename = decrypt(file.filename).decode()
if file.filename is not None and file.filepath is not None:
results.append(file)
return results
def insert_file_db(user_id,filepath,filename,title):
new_file = File(user_id=user_id,
filepath=encrypt(filepath).decode(),
filename=encrypt(filename).decode(),
title=title)
db.session.add(new_file)
db.session.commit()
return True
def delete_file_db(user_id,file_id):
file = File.query.filter_by(id=file_id, user_id=user_id).first()
if file:
filepath = decrypt(file.filepath)
filename = decrypt(file.filename)
if file.user_id == user_id and filepath is not None and filename is not None:
db.session.delete(file)
db.session.commit()
return {"id":file.id,"filepath":filepath.decode(),"filename":filename.decode(),"title":file.title}
return False
return False
def fetch_file_db(user_id,file_id):
try:
file = db.session.execute(text(f"SELECT * FROM File WHERE id = {file_id}")).first()
if file:
filepath = decrypt(file.filepath)
filename = decrypt(file.filename)
if file.user_id == user_id and filepath is not None and filename is not None:
return {"id":file.id,"filepath":filepath.decode(),"filename":filename.decode(),"title":file.title}
return False
except Exception as e:
logging.error(e)
return False