lan-air
150 lines
4.7 KiB
Python
150 lines
4.7 KiB
Python
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():
|
|
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': []}
|