From ea382227f5aae159703d8edffb741d191c855cc6 Mon Sep 17 00:00:00 2001 From: veoco Date: Mon, 12 Dec 2022 22:30:03 +0800 Subject: [PATCH] =?UTF-8?q?=E9=87=8D=E5=86=99=E7=94=A8=E6=88=B7=E7=99=BB?= =?UTF-8?q?=E5=BD=95=E5=92=8C=20IP=20=E6=A3=80=E6=9F=A5=E5=B9=B6=E9=87=8D?= =?UTF-8?q?=E5=AE=9A=E4=B9=89=E7=8A=B6=E6=80=81=E7=A0=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- depends.py | 37 +++++++++++++++++++ main.py | 87 +++++++++++++++----------------------------- templates/admin.html | 31 ++++++---------- templates/index.html | 83 +++++++++++++++++++----------------------- 4 files changed, 116 insertions(+), 122 deletions(-) create mode 100644 depends.py diff --git a/depends.py b/depends.py new file mode 100644 index 0000000..72bda9c --- /dev/null +++ b/depends.py @@ -0,0 +1,37 @@ +from typing import Union +from datetime import datetime, timedelta + +from fastapi import Header, HTTPException, Request + +import settings + + +async def admin_required(pwd: Union[str, None] = Header(default=None)): + if not pwd or pwd != settings.ADMIN_PASSWORD: + raise HTTPException(status_code=401, detail="密码错误") + + +class IPRateLimit: + ips = {} + + 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(): + 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'] + + def __call__(self, request: Request): + ip = request.client.host + if not self.check_ip(ip): + raise HTTPException(status_code=400, detail="错误次数过多，请稍后再试") + return ip diff --git a/main.py b/main.py index 6be2abe..4c6b54b 100644 --- a/main.py +++ b/main.py @@ -4,8 +4,7 @@ import random import asyncio from pathlib import Path -from fastapi import FastAPI, Depends, UploadFile, Form, File -from starlette.requests import Request +from fastapi import FastAPI, Depends, UploadFile, Form, File, HTTPException from starlette.responses import HTMLResponse, FileResponse from starlette.staticfiles import StaticFiles @@ -15,6 +14,7 @@ from sqlalchemy.ext.asyncio.session import AsyncSession import settings from database import get_session, Codes, init_models, engine from storage import STORAGE_ENGINE +from depends import admin_required, IPRateLimit app = FastAPI(debug=settings.DEBUG) @@ -43,7 +43,7 @@ admin_html = open('templates/admin.html', 'r', encoding='utf-8').read() \ .replace('{{description}}', settings.DESCRIPTION) \ .replace('{{keywords}}', settings.KEYWORDS) -error_ip_count = {} +ip_limit = IPRateLimit() async def delete_expire_files(): @@ -72,27 +72,21 @@ async def admin(): return HTMLResponse(admin_html) -@app.post(f'/{settings.ADMIN_ADDRESS}') -async def admin_post(request: Request, s: AsyncSession = Depends(get_session)): - if request.headers.get('pwd') == settings.ADMIN_PASSWORD: - query = select(Codes) - codes = (await s.execute(query)).scalars().all() - return {'code': 200, 'msg': '查询成功', 'data': codes} - else: - return {'code': 404, 'msg': '密码错误'} +@app.post(f'/{settings.ADMIN_ADDRESS}', dependencies=[Depends(admin_required)]) +async def admin_post(s: AsyncSession = Depends(get_session)): + query = select(Codes) + codes = (await s.execute(query)).scalars().all() + return {'msg': '查询成功', 'data': codes} -@app.delete(f'/{settings.ADMIN_ADDRESS}') -async def admin_delete(request: Request, code: str, s: AsyncSession = Depends(get_session)): - if request.headers.get('pwd') == settings.ADMIN_PASSWORD: - query = select(Codes).where(Codes.code == code) - file = (await s.execute(query)).scalars().first() - await storage.delete_file({'type': file.type, 'text': file.text}) - await s.delete(file) - await s.commit() - return {'code': 200, 'msg': '删除成功'} - else: - return {'code': 404, 'msg': '密码错误'} +@app.delete(f'/{settings.ADMIN_ADDRESS}', dependencies=[Depends(admin_required)]) +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() + await storage.delete_file({'type': file.type, 'text': file.text}) + await s.delete(file) + await s.commit() + return {'msg': '删除成功'} @app.get('/') @@ -100,52 +94,31 @@ async def index(): return HTMLResponse(index_html) -def check_ip(ip): - # 检查ip是否被禁止 - if ip in error_ip_count: - if error_ip_count[ip]['count'] >= settings.ERROR_COUNT: - if error_ip_count[ip]['time'] + datetime.timedelta(minutes=settings.ERROR_MINUTE) > datetime.datetime.now(): - return False - else: - error_ip_count.pop(ip) - return True - - -def ip_error(ip): - ip_info = error_ip_count.get(ip, {'count': 0, 'time': datetime.datetime.now()}) - ip_info['count'] += 1 - error_ip_count[ip] = ip_info - return ip_info['count'] - - @app.get('/select') async def get_file(code: str, s: AsyncSession = Depends(get_session)): query = select(Codes).where(Codes.code == code) info = (await s.execute(query)).scalars().first() - if info: - if info.type == 'text': - return {'code': code, 'msg': '查询成功', 'data': info.text} - else: - filepath = await storage.get_filepath(info.text) - return FileResponse(filepath, filename=info.name) + if not info: + raise HTTPException(status_code=404, detail="口令不存在") + if info.type == 'text': + return {'msg': '查询成功', 'data': info.text} else: - return {'code': 404, 'msg': '口令不存在'} + filepath = await storage.get_filepath(info.text) + return FileResponse(filepath, filename=info.name) @app.post('/') -async def index(request: Request, code: str, s: AsyncSession = Depends(get_session)): - ip = request.client.host - if not check_ip(ip): - return {'code': 404, 'msg': '错误次数过多，请稍后再试'} +async def index(code: str, ip: str = Depends(ip_limit), s: AsyncSession = Depends(get_session)): query = select(Codes).where(Codes.code == code) info = (await s.execute(query)).scalars().first() if not info: - return {'code': 404, 'msg': f'取件码错误，错误{settings.ERROR_COUNT - ip_error(ip)}次将被禁止10分钟'} + error_count = ip_limit.add_ip(ip) + raise HTTPException(status_code=404, detail=f"取件码错误，错误{settings.ERROR_COUNT - error_count}次将被禁止10分钟") if info.exp_time < datetime.datetime.now() or info.count == 0: await storage.delete_file({'type': info.type, 'text': info.text}) await s.delete(info) await s.commit() - return {'code': 404, 'msg': '取件码已过期，请联系寄件人'} + raise HTTPException(status_code=404, detail="取件码已过期，请联系寄件人") count = info.count - 1 query = update(Codes).where(Codes.id == info.id).values(count=count) await s.execute(query) @@ -153,7 +126,6 @@ async def index(request: Request, code: str, s: AsyncSession = Depends(get_sessi if info.type != 'text': info.text = f'/select?code={code}' return { - 'code': 200, 'msg': '取件成功，请点击"取"查看', 'data': {'type': info.type, 'text': info.text, 'name': info.name, 'code': info.code} } @@ -165,12 +137,12 @@ async def share(text: str = Form(default=None), style: str = Form(default='2'), code = await get_code(s) if style == '2': if value > 7: - return {'code': 404, 'msg': '最大有效天数为7天'} + raise HTTPException(status_code=400, detail="最大有效天数为7天") exp_time = datetime.datetime.now() + datetime.timedelta(days=value) exp_count = -1 elif style == '1': if value < 1: - return {'code': 404, 'msg': '最小有效次数为1次'} + raise HTTPException(status_code=400, detail="最小有效次数为1次") exp_time = datetime.datetime.now() + datetime.timedelta(days=1) exp_count = value else: @@ -181,7 +153,7 @@ async def share(text: str = Form(default=None), style: str = Form(default='2'), file_bytes = await file.read() size = len(file_bytes) if size > settings.FILE_SIZE_LIMIT: - return {'code': 404, 'msg': '文件过大'} + raise HTTPException(status_code=400, detail="文件过大") _text, _type, name = await storage.save_file(file, file_bytes, key), file.content_type, file.filename else: size, _text, _type, name = len(text), text, 'text', '文本分享' @@ -198,7 +170,6 @@ async def share(text: str = Form(default=None), style: str = Form(default='2'), s.add(info) await s.commit() return { - 'code': 200, 'msg': '分享成功，请点击文件箱查看取件码', 'data': {'code': code, 'key': key, 'name': name, 'text': _text} } diff --git a/templates/admin.html b/templates/admin.html index 7c885bb..2050dbc 100644 --- a/templates/admin.html +++ b/templates/admin.html @@ -87,28 +87,21 @@ methods: { loginAdmin: function () { axios.post('', {}, {'headers': {'pwd': this.pwd}}).then(res => { - if (res.data.code === 200) { - this.files = res.data.data; - this.login = true; - localStorage.setItem('pwd', this.pwd); - } else { - localStorage.clear() - this.$message({ - message: res.data.msg, - type: 'error' - }); - } - }) + this.files = res.data.data; + this.login = true; + localStorage.setItem('pwd', this.pwd); + }).catch(e => { + localStorage.clear() + this.$message({'message': e.response.data.detail, 'type': 'error'}); + }); }, deleteFile: function (code) { axios.delete('?code=' + code, {'headers': {'pwd': this.pwd}}).then(res => { - if (res.data.code === 200) { - this.files = this.files.filter(item => item.code !== code) - this.$message({ - message: res.data.msg, - type: 'success' - }); - } + this.files = this.files.filter(item => item.code !== code) + this.$message({ + message: res.data.msg, + type: 'success' + }); }) }, } diff --git a/templates/index.html b/templates/index.html index 352453e..57d6aaf 100644 --- a/templates/index.html +++ b/templates/index.html @@ -94,6 +94,7 @@ multiple :data="destoryData" :on-success="successUpload" + :on-error="errorUpload" >
将文字、文件拖、粘贴到此处，或点击上传
@@ -206,16 +207,12 @@ FileData.append('style', that.destoryData.style); FileData.append('value', that.destoryData.value); axios.post('/share', FileData) - .then(function (res) { - if (res.data.code === 200) { - that.$message({'message': res.data.msg, 'type': 'success'}); - that.files.push(res.data.data) - } else { - that.$message({'message': res.data.msg, 'type': 'error'}); - } + .then(res => { + that.$message({'message': res.data.msg, 'type': 'success'}); + that.files.push(res.data.data) }) - .catch(function (error) { - that.$message({'message': error.data.msg, 'type': 'error'}); + .catch(e => { + that.$message({'message': e.response.data.detail, 'type': 'error'}); }); }); } @@ -226,14 +223,14 @@ FileData.append('file', file || ''); FileData.append('style', that.destoryData.style); FileData.append('value', that.destoryData.value); - axios.post('/share', FileData).then(async res => { - if (res.data.code === 200) { - that.$message({'message': res.data.msg, 'type': 'success'}); - that.files.push(res.data.data) - } else { - that.$message({'message': res.data.msg, 'type': 'error'}); - } + axios.post('/share', FileData) + .then(res => { + that.$message({'message': res.data.msg, 'type': 'success'}); + that.files.push(res.data.data) }) + .catch(e => { + that.$message({'message': e.response.data.detail, 'type': 'error'}); + }); } } @@ -250,25 +247,21 @@ }, get_file: function () { const that = this; - axios.post('?code=' + this.code).then(function (response) { - if (response.data.code === 404) { - that.$message({ - message: response.data.msg, - type: 'error' - }); - } else { - that.files.push(response.data.data); - that.$message({ - message: response.data.msg, - type: 'success' - }); - that.quDrawer = true; - } + axios.post('?code=' + this.code).then(response => { + that.files.push(response.data.data); + that.$message({ + message: response.data.msg, + type: 'success' + }); + that.quDrawer = true; that.code = ''; - that.inout_disable = false - }).catch(function (error) { - console.log(error); + }).catch(e => { + that.$message({ + message: e.response.data.detail, + type: 'error' + }); }); + that.inout_disable = false }, inputNumber: function (number) { if (number === 'C') { @@ -286,18 +279,18 @@ } }, successUpload(response, file, fileList) { - if (response.code === 200) { - this.$message({ - message: response.msg, - type: 'success' - }); - this.files.push(response.data); - } else { - this.$message({ - message: response.msg, - type: 'error' - }); - } + this.$message({ + message: response.msg, + type: 'success' + }); + this.files.push(response.data); + }, + errorUpload(error, file, fileList){ + e = JSON.parse(error.message) + this.$message({ + message: e.detail, + type: 'error' + }); }, qrcodeUrl(file) { return 'https://api.qrserver.com/v1/create-qr-code/?data=' + window.location.origin + '/?code=' + file.code