使用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 -*-
__author__ = "aleimu"
__date__ = "2018.9.28"

# 限制接口短時間調用次數

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)


@app.before_request  # app 裝飾器，會對每一個請求都生效
def before_request():
    ip = request.remote_addr
    ip_count = r.get(ip)
    print("ip: %s, ip_count: %s" % (ip, ip_count))
    if not ip_count:
        r.set(ip, 1)
        r.expire(ip, TIME_LIMIT)
    else:
        r.incr(ip)  # 將 key 中儲存的數字值增一
        if int(ip_count) > IP_LIMIT:
            return jsonify({'code': 401, 'status': "reach the ip limit", 'message': {}})


# 裝飾器 一份鍾內限制訪問10次，本地緩存
def stat_called_time(func):
    limit_times = [10]  # 這是一個技巧，裝飾器內的變量繼承每次調用後的變化，變量就必須設置爲可變類型
    cache = {}

    @wraps(func)
    def _called_time(*args, **kwargs):
        key = func.__name__
        if key in cache.keys():
            [call_times, updatetime] = cache[key]
            if time.time() - updatetime < TIME_LIMIT:
                cache[key][0] += 1
            else:
                cache[key] = [1, time.time()]
        else:
            call_times = 1
            cache[key] = [call_times, time.time()]
        print('調用次數: %s' % cache[key][0])
        print('限制次數: %s' % limit_times[0])
        if cache[key][0] <= limit_times[0]:
            res = func(*args, **kwargs)
            cache[key][1] = time.time()
            return res
        else:
            print("超過調用次數了")
            return jsonify({'code': 401, 'status': "reach the limit", 'message': {}})

    return _called_time


@app.route("/call")
@stat_called_time
def home():
    return jsonify({'code': 200, 'status': "", 'message': {}})


@app.route("/")
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 # 參考