Skip to content

Commit 68881c6

Browse files
veocoveoco
committed
重写用户登录和 IP 检查并重定义状态码
1 parent fb80d26 commit 68881c6

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>
@@ -206,16 +207,12 @@
206207
FileData.append('style', that.destoryData.style);
207208
FileData.append('value', that.destoryData.value);
208209
axios.post('/share', FileData)
209-
.then(function (res) {
210-
if (res.data.code === 200) {
211-
that.$message({'message': res.data.msg, 'type': 'success'});
212-
that.files.push(res.data.data)
213-
} else {
214-
that.$message({'message': res.data.msg, 'type': 'error'});
215-
}
210+
.then(res => {
211+
that.$message({'message': res.data.msg, 'type': 'success'});
212+
that.files.push(res.data.data)
216213
})
217-
.catch(function (error) {
218-
that.$message({'message': error.data.msg, 'type': 'error'});
214+
.catch(e => {
215+
that.$message({'message': e.response.data.detail, 'type': 'error'});
219216
});
220217
});
221218
}
@@ -226,14 +223,14 @@
226223
FileData.append('file', file || '');
227224
FileData.append('style', that.destoryData.style);
228225
FileData.append('value', that.destoryData.value);
229-
axios.post('/share', FileData).then(async res => {
230-
if (res.data.code === 200) {
231-
that.$message({'message': res.data.msg, 'type': 'success'});
232-
that.files.push(res.data.data)
233-
} else {
234-
that.$message({'message': res.data.msg, 'type': 'error'});
235-
}
226+
axios.post('/share', FileData)
227+
.then(res => {
228+
that.$message({'message': res.data.msg, 'type': 'success'});
229+
that.files.push(res.data.data)
236230
})
231+
.catch(e => {
232+
that.$message({'message': e.response.data.detail, 'type': 'error'});
233+
});
237234
}
238235

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

0 commit comments

Comments
 (0)