Move storage to classes, support double suffixes

This commit is contained in:
Mia Herkt 2020-12-29 14:19:41 +01:00
parent cd083a7f83
commit 73045dc5e5
No known key found for this signature in database
GPG key ID: 72E154B8622EC191

181
fhost.py
View file

@ -26,7 +26,7 @@ from jinja2.exceptions import *
from hashlib import sha256 from hashlib import sha256
from magic import Magic from magic import Magic
from mimetypes import guess_extension from mimetypes import guess_extension
import os, sys import sys
import requests import requests
from short_url import UrlEncoder from short_url import UrlEncoder
from validators import url as url_valid from validators import url as url_valid
@ -66,10 +66,11 @@ app.config.update(
URL_ALPHABET = "DEQhd2uFteibPwq0SWBInTpA_jcZL5GKz3YCR14Ulk87Jors9vNHgfaOmMXy6Vx-", URL_ALPHABET = "DEQhd2uFteibPwq0SWBInTpA_jcZL5GKz3YCR14Ulk87Jors9vNHgfaOmMXy6Vx-",
) )
app.config.from_pyfile("config.py") if not app.config["TESTING"]:
app.config.from_pyfile("config.py")
if app.config["DEBUG"]: if app.config["DEBUG"]:
app.config["FHOST_USE_X_ACCEL_REDIRECT"] = False app.config["FHOST_USE_X_ACCEL_REDIRECT"] = False
if app.config["NSFW_DETECT"]: if app.config["NSFW_DETECT"]:
from nsfw_detect import NSFWDetector from nsfw_detect import NSFWDetector
@ -82,9 +83,6 @@ except:
Please install python-magic.""") Please install python-magic.""")
sys.exit(1) sys.exit(1)
storage = Path(app.config["FHOST_STORAGE_PATH"])
storage.mkdir(parents=True, exist_ok=True)
db = SQLAlchemy(app) db = SQLAlchemy(app)
migrate = Migrate(app, db) migrate = Migrate(app, db)
@ -103,6 +101,16 @@ class URL(db.Model):
def geturl(self): def geturl(self):
return url_for("get", path=self.getname(), _external=True) + "\n" return url_for("get", path=self.getname(), _external=True) + "\n"
def get(url):
u = URL.query.filter_by(url=url).first()
if not u:
u = URL(url)
db.session.add(u)
db.session.commit()
return u
class File(db.Model): class File(db.Model):
id = db.Column(db.Integer, primary_key = True) id = db.Column(db.Integer, primary_key = True)
sha256 = db.Column(db.String, unique = True) sha256 = db.Column(db.String, unique = True)
@ -112,12 +120,11 @@ class File(db.Model):
removed = db.Column(db.Boolean, default=False) removed = db.Column(db.Boolean, default=False)
nsfw_score = db.Column(db.Float) nsfw_score = db.Column(db.Float)
def __init__(self, sha256, ext, mime, addr, nsfw_score): def __init__(self, sha256, ext, mime, addr):
self.sha256 = sha256 self.sha256 = sha256
self.ext = ext self.ext = ext
self.mime = mime self.mime = mime
self.addr = addr self.addr = addr
self.nsfw_score = nsfw_score
def getname(self): def getname(self):
return u"{0}{1}".format(su.enbase(self.id, 1), self.ext) return u"{0}{1}".format(su.enbase(self.id, 1), self.ext)
@ -130,6 +137,71 @@ class File(db.Model):
else: else:
return url_for("get", path=n, _external=True) + "\n" return url_for("get", path=n, _external=True) + "\n"
def store(file_, addr):
data = file_.stream.read()
digest = sha256(data).hexdigest()
def get_mime():
guess = mimedetect.from_buffer(data)
app.logger.debug(f"MIME - specified: '{file_.content_type}' - detected: '{guess}'")
if not file_.content_type or not "/" in file_.content_type or file_.content_type == "application/octet-stream":
mime = guess
else:
mime = file_.content_type
if mime in app.config["FHOST_MIME_BLACKLIST"] or guess in app.config["FHOST_MIME_BLACKLIST"]:
abort(415)
if mime.startswith("text/") and not "charset" in mime:
mime += "; charset=utf-8"
return mime
def get_ext(mime):
ext = "".join(Path(file_.filename).suffixes[-2:])
gmime = mime[:mime.find(";")]
guess = guess_extension(gmime)
app.logger.debug(f"extension - specified: '{ext}' - detected: '{guess}'")
if not ext:
if gmime in app.config["FHOST_EXT_OVERRIDE"]:
ext = app.config["FHOST_EXT_OVERRIDE"][gmime]
else:
ext = guess_extension(gmime)
return ext[:app.config["FHOST_MAX_EXT_LENGTH"]] or ".bin"
f = File.query.filter_by(sha256=digest).first()
if f:
if f.removed:
abort(451)
else:
mime = get_mime()
ext = get_ext(mime)
f = File(digest, ext, mime, addr)
f.addr = addr
storage = Path(app.config["FHOST_STORAGE_PATH"])
storage.mkdir(parents=True, exist_ok=True)
p = storage / digest
if not p.is_file():
file_.stream.seek(0)
file_.save(p)
else:
p.touch()
if not f.nsfw_score and app.config["NSFW_DETECT"]:
f.nsfw_score = nsfw.detect(p)
db.session.add(f)
db.session.commit()
return f
def fhost_url(scheme=None): def fhost_url(scheme=None):
if not scheme: if not scheme:
return url_for(".fhost", _external=True).rstrip("/") return url_for(".fhost", _external=True).rstrip("/")
@ -146,16 +218,9 @@ def shorten(url):
if not url_valid(url) or is_fhost_url(url) or "\n" in url: if not url_valid(url) or is_fhost_url(url) or "\n" in url:
abort(400) abort(400)
existing = URL.query.filter_by(url=url).first() u = URL.get(url)
if existing: return u.geturl()
return existing.geturl()
else:
u = URL(url)
db.session.add(u)
db.session.commit()
return u.geturl()
def in_upload_bl(addr): def in_upload_bl(addr):
if app.config["FHOST_UPLOAD_BLACKLIST"]: if app.config["FHOST_UPLOAD_BLACKLIST"]:
@ -172,71 +237,9 @@ def store_file(f, addr):
if in_upload_bl(addr): if in_upload_bl(addr):
return "Your host is blocked from uploading files.\n", 451 return "Your host is blocked from uploading files.\n", 451
data = f.stream.read() sf = File.store(f, addr)
digest = sha256(data).hexdigest()
existing = File.query.filter_by(sha256=digest).first()
if existing: return sf.geturl()
if existing.removed:
abort(451)
epath = storage / existing.sha256
if not epath.is_file():
f.save(epath)
if existing.nsfw_score == None:
if app.config["NSFW_DETECT"]:
existing.nsfw_score = nsfw.detect(epath)
epath.touch()
existing.addr = addr
db.session.commit()
return existing.geturl()
else:
guessmime = mimedetect.from_buffer(data)
if not f.content_type or not "/" in f.content_type or f.content_type == "application/octet-stream":
mime = guessmime
else:
mime = f.content_type
if mime in app.config["FHOST_MIME_BLACKLIST"] or guessmime in app.config["FHOST_MIME_BLACKLIST"]:
abort(415)
if mime.startswith("text/") and not "charset" in mime:
mime += "; charset=utf-8"
ext = os.path.splitext(f.filename)[1]
if not ext:
gmime = mime.split(";")[0]
if not gmime in app.config["FHOST_EXT_OVERRIDE"]:
ext = guess_extension(gmime)
else:
ext = app.config["FHOST_EXT_OVERRIDE"][gmime]
else:
ext = ext[:8]
if not ext:
ext = ".bin"
spath = storage / digest
f.save(spath)
if app.config["NSFW_DETECT"]:
nsfw_score = nsfw.detect(spath)
else:
nsfw_score = None
sf = File(digest, ext, mime, addr, nsfw_score)
db.session.add(sf)
db.session.commit()
return sf.geturl()
def store_url(url, addr): def store_url(url, addr):
if is_fhost_url(url): if is_fhost_url(url):
@ -267,17 +270,19 @@ def store_url(url, addr):
@app.route("/<path:path>") @app.route("/<path:path>")
def get(path): def get(path):
p = os.path.splitext(path) path = Path(path)
id = su.debase(p[0]) sufs = "".join(path.suffixes[-2:])
name = path.name[:-len(sufs) or None]
id = su.debase(name)
if p[1]: if sufs:
f = File.query.get(id) f = File.query.get(id)
if f and f.ext == p[1]: if f and f.ext == sufs:
if f.removed: if f.removed:
abort(451) abort(451)
fpath = storage / f.sha256 fpath = Path(app.config["FHOST_STORAGE_PATH"]) / f.sha256
if not fpath.is_file(): if not fpath.is_file():
abort(404) abort(404)
@ -286,7 +291,7 @@ def get(path):
response = make_response() response = make_response()
response.headers["Content-Type"] = f.mime response.headers["Content-Type"] = f.mime
response.headers["Content-Length"] = fpath.stat().st_size response.headers["Content-Length"] = fpath.stat().st_size
response.headers["X-Accel-Redirect"] = "/" + fpath response.headers["X-Accel-Redirect"] = "/" + str(fpath)
return response return response
else: else:
return send_from_directory(app.config["FHOST_STORAGE_PATH"], f.sha256, mimetype = f.mime) return send_from_directory(app.config["FHOST_STORAGE_PATH"], f.sha256, mimetype = f.mime)