重构中，配置存入数据库，第一次智能生成随机后台地址，登录密码

core/database.py

@@ -0,0 +1,76 @@
+import datetime
+
+from sqlalchemy import Boolean, Column, Integer, String, DateTime, JSON, Text, select, insert, delete
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.ext.asyncio import create_async_engine
+from sqlalchemy.ext.asyncio.session import AsyncSession
+from settings import settings
+
+engine = create_async_engine(settings.DATABASE_URL)
+
+Base = declarative_base()
+
+
+class Options(Base):
+    __tablename__ = 'options'
+    id = Column(Integer, primary_key=True, index=True)
+    key = Column(String, unique=True, index=True)
+    value = Column(JSON)
+
+
+class Codes(Base):
+    __tablename__ = "codes"
+    id = Column(Integer, primary_key=True, index=True)
+    code = Column(String(10), unique=True, index=True)
+    key = Column(String(30), unique=True)
+    name = Column(String(500))
+    size = Column(Integer)
+    type = Column(String(20))
+    text = Column(Text)
+    used = Column(Boolean, default=False)
+    count = Column(Integer, default=-1)
+    use_time = Column(DateTime, default=datetime.datetime.now, onupdate=datetime.datetime.now)
+    exp_time = Column(DateTime, nullable=True)
+
+
+async def init_models():
+    async with engine.begin() as conn:
+        await conn.run_sync(Base.metadata.create_all)
+        if await conn.scalar(select(Options).filter(Options.key == 'INSTALL')) is None:
+            # 如果没有存在install，则清空表，并插入默认数据
+            await conn.execute(delete(table=Options))
+            await conn.execute(insert(table=Options, values=[
+                {'key': 'INSTALL', 'value': settings.VERSION},
+                {'key': 'BANNERS', 'value': settings.BANNERS},
+                {'key': 'ENABLE_UPLOAD', 'value': settings.ENABLE_UPLOAD},
+                {'key': 'MAX_DAYS', 'value': settings.MAX_DAYS},
+                {'key': 'ERROR_COUNT', 'value': settings.ERROR_COUNT},
+                {'key': 'ERROR_MINUTE', 'value': settings.ERROR_COUNT},
+                {'key': 'UPLOAD_COUNT', 'value': settings.UPLOAD_COUNT},
+                {'key': 'UPLOAD_MINUTE', 'value': settings.UPLOAD_MINUTE},
+                {'key': 'DELETE_EXPIRE_FILES_INTERVAL', 'value': settings.DELETE_EXPIRE_FILES_INTERVAL},
+                {'key': 'ADMIN_ADDRESS', 'value': settings.ADMIN_ADDRESS},
+                {'key': 'ADMIN_PASSWORD', 'value': settings.ADMIN_PASSWORD},
+                {'key': 'FILE_SIZE_LIMIT', 'value': settings.FILE_SIZE_LIMIT},
+                {'key': 'TITLE', 'value': settings.TITLE},
+                {'key': 'DESCRIPTION', 'value': settings.DESCRIPTION},
+                {'key': 'KEYWORDS', 'value': settings.KEYWORDS},
+                {'key': 'STORAGE_ENGINE', 'value': settings.STORAGE_ENGINE},
+                {'key': 'STORAGE_CONFIG', 'value': {}},
+            ]))
+            print(
+                f'初始化数据库成功！\n'
+                f'如您未配置.env文件，将为您随机生成信息\n'
+                f'您的后台地址为：/{settings.ADMIN_ADDRESS}\n'
+                f'您的管理员密码为：{settings.ADMIN_PASSWORD}\n'
+                f'请尽快修改后台信息！\n'
+                f'FileCodeBox https://github.com/vastsa/FileCodeBox'
+            )
+        else:
+            # 从数据库更新缓存中的setting
+            await settings.update(await conn.execute(select(Options).filter()))
+
+
+async def get_session():
+    async with AsyncSession(engine, expire_on_commit=False) as s:
+        yield s

core/depends.py

@@ -0,0 +1,52 @@
+from typing import Union
+from datetime import datetime, timedelta
+
+from fastapi import Header, HTTPException, Request
+
+from settings import settings
+
+
+async def admin_required(pwd: Union[str, None] = Header(default=None), request: Request = None):
+    if 'share' in request.url.path:
+        if pwd != settings.ADMIN_PASSWORD and not settings.ENABLE_UPLOAD:
+            raise HTTPException(status_code=403, detail='本站上传功能已关闭，仅管理员可用')
+    else:
+        print(settings.ADMIN_PASSWORD)
+        if settings.ADMIN_PASSWORD is None:
+            raise HTTPException(status_code=404, detail='您未设置管理员密码，无法使用此功能，请更新配置文件后，重启系统')
+        if not pwd or pwd != settings.ADMIN_PASSWORD:
+            raise HTTPException(status_code=401, detail="密码错误，请重新登录")
+
+
+class IPRateLimit:
+    def __init__(self, count, minutes):
+        self.ips = {}
+        self.count = count
+        self.minutes = minutes
+
+    def check_ip(self, ip):
+        # 检查ip是否被禁止
+        if ip in self.ips:
+            if self.ips[ip]['count'] >= self.count:
+                if self.ips[ip]['time'] + timedelta(minutes=self.minutes) > datetime.now():
+                    return False
+                else:
+                    self.ips.pop(ip)
+        return True
+
+    def add_ip(self, ip):
+        ip_info = self.ips.get(ip, {'count': 0, 'time': datetime.now()})
+        ip_info['count'] += 1
+        self.ips[ip] = ip_info
+        return ip_info['count']
+
+    async def remove_expired_ip(self):
+        for ip in list(self.ips.keys()):
+            if self.ips[ip]['time'] + timedelta(minutes=self.minutes) < datetime.now():
+                self.ips.pop(ip)
+
+    def __call__(self, request: Request):
+        ip = request.headers.get('X-Real-IP', request.headers.get('X-Forwarded-For', request.client.host))
+        if not self.check_ip(ip):
+            raise HTTPException(status_code=400, detail=f"请求次数过多，请稍后再试")
+        return ip

core/storage.py

@@ -0,0 +1,138 @@
+import os
+import asyncio
+from datetime import datetime
+from pathlib import Path
+from typing import BinaryIO
+from fastapi import UploadFile
+
+from core.database import Codes
+from settings import settings
+
+if settings.STORAGE_ENGINE == 'aliyunsystem':
+    try:
+        import oss2
+    except ImportError:
+        os.system('pip install oss2')
+        import oss2
+
+
+class AliyunFileStorage:
+    def __init__(self):
+        auth = oss2.Auth(settings.KeyId, settings.KeySecret)
+        self.bucket = oss2.Bucket(auth, settings.OSS_ENDPOINT, settings.BUCKET_NAME)
+
+    def upload_file(self, upload_filepath, remote_filepath):
+        self.bucket.put_object_from_file(remote_filepath, upload_filepath)
+
+    async def get_text(self, file: UploadFile, key: str):
+        ext = file.filename.split('.')[-1]
+        now = datetime.now()
+        path = f"FileCodeBox/upload/{now.year}/{now.month}/{now.day}"
+        text = f"{path}/{f'{key}.{ext}'}"
+        return f"https://{settings.BUCKET_NAME}.{settings.OSS_ENDPOINT}/{text}"
+
+    async def get_url(self, info: Codes):
+        text = info.text.strip(f"https://{settings.BUCKET_NAME}.{settings.OSS_ENDPOINT}/")
+        url = self.bucket.sign_url('GET', text, settings.ACCESSTIME, slash_safe=True)
+        return url
+
+    @staticmethod
+    async def get_size(file: UploadFile):
+        f = file.file
+        f.seek(0, os.SEEK_END)
+        size = f.tell()
+        f.seek(0, os.SEEK_SET)
+        return size
+
+    @staticmethod
+    def _save(filepath, file: BinaryIO):
+        with open(filepath, 'wb') as f:
+            chunk_size = 256 * 1024
+            chunk = file.read(chunk_size)
+            while chunk:
+                f.write(chunk)
+                chunk = file.read(chunk_size)
+
+    async def save_file(self, file: UploadFile, remote_filepath: str):
+        now = int(datetime.now().timestamp())
+        upload_filepath = settings.DATA_ROOT + str(now)
+        await asyncio.to_thread(self._save, upload_filepath, file.file)
+        self.upload_file(upload_filepath, remote_filepath)
+        await asyncio.to_thread(os.remove, upload_filepath)
+
+    async def delete_files(self, texts):
+        tasks = [self.delete_file(text) for text in texts]
+        await asyncio.gather(*tasks)
+
+    async def delete_file(self, text: str):
+        text = text.strip(f"https://{settings.BUCKET_NAME}.{settings.OSS_ENDPOINT}/")
+        self.bucket.delete_object(text)
+
+
+class FileSystemStorage:
+    def __init__(self):
+        self.DATA_ROOT = Path(settings.DATA_ROOT)
+        self.STATIC_URL = settings.STATIC_URL
+        self.NAME = "filesystem"
+        self.DOWN_PATH = '/select'
+
+    async def get_filepath(self, text: str):
+        return self.DATA_ROOT / text.lstrip(self.STATIC_URL + '/')
+
+    async def get_url(self, info: Codes):
+        return f'{self.DOWN_PATH}?code={info.code}'
+
+    async def get_text(self, file: UploadFile, key: str):
+        ext = file.filename.split('.')[-1]
+        now = datetime.now()
+        path = self.DATA_ROOT / f"upload/{now.year}/{now.month}/{now.day}/"
+        if not path.exists():
+            path.mkdir(parents=True)
+        text = f"{self.STATIC_URL}/{(path / f'{key}.{ext}').relative_to(self.DATA_ROOT)}"
+        return text
+
+    @staticmethod
+    async def get_size(file: UploadFile):
+        f = file.file
+        f.seek(0, os.SEEK_END)
+        size = f.tell()
+        f.seek(0, os.SEEK_SET)
+        return size
+
+    @staticmethod
+    def _save(filepath, file: BinaryIO):
+        with open(filepath, 'wb') as f:
+            chunk_size = 256 * 1024
+            chunk = file.read(chunk_size)
+            while chunk:
+                f.write(chunk)
+                chunk = file.read(chunk_size)
+
+    async def save_file(self, file: UploadFile, text: str):
+        filepath = await self.get_filepath(text)
+        await asyncio.to_thread(self._save, filepath, file.file)
+
+    async def delete_file(self, text: str):
+        filepath = await self.get_filepath(text)
+        if filepath.exists():
+            await asyncio.to_thread(os.remove, filepath)
+            await asyncio.to_thread(self.judge_delete_folder, filepath)
+
+    async def delete_files(self, texts):
+        tasks = [self.delete_file(text) for text in texts]
+        await asyncio.gather(*tasks)
+
+    def judge_delete_folder(self, filepath):
+        current = filepath.parent
+        while current != self.DATA_ROOT:
+            if not list(current.iterdir()):
+                os.rmdir(current)
+                current = current.parent
+            else:
+                break
+
+
+STORAGE_ENGINE = {
+    "filesystem": FileSystemStorage,
+    "aliyunsystem": AliyunFileStorage
+}

core/utils.py

@@ -0,0 +1,43 @@
+import datetime
+import random
+import asyncio
+from sqlalchemy import or_, select, delete
+from sqlalchemy.ext.asyncio.session import AsyncSession
+from .database import Codes, engine
+from .depends import IPRateLimit
+from .storage import STORAGE_ENGINE
+from settings import settings
+
+storage = STORAGE_ENGINE[settings.STORAGE_ENGINE]()
+
+# 错误IP限制器
+error_ip_limit = IPRateLimit(settings.ERROR_COUNT, settings.ERROR_MINUTE)
+# 上传文件限制器
+upload_ip_limit = IPRateLimit(settings.UPLOAD_COUNT, settings.UPLOAD_MINUTE)
+
+
+async def delete_expire_files():
+    while True:
+        async with AsyncSession(engine, expire_on_commit=False) as s:
+            await error_ip_limit.remove_expired_ip()
+            await upload_ip_limit.remove_expired_ip()
+            query = select(Codes).where(or_(Codes.exp_time < datetime.datetime.now(), Codes.count == 0))
+            exps = (await s.execute(query)).scalars().all()
+            files = []
+            exps_ids = []
+            for exp in exps:
+                if exp.type != "text":
+                    files.append(exp.text)
+                exps_ids.append(exp.id)
+            await storage.delete_files(files)
+            query = delete(Codes).where(Codes.id.in_(exps_ids))
+            await s.execute(query)
+            await s.commit()
+        await asyncio.sleep(settings.DELETE_EXPIRE_FILES_INTERVAL * 60)
+
+
+async def get_code(s: AsyncSession):
+    code = random.randint(10000, 99999)
+    while (await s.execute(select(Codes.id).where(Codes.code == code))).scalar():
+        code = random.randint(10000, 99999)
+    return str(code)