保護 API 被濫用或避免駭客攻擊,也等於保護好您的錢包。不久前看 Twtter 有人分享因為被駭客阻斷攻擊導致,服務放在 AWS S3 沒有開啟網頁防火牆,在很短的時間內被攻擊了幾百萬次,導致這惡意攻擊的流量變成了昂貴的帳單需要支付,基本的保護需要做好做滿,為我們的 API 加上存取限制。採用 flask_limiter 簡單加上存取限制。
在 app 目錄新增 limiter_flask.py
檔案,將儲存設定為記憶體,這個範例不要使用在正式產品中,正式產品需要改用 為儲存媒介。
# z5525/app/limiter_flask.py
from flask_limiter import Limiter
from flask_limiter.util import get_remote_address
# 這是測試使用請勿在正式產品使用
limiter = Limiter(
key_func=get_remote_address,
default_limits=["10 per minute", "1 per second"], # 每分鐘最多 10 次和每秒最多 1次請求
storage_uri="memory://",
)
在 init.py 加入 limiter_flask
# z5525/app/__init__.py
from z5525.app.limiter_flask import limiter
在 main.py 的 create_app
中使用,匯入 以及使用
# z5525/app/main.py
# 省略
from z5525.app import (
limiter
)
# 省略
def create_app():
app = Flask(__name__)
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///:memory:'
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
app.json.compact = False
db.init_app(app)
api.init_app(app)
limiter.init_app(app) # 只要加入這一行
# 每天最多 1 次請求
@app.route("/api/day")
@limiter.limit("1 per day")
def day():
return {"message": "1 per day"}
# 每分鐘最多 10 次和每秒最多 1 次請求
@app.route("/api/second")
def second():
return {"message": "1 per second"}
# 不受限制的 API
@app.route("/api/ping")
@limiter.exempt
def ping():
return {"message": "PONG"}
return app
# 省略
在 __init__.py
檔案中修改,加入一個 TestResource
,只接受 get 並且限制每分鐘最多 5 次請求。
# z5525/handlers/__init__.py
from flask import Flask, request, jsonify
from flask_restful import Api, Resource
from z5525.handlers.users import UserResource
from z5525.app import (
limiter, # 不要忘記匯入
)
APP_PREFIX = "/api"
limit_5 = limiter.shared_limit("5 per minute", scope="five_scope")
limit_3 = limiter.shared_limit("3 per minute", scope="three_scope")
class TestResource(Resource):
@limiter_5
def get(self):
return jsonify({"message": "Success"})
api = Api()
# 上一篇完成的 UserResource
api.add_resource(UserResource, f'{APP_PREFIX}/users', f'{APP_PREFIX}/users/<int:user_id>')
api.add_resource(TestResource, f'{APP_PREFIX}/five', endpoint="five_scope")
範例中使用兩種呼叫 flask_limiter 的方法,兩種方法都會替代掉內定值,最後測試的時候都需要達成該條件才算正常。
1.flask + flask_limiter
沒有使用 flask_restful 就依照官網的說明設定即可使用。
# flask + flask_limiter
# 加入不同條件限制 API , 以及不受限制的 API 使用 `@limiter.exempt`
# 每天最多 1 次請求
@app.route("/api/day")
@limiter.limit("1 per day")
def day():
return {"message": "1 per day"}
2.flask_restful + flask_limiter
當使用了 flask_restful 程式庫,必須將使用 來完成設定,使用官網的方法是不會生效的。注意在 flask_restful 一定要使用這種方法才能使用。
# flask_restful + flask_limiter
limit_5 = limiter.shared_limit("5 per minute", scope="five_scope")
class TestResource(Resource):
@limiter_5
def get(self):
return jsonify({"message": "Success"})
api.add_resource(TestResource, '/api/five', endpoint="five_scope")
原本的 API 完成後,只要小修改兩個步驟,就可以限制 API 的使用。開始測試執行 python -m z5525.app.main
透過瀏覽器測試 http://127.0.0.1:5525/api/five
或是 curl http://127.0.0.1:5525/api/five
curl http://127.0.0.1:5525/api/five
{"message": "Success"}
curl http://127.0.0.1:5525/api/five
{"message": "5 per 1 minute"}
在目錄 新增一個測試檔案命名為 test_limiting.py
# tests/test_limiting.py
import pytest
import os
from flask import Flask
from z5525.app.main import create_app
from z5525.handlers import api, UserResource
from z5525.models import db, User
APP_PREFIX = "/api"
@pytest.fixture
def client():
""" 建立 Flask 測試客戶端 """
app = create_app()
with app.app_context():
db.create_all()
yield app.test_client()
@pytest.fixture
def sample_user():
""" 插入一筆測試產品 """
user = User(
username="Test user",
email="[email protected]",
password="pass",
)
db.session.add(user)
db.session.commit()
return user
def test_limiter_day(client):
""" 測試 limiter API,並檢查 JSON 格式 """
response = client.get(f'{APP_PREFIX}/day')
assert response.status_code == 200
assert "message" in response.json
response = client.get(f'{APP_PREFIX}/day')
assert response.status_code == 429
assert response.data == b'<!doctype html>\n<html lang=en>\n<title>429 Too Many Requests</title>\n<h1>Too Many Requests</h1>\n<p>1 per 1 day</p>\n'
def test_limiter_five(client, sample_user):
""" 測試 limiter API,並檢查 JSON 格式 """
response = client.get(f'{APP_PREFIX}/five')
assert response.status_code == 200
assert "message" in response.json
response = client.get(f'{APP_PREFIX}/five')
assert response.status_code == 200
assert "message" in response.json
response = client.get(f'{APP_PREFIX}/five')
assert response.status_code == 200
response = client.get(f'{APP_PREFIX}/five')
assert response.status_code == 200
response = client.get(f'{APP_PREFIX}/five')
# 429 TOO MANY REQUESTS
assert response.status_code == 429
assert "message" in response.json
assert "5 per 1 minute" == response.json["message"]
# 可以請求未達到限制的 API
response = client.get(f'{APP_PREFIX}/users')
assert response.status_code == 200
以上需要先完成主機端 Flask REST APIs 設計的所有片段程式,最後再加上本篇的小修改才可以使用。
Flask
flask_restful
flask_sqlalchemy
Flask-Limiter
python 3.11
flask 3.0.3
flask_restful 0.3.10
flask_sqlalchemy 3.1.1
Flask-Limiter 3.10.0
pytest 8.4.1