Merge pull request #10 from veoco/master

重写用户登录和 IP 检查并重定义状态码

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

main.py

@@ -4,8 +4,7 @@ import random
 import asyncio
 from pathlib import Path
 
-from fastapi import FastAPI, Depends, UploadFile, Form, File
+from fastapi import FastAPI, Depends, UploadFile, Form, File, HTTPException
-from starlette.requests import Request
 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}')
+@app.post(f'/{settings.ADMIN_ADDRESS}', dependencies=[Depends(admin_required)])
-async def admin_post(request: Request, s: AsyncSession = Depends(get_session)):
+async def admin_post(s: AsyncSession = Depends(get_session)):
-    if request.headers.get('pwd') == settings.ADMIN_PASSWORD:
+    query = select(Codes)
-        query = select(Codes)
+    codes = (await s.execute(query)).scalars().all()
-        codes = (await s.execute(query)).scalars().all()
+    return {'msg': '查询成功', 'data': codes}
-        return {'code': 200, 'msg': '查询成功', 'data': codes}
-    else:
-        return {'code': 404, 'msg': '密码错误'}
 
 
-@app.delete(f'/{settings.ADMIN_ADDRESS}')
+@app.delete(f'/{settings.ADMIN_ADDRESS}', dependencies=[Depends(admin_required)])
-async def admin_delete(request: Request, code: str, s: AsyncSession = Depends(get_session)):
+async def admin_delete(code: str, s: AsyncSession = Depends(get_session)):
-    if request.headers.get('pwd') == settings.ADMIN_PASSWORD:
+    query = select(Codes).where(Codes.code == code)
-        query = select(Codes).where(Codes.code == code)
+    file = (await s.execute(query)).scalars().first()
-        file = (await s.execute(query)).scalars().first()
+    await storage.delete_file({'type': file.type, 'text': file.text})
-        await storage.delete_file({'type': file.type, 'text': file.text})
+    await s.delete(file)
-        await s.delete(file)
+    await s.commit()
-        await s.commit()
+    return {'msg': '删除成功'}
-        return {'code': 200, 'msg': '删除成功'}
-    else:
-        return {'code': 404, '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 not info:
-        if info.type == 'text':
+        raise HTTPException(status_code=404, detail="口令不存在")
-            return {'code': code, 'msg': '查询成功', 'data': info.text}
+    if info.type == 'text':
-        else:
+        return {'msg': '查询成功', 'data': info.text}
-            filepath = await storage.get_filepath(info.text)
-            return FileResponse(filepath, filename=info.name)
     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)):
+async def index(code: str, ip: str = Depends(ip_limit), s: AsyncSession = Depends(get_session)):
-    ip = request.client.host
-    if not check_ip(ip):
-        return {'code': 404, 'msg': '错误次数过多，请稍后再试'}
     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}
     }

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.files = res.data.data;
+                    this.login = true;
-                        this.login = true;
+                    localStorage.setItem('pwd', this.pwd);
-                        localStorage.setItem('pwd', this.pwd);
+                }).catch(e => {
-                    } else {
+                    localStorage.clear()
-                        localStorage.clear()
+                    this.$message({'message': e.response.data.detail, 'type': 'error'});
-                        this.$message({
+                });
-                            message: res.data.msg,
-                            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.files = this.files.filter(item => item.code !== code)
+                    this.$message({
-                        this.$message({
+                        message: res.data.msg,
-                            message: res.data.msg,
+                        type: 'success'
-                            type: 'success'
+                    });
-                        });
-                    }
                 })
             },
         }

templates/index.html

@@ -94,6 +94,7 @@
                         multiple
                         :data="destoryData"
                         :on-success="successUpload"
+                        :on-error="errorUpload"
                 >
                     <i class="el-icon-upload"></i>
                     <div class="el-upload__text">将文字、文件拖、粘贴到此处，或<em>点击上传</em></div>
@@ -205,16 +206,12 @@
                                         FileData.append('style', that.destoryData.style);
                                         FileData.append('value', that.destoryData.value);
                                         axios.post('/share', FileData)
-                                                .then(function (res) {
+                                                .then(res => {
-                                                    if (res.data.code === 200) {
+                                                    that.$message({'message': res.data.msg, 'type': 'success'});
-                                                        that.$message({'message': res.data.msg, 'type': 'success'});
+                                                    that.files.push(res.data.data)
-                                                        that.files.push(res.data.data)
-                                                    } else {
-                                                        that.$message({'message': res.data.msg, 'type': 'error'});
-                                                    }
                                                 })
-                                                .catch(function (error) {
+                                                .catch(e => {
-                                                    that.$message({'message': error.data.msg, 'type': 'error'});
+                                                    that.$message({'message': e.response.data.detail, 'type': 'error'});
                                                 });
                                     });
                                 }
@@ -225,14 +222,14 @@
                                 FileData.append('file', file || '');
                                 FileData.append('style', that.destoryData.style);
                                 FileData.append('value', that.destoryData.value);
-                                axios.post('/share', FileData).then(async res => {
+                                axios.post('/share', FileData)
-                                    if (res.data.code === 200) {
+                                .then(res => {
-                                        that.$message({'message': res.data.msg, 'type': 'success'});
+                                    that.$message({'message': res.data.msg, 'type': 'success'});
-                                        that.files.push(res.data.data)
+                                    that.files.push(res.data.data)
-                                    } else {
-                                        that.$message({'message': res.data.msg, 'type': 'error'});
-                                    }
                                 })
+                                .catch(e => {
+                                    that.$message({'message': e.response.data.detail, 'type': 'error'});
+                                });
                             }
 
                         }
@@ -249,25 +246,21 @@
             },
             get_file: function () {
                 const that = this;
-                axios.post('?code=' + this.code).then(function (response) {
+                axios.post('?code=' + this.code).then(response => {
-                    if (response.data.code === 404) {
+                    that.files.push(response.data.data);
-                        that.$message({
+                    that.$message({
-                            message: response.data.msg,
+                        message: response.data.msg,
-                            type: 'error'
+                        type: 'success'
-                        });
+                    });
-                    } else {
+                    that.quDrawer = true;
-                        that.files.push(response.data.data);
-                        that.$message({
-                            message: response.data.msg,
-                            type: 'success'
-                        });
-                        that.quDrawer = true;
-                    }
                     that.code = '';
-                    that.inout_disable = false
+                }).catch(e => {
-                }).catch(function (error) {
+                    that.$message({
-                    console.log(error);
+                        message: e.response.data.detail,
+                        type: 'error'
+                    });
                 });
+                that.inout_disable = false
             },
             inputNumber: function (number) {
                 if (number === 'C') {
@@ -285,18 +278,18 @@
                 }
             },
             successUpload(response, file, fileList) {
-                if (response.code === 200) {
+                this.$message({
-                    this.$message({
+                    message: response.msg,
-                        message: response.msg,
+                    type: 'success'
-                        type: 'success'
+                });
-                    });
+                this.files.push(response.data);
-                    this.files.push(response.data);
+            },
-                } else {
+            errorUpload(error, file, fileList){
-                    this.$message({
+                e = JSON.parse(error.message)
-                        message: response.msg,
+                this.$message({
-                        type: 'error'
+                    message: e.detail,
-                    });
+                    type: 'error'
-                }
+                });
             },
             qrcodeUrl(file) {
                 return 'https://api.qrserver.com/v1/create-qr-code/?data=' + window.location.origin + '/?code=' + file.code