import datetime import os import uuid from fastapi import FastAPI, Depends, UploadFile, Form, File from sqlalchemy.orm import Session from starlette.requests import Request from starlette.responses import HTMLResponse import random from starlette.staticfiles import StaticFiles import database from database import engine, SessionLocal, Base Base.metadata.create_all(bind=engine) app = FastAPI() if not os.path.exists('./static'): os.makedirs('./static') app.mount("/static", StaticFiles(directory="static"), name="static") index_html = open('templates/index.html', 'r').read() admin_html = open('templates/admin.html', 'r').read() # 过期时间 exp_hour = 24 # 允许错误次数 error_count = 5 # 禁止分钟数 error_minute = 60 # 后台地址 admin_address = 'admin' # 管理密码 admin_password = 'admin' error_ip_count = {} def get_db(): db = SessionLocal() try: yield db finally: db.close() def get_code(db: Session = Depends(get_db)): code = random.randint(10000, 99999) while db.query(database.Codes).filter(database.Codes.code == code).first(): code = random.randint(10000, 99999) return str(code) def get_file_name(key, ext, file): now = datetime.datetime.now() path = f'./static/upload/{now.year}/{now.month}/{now.day}/' name = f'{key}.{ext}' if not os.path.exists(path): os.makedirs(path) file = file.file.read() with open(f'{os.path.join(path, name)}', 'wb') as f: f.write(file) return key, len(file), path[1:] + name @app.get(f'/{admin_address}') async def admin(request: Request): return HTMLResponse(admin_html) @app.post(f'/{admin_address}') async def admin_post(request: Request, db: Session = Depends(get_db)): if request.headers.get('pwd') == admin_password: codes = db.query(database.Codes).all() return {'code': 200, 'msg': '查询成功', 'data': codes} else: return {'code': 400, 'msg': '密码错误'} @app.delete(f'/{admin_address}') async def admin_delete(request: Request, code: str, db: Session = Depends(get_db)): if request.headers.get('pwd') == admin_password: file = db.query(database.Codes).filter(database.Codes.code == code) if file.first().type != 'text/plain': os.remove('.' + file.first().text) file.delete() db.commit() return {'code': 200, 'msg': '删除成功'} else: return {'code': 400, 'msg': '密码错误'} @app.get('/') async def index(): return HTMLResponse(index_html) @app.post('/') async def index(request: Request, code: str, db: Session = Depends(get_db)): info = db.query(database.Codes).filter(database.Codes.code == code).first() error = error_ip_count.get(request.client.host, {'count': 0, 'time': datetime.datetime.now()}) if error['count'] > error_count: if datetime.datetime.now() - error['time'] < datetime.timedelta(minutes=error_minute): return {'code': 404, 'msg': '请求过于频繁,请稍后再试'} else: error['count'] = 0 else: if not info: error['count'] += 1 error_ip_count[request.client.host] = error return {'code': 404, 'msg': f'取件码错误,错误5次将被禁止10分钟'} else: return {'code': 200, 'msg': '取件成功,请点击库查看', 'data': info} @app.post('/share') async def share(text: str = Form(default=None), file: UploadFile = File(default=None), db: Session = Depends(get_db)): cutoff_time = datetime.datetime.now() - datetime.timedelta(hours=exp_hour) db.query(database.Codes).filter(database.Codes.use_time < cutoff_time).delete() db.commit() code = get_code(db) if text: info = database.Codes( code=code, text=text, type='text/plain', key=uuid.uuid4().hex, size=len(text), used=True, name='分享文本' ) db.add(info) db.commit() return {'code': 200, 'msg': '上传成功,请点击文件库查看', 'data': {'code': code, 'name': '分享文本', 'text': text}} elif file: key, size, full_path = get_file_name(uuid.uuid4().hex, file.filename.split('.')[-1], file) info = database.Codes( code=code, text=full_path, type=file.content_type, key=key, size=size, used=True, name=file.filename ) db.add(info) db.commit() return {'code': 200, 'msg': '上传成功,请点击文件库查看', 'data': {'code': code, 'name': file.filename, 'text': full_path}} else: return {'code': 422, 'msg': '参数错误', 'data': []}