From aaf0e4492adb89394ee2acd88bda12413107c5be Mon Sep 17 00:00:00 2001
From: Mia Herkt <mia@0x0.st>
Date: Tue, 13 Dec 2022 23:02:41 +0100
Subject: [PATCH] Record file sizes in db

Moderation interface is going to use this.
---
 fhost.py                                      |  8 ++--
 .../30bfe33aa328_add_file_size_field.py       | 46 +++++++++++++++++++
 2 files changed, 51 insertions(+), 3 deletions(-)
 create mode 100644 migrations/versions/30bfe33aa328_add_file_size_field.py

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')