orion/filecodebox
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update:新增限制时间内上传次数，移除IP黑名单列表过期IP，防止炸内存

Browse Source
This commit is contained in:
lan
2022-12-13 19:16:29 +08:00
parent 97cff6dad2
commit d23cffb570
4 changed files with 79 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
depends.py
+16 -5
View File
@@ -13,25 +13,36 @@ async def admin_required(pwd: Union[str, None] = Header(default=None)):
class IPRateLimit:
ips = {}
count = 0
minutes = 0
def __init__(self, count, minutes):
self.count = count
self.minutes = minutes
def check_ip(self, ip):
# 检查ip是否被禁止
if ip in self.ips:
if self.ips[ip]['count'] >= settings.ERROR_COUNT:
if self.ips[ip]['time'] + timedelta(minutes=settings.ERROR_MINUTE) > datetime.now():
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(cls, ip):
ip_info = cls.ips.get(ip, {'count': 0, 'time': datetime.now()})
ip_info['count'] += 1
cls.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.client.host
if not self.check_ip(ip):
raise HTTPException(status_code=400, detail="错误次数过多，请稍后再试")
raise HTTPException(status_code=400, detail=f"请求次数过多，请稍后再试")
return ip
main.py
+35 -21
View File
@@ -11,58 +11,66 @@ from sqlalchemy import select, update
from sqlalchemy.ext.asyncio.session import AsyncSession
import settings
from utils import delete_expire_files, storage, get_code
from utils import delete_expire_files, storage, get_code, error_ip_limit, upload_ip_limit
from database import get_session, Codes, init_models
from depends import admin_required, IPRateLimit
from depends import admin_required
# 实例化FastAPI
app = FastAPI(debug=settings.DEBUG)
# 数据存储文件夹
DATA_ROOT = Path(settings.DATA_ROOT)
if not DATA_ROOT.exists():
DATA_ROOT.mkdir(parents=True)
# 静态文件夹
app.mount(settings.STATIC_URL, StaticFiles(directory=DATA_ROOT), name="static")
@app.on_event('startup')
async def startup():
# 初始化数据库
await init_models()
# 启动后台任务，不定时删除过期文件
asyncio.create_task(delete_expire_files())
# 首页页面
index_html = open('templates/index.html', 'r', encoding='utf-8').read() \
.replace('{{title}}', settings.TITLE) \
.replace('{{description}}', settings.DESCRIPTION) \
.replace('{{keywords}}', settings.KEYWORDS)
# 管理页面
admin_html = open('templates/admin.html', 'r', encoding='utf-8').read() \
.replace('{{title}}', settings.TITLE) \
.replace('{{description}}', settings.DESCRIPTION) \
.replace('{{keywords}}', settings.KEYWORDS)
ip_limit = IPRateLimit()
@app.get(f'/{settings.ADMIN_ADDRESS}')
@app.get(f'/{settings.ADMIN_ADDRESS}', description='管理页面', response_class=HTMLResponse)
async def admin():
return HTMLResponse(admin_html)
@app.post(f'/{settings.ADMIN_ADDRESS}', dependencies=[Depends(admin_required)])
@app.post(f'/{settings.ADMIN_ADDRESS}', dependencies=[Depends(admin_required)], description='查询数据库列表')
async def admin_post(s: AsyncSession = Depends(get_session)):
query = select(Codes)
codes = (await s.execute(query)).scalars().all()
# 查询数据库列表
codes = (await s.execute(select(Codes))).scalars().all()
return {'detail': '查询成功', 'data': codes}
@app.delete(f'/{settings.ADMIN_ADDRESS}', dependencies=[Depends(admin_required)])
@app.delete(f'/{settings.ADMIN_ADDRESS}', dependencies=[Depends(admin_required)], description='删除数据库记录')
async def admin_delete(code: str, s: AsyncSession = Depends(get_session)):
# 找到相应记录
query = select(Codes).where(Codes.code == code)
# 找到第一条记录
file = (await s.execute(query)).scalars().first()
if file:
if file.type != 'text':
await storage.delete_file(file.text)
await s.delete(file)
await s.commit()
# 如果记录存在，并且不是文本
if file and file.type != 'text':
# 删除文件
await storage.delete_file(file.text)
# 删除数据库记录
await s.delete(file)
await s.commit()
return {'detail': '删除成功'}
@@ -72,25 +80,30 @@ async def index():
@app.get('/select')
async def get_file(code: str, s: AsyncSession = Depends(get_session)):
async def get_file(code: str, ip: str = Depends(error_ip_limit), s: AsyncSession = Depends(get_session)):
# 查出数据库记录
query = select(Codes).where(Codes.code == code)
info = (await s.execute(query)).scalars().first()
# 如果记录不存在，IP错误次数+1
if not info:
raise HTTPException(status_code=404, detail="口令不存在")
error_ip_limit.add_ip(ip)
raise HTTPException(status_code=404, detail="口令不存在，次数过多将被禁止访问")
# 如果是文本，直接返回
if info.type == 'text':
return {'detail': '查询成功', 'data': info.text}
# 如果是文件，返回文件
else:
filepath = await storage.get_filepath(info.text)
return FileResponse(filepath, filename=info.name)
@app.post('/')
async def index(code: str, ip: str = Depends(ip_limit), s: AsyncSession = Depends(get_session)):
async def index(code: str, ip: str = Depends(error_ip_limit), s: AsyncSession = Depends(get_session)):
query = select(Codes).where(Codes.code == code)
info = (await s.execute(query)).scalars().first()
if not info:
error_count = settings.ERROR_COUNT - ip_limit.add_ip(ip)
raise HTTPException(status_code=404, detail=f"取件码错误，错误{error_count}次将被禁止10分钟")
error_count = settings.ERROR_COUNT - error_ip_limit.add_ip(ip)
raise HTTPException(status_code=404, detail=f"取件码错误，{error_count}次后将被禁止{settings.ERROR_MINUTE}分钟")
if info.exp_time < datetime.datetime.now() or info.count == 0:
if info.type != "text":
await storage.delete_file(info.text)
@@ -109,7 +122,7 @@ async def index(code: str, ip: str = Depends(ip_limit), s: AsyncSession = Depend
@app.post('/share')
async def share(background_tasks: BackgroundTasks, text: str = Form(default=None), style: str = Form(default='2'),
value: int = Form(default=1), file: UploadFile = File(default=None),
value: int = Form(default=1), file: UploadFile = File(default=None), ip: str = Depends(upload_ip_limit),
s: AsyncSession = Depends(get_session)):
code = await get_code(s)
if style == '2':
@@ -137,6 +150,7 @@ async def share(background_tasks: BackgroundTasks, text: str = Form(default=None
info = Codes(code=code, text=_text, size=size, type=_type, name=name, count=exp_count, exp_time=exp_time, key=key)
s.add(info)
await s.commit()
upload_ip_limit.add_ip(ip)
return {
'detail': '分享成功，请点击取件码按钮查看上传列表',
'data': {'code': code, 'key': key, 'name': name, 'text': _text}
settings.py
+19 -14
View File
@@ -1,31 +1,36 @@
from starlette.config import Config
# 配置文件.env
config = Config(".env")
# 是否开启DEBUG模式
DEBUG = config('DEBUG', cast=bool, default=False)
# 端口
PORT = config('PORT', cast=int, default=12345)
# Sqlite数据库文件
DATABASE_URL = config('DATABASE_URL', cast=str, default="sqlite+aiosqlite:///database.db")
# 静态文件夹
DATA_ROOT = config('DATA_ROOT', cast=str, default="./static")
# 静态文件夹URL
STATIC_URL = config('STATIC_URL', cast=str, default="/static")
# 错误次数
ERROR_COUNT = config('ERROR_COUNT', cast=int, default=5)
# 错误限制分钟数
ERROR_MINUTE = config('ERROR_MINUTE', cast=int, default=10)
# 上传次数
UPLOAD_COUNT = config('UPLOAD_COUNT', cast=int, default=60)
# 上传限制分钟数
UPLOAD_MINUTE = config('UPLOAD_MINUTE', cast=int, default=1)
# 管理地址
ADMIN_ADDRESS = config('ADMIN_ADDRESS', cast=str, default="admin")
# 管理密码
ADMIN_PASSWORD = config('ADMIN_PASSWORD', cast=str, default="admin")
# 文件大小限制，默认10MB
FILE_SIZE_LIMIT = config('FILE_SIZE_LIMIT', cast=int, default=10) * 1024 * 1024
# 网站标题
TITLE = config('TITLE', cast=str, default="文件快递柜")
# 网站描述
DESCRIPTION = config('DESCRIPTION', cast=str, default="FileCodeBox，文件快递柜，口令传送箱，匿名口令分享文本，文件等文件")
# 网站关键词
KEYWORDS = config('KEYWORDS', cast=str, default="FileCodeBox，文件快递柜，口令传送箱，匿名口令分享文本，文件等文件")
# 存储引擎
STORAGE_ENGINE = config('STORAGE_ENGINE', cast=str, default="filesystem")
utils.py
+9 -1
View File
@@ -5,14 +5,22 @@ from sqlalchemy import or_, select, delete
from sqlalchemy.ext.asyncio.session import AsyncSession
import settings
from database import Codes, engine
from depends import IPRateLimit
from storage import STORAGE_ENGINE
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 = []
@@ -25,7 +33,7 @@ async def delete_expire_files():
query = delete(Codes).where(Codes.id.in_(exps_ids))
await s.execute(query)
await s.commit()
await asyncio.sleep(random.randint(60, 300))
await asyncio.sleep(random.randint(2, 2))
async def get_code(s: AsyncSession):