使用flask_limiter設定API配額

前言

閑來無事，突然想到了以前做過的關於后台API安全方面的事，關於接口訪問配額的設置，flask有沒有很好的庫支持呢？一找還真有！主要是對照了庫的官方文檔自己寫了下dome，以供參考。

# -*- coding:utf-8 -*-

import json
from flask import Flask, jsonify, request
from flask_limiter import Limiter, HEADERS  # https://github.com/alisaifee/flask-limiter
from flask_limiter.util import get_remote_address

# import limits.storage   # https://github.com/alisaifee/limits/tree/master/limits    依賴了這個limits庫

RATELIMIT_STORAGE_URL = "redis://172.16.4.120:6379"  # 將被限制不可以再正常訪問的請求放入緩存
app = Flask(__name__)


@app.after_request
def after_request(response):
    response.headers['Content-Type'] = 'text/html'  # 避免ie8把json數據以下載方式打開
    return response


limiter = Limiter(
    app,
    key_func=get_remote_address,
    default_limits=["200 per day", "50 per hour"], # 全局配置,一般不要配置
    storage_uri=RATELIMIT_STORAGE_URL,
    headers_enabled=True  # X-RateLimit寫入響應頭。
)
"""
root@(none):~# curl -v "http://172.16.2.197:6000/index4"
* Hostname was NOT found in DNS cache
*   Trying 172.16.2.197...
* Connected to 172.16.2.197 (172.16.2.197) port 6000 (#0)
> GET /index4 HTTP/1.1
> User-Agent: curl/7.35.0
> Host: 172.16.2.197:6000
> Accept: */*
> 
* HTTP 1.0, assume close after body
< HTTP/1.0 200 OK
< Content-Type: text/html
< X-RateLimit-Limit: 2
< X-RateLimit-Remaining: 1
< X-RateLimit-Reset: 1539849353
# X-RateLimit-Limit	活動窗口允許的請求總數
# X-RateLimit-Remaining	活動窗口中剩余的請求數。
# X-RateLimit-Reset	自重建窗口以來的UTC時間以來的UTC秒數。
< Retry-After: 3600
< Content-Length: 32
< Server: Werkzeug/0.14.1 Python/2.7.9
< Date: Thu, 18 Oct 2018 06:55:52 GMT
< 
{
  "response": "mysql_limit"
}
"""


@app.route("/slow")
@limiter.limit("1 per day")
def slow():
    return "24"


# curl -v "http://172.16.2.197:6000/slow"
# 127.0.0.1:6379> keys *
# 1) "LIMITER/172.16.4.120/slow/1/1/day"
# 127.0.0.1:6379> TTL "LIMITER/172.16.4.120/slow/1/1/day"
# (integer) 86285       # 設定的失效時間是 1 天

@app.route("/fast")  # 默認訪問速率
def fast():
    return "42"


@app.route("/ping")
@limiter.exempt  # 無訪問速率限制
def ping():
    return "PONG"


@app.route("/")
@limiter.limit("1/second", error_message='chill!')
@limiter.limit("100/day")
@limiter.limit("10/hour")
@limiter.limit("1/minute")
def index():
    return "index"


# curl -v "http://172.16.2.197:6000/" # 訪問了一次就觸發了下面三條redis記錄，清理掉redis記錄，可以一分鍾內再次訪問
# 127.0.0.1:6379> keys *
# 2) "LIMITER/172.16.4.120/index/10/1/hour"
# 3) "LIMITER/172.16.4.120/index/100/1/day"
# 4) "LIMITER/172.16.4.120/index/1/1/minute"


@app.route('/index0')
@limiter.limit("100/30seconds", error_message=json.dumps({"data": "You hit the rate limit", "error": 429}))# 支持設定時間跨度,和報錯信息!
def index0():
    return jsonify({'response': 'This is a rate limited response'})


@app.route('/index2')
@limiter.limit("100/day;10/hour;1/minute")  # 與index()同功效
def index2():
    return jsonify({'response': 'Are we rated limited?'})


@app.route('/index3')
@limiter.exempt
def index3():
    return jsonify({'response': 'We are not rate limited'})


@app.route("/expensive")  # exempt_when=callable 當滿足給定條件時，可以免除限制
@limiter.limit("1/day", exempt_when=lambda: get_remote_address() == "172.16.4.120")
def expensive_route():
    return jsonify({'response': 'you are wellcome!'})


# 多個路由應共享速率限制的情況（例如，訪問有限速率的相同資源保護路由時）
mysql_limit = limiter.shared_limit("2/hour", scope="mysql_flag")
# 3) "LIMITER/172.16.4.120/mysql_flag/2/1/hour"

@app.route('/index4')
@mysql_limit
def index4():
    return jsonify({'response': 'mysql_limit'})


@app.route('/index5')
@mysql_limit
def index5():
    return jsonify({'response': 'mysql_limit'})


def host_scope(endpoint_name):
    return request.host

# 動態共享限制
host_limit = limiter.shared_limit("2/hour", scope=host_scope)
# 1) "LIMITER/172.16.4.120/172.16.2.197:6000/2/1/hour"

@app.route('/index6')
@host_limit
def index6():
    return jsonify({'response': 'host_limit'})


@app.route('/index7')
@host_limit
def index7():
    return jsonify({'response': 'host_limit'})


# @limiter.request_filter這個裝飾器只是將一個函數標記為將要測試速率限制的請求的過濾器。如果任何請求過濾器返回True，
# 則不會對該請求執行速率限制。此機制可用於創建自定義白名單。

@limiter.request_filter
def ip_whitelist():
    return request.remote_addr == "127.0.0.1"


@limiter.request_filter
def header_whitelist():
    return request.headers.get("X-Internal", "") == "true"


if __name__ == "__main__":
    app.run(host='0.0.0.0', port=6000, threaded=True)

DIY簡陋版

# -*- coding:utf-8 -*-
# 限制接口短時間調用次數

import redis
import time
from flask import Flask, jsonify, request
from functools import wraps

REDIS_DB = 0
REDIS_HOST = '172.16.4.120'
REDIS_PORT = 6379
REDIS_PASSWORD = ''
IP_LIMIT = 10
TIME_LIMIT = 60

app = Flask(__name__)
r = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_DB, socket_timeout=3000)

def repeat_intercept(ex=5, ip_path=True):
    """攔截重復請求 ex默認5秒"""

    def decorator(func):
        @wraps(func)
        def function(*args, **kwargs):
            req_str = "intercept:"
            if ip_path:
                req_str += request.remote_addr + request.path   # 只對ip和url做限制
            else:
                req_dict = request.values.to_dict()
                req_str += json.dumps(sorted([(x, y) for x, y in req_dict.items()]))    #對請求的參數也做校驗
            if r.exists(req_str):
                return jsonify({"code": 1404, "errmsg": "請求頻率過高", "data": None})
            else:
                r.set(req_str, 0, ex)
                return func(*args, **kwargs)

        return function

    return decorator

@app.route("/call")
@repeat_intercept(ex=1, ip_path=True)
def home():
    return jsonify({'code': 200, 'status': "", 'message': {}})


@app.route("/")
@repeat_intercept(ex=5, ip_path=False)
def index():
    return jsonify({'code': 200, 'status': "", 'message': {}})



if __name__ == "__main__":
    app.run(host='0.0.0.0', port=5000, threaded=True)

參考文檔

https://flask-limiter.readthedocs.io/en/stable/ # 主要來源
https://github.com/search?l=Python&q=flask_limiter+&type=Code # 參考