diff --git a/fhost.py b/fhost.py index 641729d..eb30a9d 100755 --- a/fhost.py +++ b/fhost.py @@ -140,6 +140,7 @@ class File(db.Model): mgmt_token = db.Column(db.String) secret = db.Column(db.String) last_vscan = db.Column(db.DateTime) + size = db.Column(db.BigInteger) def __init__(self, sha256, ext, mime, addr, expiration, mgmt_token): self.sha256 = sha256 @@ -288,6 +289,8 @@ class File(db.Model): with open(p, "wb") as of: of.write(data) + f.size = len(data) + if not f.nsfw_score and app.config["NSFW_DETECT"]: f.nsfw_score = nsfw.detect(p) @@ -416,8 +419,7 @@ def manage_file(f): except ValueError: abort(400) - fsize = f.getpath().stat().st_size - f.expiration = File.get_expiration(requested_expiration, fsize) + f.expiration = File.get_expiration(requested_expiration, f.size) db.session.commit() return "", 202 @@ -455,7 +457,7 @@ def get(path, secret=None): if app.config["FHOST_USE_X_ACCEL_REDIRECT"]: response = make_response() response.headers["Content-Type"] = f.mime - response.headers["Content-Length"] = fpath.stat().st_size + response.headers["Content-Length"] = f.size response.headers["X-Accel-Redirect"] = "/" + str(fpath) else: response = send_from_directory(app.config["FHOST_STORAGE_PATH"], f.sha256, mimetype = f.mime) diff --git a/migrations/versions/30bfe33aa328_add_file_size_field.py b/migrations/versions/30bfe33aa328_add_file_size_field.py new file mode 100644 index 0000000..e6ac279 --- /dev/null +++ b/migrations/versions/30bfe33aa328_add_file_size_field.py @@ -0,0 +1,46 @@ +"""add file size field + +Revision ID: 30bfe33aa328 +Revises: 5cee97aab219 +Create Date: 2022-12-13 22:32:12.242394 + +""" + +# revision identifiers, used by Alembic. +revision = '30bfe33aa328' +down_revision = '5cee97aab219' + +from alembic import op +import sqlalchemy as sa +from sqlalchemy.ext.automap import automap_base +from sqlalchemy.orm import Session +from flask import current_app +from pathlib import Path + +Base = automap_base() + +def upgrade(): + op.add_column('file', sa.Column('size', sa.BigInteger(), nullable=True)) + bind = op.get_bind() + Base.prepare(autoload_with=bind) + File = Base.classes.file + session = Session(bind=bind) + + storage = Path(current_app.config["FHOST_STORAGE_PATH"]) + + updates = [] + files = session.scalars(sa.select(File).where(sa.not_(File.removed))) + for f in files: + p = storage / f.sha256 + if p.is_file(): + updates.append({ + "id" : f.id, + "size" : p.stat().st_size + }) + + session.bulk_update_mappings(File, updates) + session.commit() + + +def downgrade(): + op.drop_column('file', 'size')