Skip to content

Commit 31f67fa

Browse files
vastsavastsa
authored
Merge pull request #10 from veoco/master
重写用户登录和 IP 检查并重定义状态码
2 parents 27fc665 + 68881c6 commit 31f67fa

File tree

4 files changed

+116
-122
lines changed

4 files changed

+116
-122
lines changed

‎depends.py‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
from typing import Union
2+
from datetime import datetime, timedelta
3+
4+
from fastapi import Header, HTTPException, Request
5+
6+
import settings
7+
8+
9+
async def admin_required(pwd: Union[str, None] = Header(default=None)):
10+
if not pwd or pwd != settings.ADMIN_PASSWORD:
11+
raise HTTPException(status_code=401, detail="密码错误")
12+
13+
14+
class IPRateLimit:
15+
ips = {}
16+
17+
def check_ip(self, ip):
18+
# 检查ip是否被禁止
19+
if ip in self.ips:
20+
if self.ips[ip]['count'] >= settings.ERROR_COUNT:
21+
if self.ips[ip]['time'] + timedelta(minutes=settings.ERROR_MINUTE) > datetime.now():
22+
return False
23+
else:
24+
self.ips.pop(ip)
25+
return True
26+
27+
def add_ip(cls, ip):
28+
ip_info = cls.ips.get(ip, {'count': 0, 'time': datetime.now()})
29+
ip_info['count'] += 1
30+
cls.ips[ip] = ip_info
31+
return ip_info['count']
32+
33+
def __call__(self, request: Request):
34+
ip = request.client.host
35+
if not self.check_ip(ip):
36+
raise HTTPException(status_code=400, detail="错误次数过多，请稍后再试")
37+
return ip

‎main.py‎

Lines changed: 29 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,7 @@
44
import asyncio
55
from pathlib import Path
66

7-
from fastapi import FastAPI, Depends, UploadFile, Form, File
8-
from starlette.requests import Request
7+
from fastapi import FastAPI, Depends, UploadFile, Form, File, HTTPException
98
from starlette.responses import HTMLResponse, FileResponse
109
from starlette.staticfiles import StaticFiles
1110

@@ -15,6 +14,7 @@
1514
import settings
1615
from database import get_session, Codes, init_models, engine
1716
from storage import STORAGE_ENGINE
17+
from depends import admin_required, IPRateLimit
1818

1919
app = FastAPI(debug=settings.DEBUG)
2020

@@ -43,7 +43,7 @@ async def startup():
4343
.replace('{{description}}', settings.DESCRIPTION) \
4444
.replace('{{keywords}}', settings.KEYWORDS)
4545

46-
error_ip_count = {}
46+
ip_limit = IPRateLimit()
4747

4848

4949
async def delete_expire_files():
@@ -72,88 +72,60 @@ async def admin():
7272
return HTMLResponse(admin_html)
7373

7474

75-
@app.post(f'/{settings.ADMIN_ADDRESS}')
76-
async def admin_post(request: Request, s: AsyncSession = Depends(get_session)):
77-
if request.headers.get('pwd') == settings.ADMIN_PASSWORD:
78-
query = select(Codes)
79-
codes = (await s.execute(query)).scalars().all()
80-
return {'code': 200, 'msg': '查询成功', 'data': codes}
81-
else:
82-
return {'code': 404, 'msg': '密码错误'}
75+
@app.post(f'/{settings.ADMIN_ADDRESS}', dependencies=[Depends(admin_required)])
76+
async def admin_post(s: AsyncSession = Depends(get_session)):
77+
query = select(Codes)
78+
codes = (await s.execute(query)).scalars().all()
79+
return {'msg': '查询成功', 'data': codes}
8380

8481

85-
@app.delete(f'/{settings.ADMIN_ADDRESS}')
86-
async def admin_delete(request: Request, code: str, s: AsyncSession = Depends(get_session)):
87-
if request.headers.get('pwd') == settings.ADMIN_PASSWORD:
88-
query = select(Codes).where(Codes.code == code)
89-
file = (await s.execute(query)).scalars().first()
90-
await storage.delete_file({'type': file.type, 'text': file.text})
91-
await s.delete(file)
92-
await s.commit()
93-
return {'code': 200, 'msg': '删除成功'}
94-
else:
95-
return {'code': 404, 'msg': '密码错误'}
82+
@app.delete(f'/{settings.ADMIN_ADDRESS}', dependencies=[Depends(admin_required)])
83+
async def admin_delete(code: str, s: AsyncSession = Depends(get_session)):
84+
query = select(Codes).where(Codes.code == code)
85+
file = (await s.execute(query)).scalars().first()
86+
await storage.delete_file({'type': file.type, 'text': file.text})
87+
await s.delete(file)
88+
await s.commit()
89+
return {'msg': '删除成功'}
9690

9791

9892
@app.get('/')
9993
async def index():
10094
return HTMLResponse(index_html)
10195

10296

103-
def check_ip(ip):
104-
# 检查ip是否被禁止
105-
if ip in error_ip_count:
106-
if error_ip_count[ip]['count'] >= settings.ERROR_COUNT:
107-
if error_ip_count[ip]['time'] + datetime.timedelta(minutes=settings.ERROR_MINUTE) > datetime.datetime.now():
108-
return False
109-
else:
110-
error_ip_count.pop(ip)
111-
return True
112-
113-
114-
def ip_error(ip):
115-
ip_info = error_ip_count.get(ip, {'count': 0, 'time': datetime.datetime.now()})
116-
ip_info['count'] += 1
117-
error_ip_count[ip] = ip_info
118-
return ip_info['count']
119-
120-
12197
@app.get('/select')
12298
async def get_file(code: str, s: AsyncSession = Depends(get_session)):
12399
query = select(Codes).where(Codes.code == code)
124100
info = (await s.execute(query)).scalars().first()
125-
if info:
126-
if info.type == 'text':
127-
return {'code': code, 'msg': '查询成功', 'data': info.text}
128-
else:
129-
filepath = await storage.get_filepath(info.text)
130-
return FileResponse(filepath, filename=info.name)
101+
if not info:
102+
raise HTTPException(status_code=404, detail="口令不存在")
103+
if info.type == 'text':
104+
return {'msg': '查询成功', 'data': info.text}
131105
else:
132-
return {'code': 404, 'msg': '口令不存在'}
106+
filepath = await storage.get_filepath(info.text)
107+
return FileResponse(filepath, filename=info.name)
133108

134109

135110
@app.post('/')
136-
async def index(request: Request, code: str, s: AsyncSession = Depends(get_session)):
137-
ip = request.client.host
138-
if not check_ip(ip):
139-
return {'code': 404, 'msg': '错误次数过多，请稍后再试'}
111+
async def index(code: str, ip: str = Depends(ip_limit), s: AsyncSession = Depends(get_session)):
140112
query = select(Codes).where(Codes.code == code)
141113
info = (await s.execute(query)).scalars().first()
142114
if not info:
143-
return {'code': 404, 'msg': f'取件码错误，错误{settings.ERROR_COUNT - ip_error(ip)}次将被禁止10分钟'}
115+
error_count = ip_limit.add_ip(ip)
116+
raise HTTPException(status_code=404, detail=f"取件码错误，错误{settings.ERROR_COUNT - error_count}次将被禁止10分钟")
144117
if info.exp_time < datetime.datetime.now() or info.count == 0:
145118
await storage.delete_file({'type': info.type, 'text': info.text})
146119
await s.delete(info)
147120
await s.commit()
148-
return {'code': 404, 'msg': '取件码已过期，请联系寄件人'}
121+
raise HTTPException(status_code=404, detail="取件码已过期，请联系寄件人")
149122
count = info.count - 1
150123
query = update(Codes).where(Codes.id == info.id).values(count=count)
151124
await s.execute(query)
152125
await s.commit()
153126
if info.type != 'text':
154127
info.text = f'/select?code={code}'
155128
return {
156-
'code': 200,
157129
'msg': '取件成功，请点击"取"查看',
158130
'data': {'type': info.type, 'text': info.text, 'name': info.name, 'code': info.code}
159131
}
@@ -165,12 +137,12 @@ async def share(text: str = Form(default=None), style: str = Form(default='2'),
165137
code = await get_code(s)
166138
if style == '2':
167139
if value > 7:
168-
return {'code': 404, 'msg': '最大有效天数为7天'}
140+
raise HTTPException(status_code=400, detail="最大有效天数为7天")
169141
exp_time = datetime.datetime.now() + datetime.timedelta(days=value)
170142
exp_count = -1
171143
elif style == '1':
172144
if value < 1:
173-
return {'code': 404, 'msg': '最小有效次数为1次'}
145+
raise HTTPException(status_code=400, detail="最小有效次数为1次")
174146
exp_time = datetime.datetime.now() + datetime.timedelta(days=1)
175147
exp_count = value
176148
else:
@@ -181,7 +153,7 @@ async def share(text: str = Form(default=None), style: str = Form(default='2'),
181153
file_bytes = await file.read()
182154
size = len(file_bytes)
183155
if size > settings.FILE_SIZE_LIMIT:
184-
return {'code': 404, 'msg': '文件过大'}
156+
raise HTTPException(status_code=400, detail="文件过大")
185157
_text, _type, name = await storage.save_file(file, file_bytes, key), file.content_type, file.filename
186158
else:
187159
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'),
198170
s.add(info)
199171
await s.commit()
200172
return {
201-
'code': 200,
202173
'msg': '分享成功，请点击文件箱查看取件码',
203174
'data': {'code': code, 'key': key, 'name': name, 'text': _text}
204175
}

‎templates/admin.html‎

Lines changed: 12 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -87,28 +87,21 @@
8787
methods: {
8888
loginAdmin: function () {
8989
axios.post('', {}, {'headers': {'pwd': this.pwd}}).then(res => {
90-
if (res.data.code === 200) {
91-
this.files = res.data.data;
92-
this.login = true;
93-
localStorage.setItem('pwd', this.pwd);
94-
} else {
95-
localStorage.clear()
96-
this.$message({
97-
message: res.data.msg,
98-
type: 'error'
99-
});
100-
}
101-
})
90+
this.files = res.data.data;
91+
this.login = true;
92+
localStorage.setItem('pwd', this.pwd);
93+
}).catch(e => {
94+
localStorage.clear()
95+
this.$message({'message': e.response.data.detail, 'type': 'error'});
96+
});
10297
},
10398
deleteFile: function (code) {
10499
axios.delete('?code=' + code, {'headers': {'pwd': this.pwd}}).then(res => {
105-
if (res.data.code === 200) {
106-
this.files = this.files.filter(item => item.code !== code)
107-
this.$message({
108-
message: res.data.msg,
109-
type: 'success'
110-
});
111-
}
100+
this.files = this.files.filter(item => item.code !== code)
101+
this.$message({
102+
message: res.data.msg,
103+
type: 'success'
104+
});
112105
})
113106
},
114107
}

‎templates/index.html‎

Lines changed: 38 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,7 @@
9494
multiple
9595
:data="destoryData"
9696
:on-success="successUpload"
97+
:on-error="errorUpload"
9798
>
9899
<i class="el-icon-upload"></i>
99100
<div class="el-upload__text">将文字、文件拖、粘贴到此处，或<em>点击上传</em></div>
@@ -205,16 +206,12 @@
205206
FileData.append('style', that.destoryData.style);
206207
FileData.append('value', that.destoryData.value);
207208
axios.post('/share', FileData)
208-
.then(function (res) {
209-
if (res.data.code === 200) {
210-
that.$message({'message': res.data.msg, 'type': 'success'});
211-
that.files.push(res.data.data)
212-
} else {
213-
that.$message({'message': res.data.msg, 'type': 'error'});
214-
}
209+
.then(res => {
210+
that.$message({'message': res.data.msg, 'type': 'success'});
211+
that.files.push(res.data.data)
215212
})
216-
.catch(function (error) {
217-
that.$message({'message': error.data.msg, 'type': 'error'});
213+
.catch(e => {
214+
that.$message({'message': e.response.data.detail, 'type': 'error'});
218215
});
219216
});
220217
}
@@ -225,14 +222,14 @@
225222
FileData.append('file', file || '');
226223
FileData.append('style', that.destoryData.style);
227224
FileData.append('value', that.destoryData.value);
228-
axios.post('/share', FileData).then(async res => {
229-
if (res.data.code === 200) {
230-
that.$message({'message': res.data.msg, 'type': 'success'});
231-
that.files.push(res.data.data)
232-
} else {
233-
that.$message({'message': res.data.msg, 'type': 'error'});
234-
}
225+
axios.post('/share', FileData)
226+
.then(res => {
227+
that.$message({'message': res.data.msg, 'type': 'success'});
228+
that.files.push(res.data.data)
235229
})
230+
.catch(e => {
231+
that.$message({'message': e.response.data.detail, 'type': 'error'});
232+
});
236233
}
237234

238235
}
@@ -249,25 +246,21 @@
249246
},
250247
get_file: function () {
251248
const that = this;
252-
axios.post('?code=' + this.code).then(function (response) {
253-
if (response.data.code === 404) {
254-
that.$message({
255-
message: response.data.msg,
256-
type: 'error'
257-
});
258-
} else {
259-
that.files.push(response.data.data);
260-
that.$message({
261-
message: response.data.msg,
262-
type: 'success'
263-
});
264-
that.quDrawer = true;
265-
}
249+
axios.post('?code=' + this.code).then(response => {
250+
that.files.push(response.data.data);
251+
that.$message({
252+
message: response.data.msg,
253+
type: 'success'
254+
});
255+
that.quDrawer = true;
266256
that.code = '';
267-
that.inout_disable = false
268-
}).catch(function (error) {
269-
console.log(error);
257+
}).catch(e => {
258+
that.$message({
259+
message: e.response.data.detail,
260+
type: 'error'
261+
});
270262
});
263+
that.inout_disable = false
271264
},
272265
inputNumber: function (number) {
273266
if (number === 'C') {
@@ -285,18 +278,18 @@
285278
}
286279
},
287280
successUpload(response, file, fileList) {
288-
if (response.code === 200) {
289-
this.$message({
290-
message: response.msg,
291-
type: 'success'
292-
});
293-
this.files.push(response.data);
294-
} else {
295-
this.$message({
296-
message: response.msg,
297-
type: 'error'
298-
});
299-
}
281+
this.$message({
282+
message: response.msg,
283+
type: 'success'
284+
});
285+
this.files.push(response.data);
286+
},
287+
errorUpload(error, file, fileList){
288+
e = JSON.parse(error.message)
289+
this.$message({
290+
message: e.detail,
291+
type: 'error'
292+
});
300293
},
301294
qrcodeUrl(file) {
302295
return 'https://api.qrserver.com/v1/create-qr-code/?data=' + window.location.origin + '/?code=' + file.code

0 commit comments

Comments
 (0)