高级技巧：利用Lua编写安全场景的测试数据生成工具

tar zxf openresty-1.19.9.1.tar.gzcd openresty-1.19.9.1/./configure --with-luajit --with-http_iconv_modulemake -j8 && sudo make install

pip install pillow

import stringimport typingfrom optparse import OptionParserimport randomimport os
from PIL import Image, ImageDraw

def generate_jpg(width: int, height: int, output: str) -> None:    """    生成一张随机的 JPG 图片    :param width: 生成的图片的宽度    :param height: 生成的图片的高度    :param output: 输出文件名称    """    img: Image = Image.new("RGB", (width, height))    pixels = img.load()    for x in range(width):        for y in range(height):            r = random.randint(0, 255)            g = random.randint(0, 255)            b = random.randint(0, 255)            pixels[x, y] = (r, g, b)    img.save(output, format="JPEG")    print(f"the generated JPEG image is stored in {output}, file size is {os.stat(output).st_size / 1024} KB")

def generate_png(width: int, height: int, output: str) -> None:    """    生成一张随机的 PNG 图片    :param width: 生成的图片的宽度    :param height: 生成的图片的高度    :param output: 输出文件名称    """    img: Image = Image.new("RGBA", (width, height))    draw: ImageDraw = ImageDraw.Draw(img)    for x in range(width):        for y in range(height):            alpha = random.randint(0, 255)            r = random.randint(0, 255)            g = random.randint(0, 255)            b = random.randint(0, 255)            draw.point((x, y), fill=(r, g, b, alpha))    img.save(output, format="PNG")    print(f"the generated PNG image is stored in {output}, file size is {os.stat(output).st_size / 1024} KB")

def generate_gif(width: int, height: int, num_frames: int, output: str) -> None:    """    生成一张随机的 GIF 图片    :param width: 生成的图片的宽度    :param height: 生成的图片的高度    :param num_frames: 生成的图片的桢数    :param output: 输出文件名称    """    frames: typing.List[Image] = []    for _ in range(num_frames):        # 生成每一帧的随机图像        image = Image.new("RGB", (width, height))        for x in range(width):            for y in range(height):                r = random.randint(0, 255)                g = random.randint(0, 255)                b = random.randint(0, 255)                image.putpixel((x, y), (r, g, b))        # 将当前帧添加到帧列表中        frames.append(image)    # 保存图像    frames[0].save(output, format="GIF", append_images=frames[1:], save_all=True, duration=200, loop=1)    print(f"the generated GIF image is stored in {output}, file size is {os.stat(output).st_size / 1024} KB")

def generate_text(size: int, output: str) -> None:    """    生成特定长度的随机文本    :param size: 生成的随机文本的长度    :param output: 输出文件名称    """    with open(output, "wb") as fd:        current_size: int = size        while current_size > 0:            # 每次生成 4K            batch: int = min(4096, current_size)            fd.write("".join([random.choice(string.printable) for _ in range(batch)]).encode())            current_size -= batch    print(f"the generated text is store in {output}, file size is {os.stat(output).st_size / 1024} KB")

def main() -> None:    parser: OptionParser = OptionParser(usage="python %prog options...")    parser.add_option("-t", "--type", dest="type", default="txt", type=str,                      help="the type of generated file, including jpg, png, gif, txt")    parser.add_option("-w", "--width", dest="width", default=200, type=int,                      help="the width of image, if type is image")    parser.add_option("-H", "--height", dest="height", default=200, type=int,                      help="the height of image, if type is image")    parser.add_option("-s", "--size", dest="size", default=1024, type=int,                      help="the size of generated file, in bytes")    parser.add_option("-o", "--output", dest="output", default="a", type=str,                      help="output file name")    parser.add_option("-n", "--num-frames", dest="num_frames", default=10, type=int,                      help="the frame number of generated GIF image")    options, _ = parser.parse_args()
    _, ext = os.path.splitext(options.output)    if options.type.lower() == "jpg":        if ext not in [".jpg", "jpeg", ".jfif"]:            options.output += ".jpg"        generate_jpg(options.width, options.height, options.output)        return    if options.type.lower() == "png":        if ext not in [".png"]:            options.output += ".png"        generate_png(options.width, options.height, options.output)        return    if options.type.lower() == "gif":        if ext not in [".gif"]:            options.output += ".gif"        generate_gif(options.width, options.height, options.num_frames, options.output)        return    if options.type.lower() == "txt":        if ext not in [".txt"]:            options.output += ".txt"        generate_text(options.size, options.output)

if __name__ == "__main__":    main()

.├── generate_image.py├── nginx.conf└── wrk.lua

mkdir -p static/

mkdir -p logs/

cp /usr/local/openresty/nginx/conf/mime.types .

python3 generate_image.py -t gif -w 100 -H 100 -o static/100x100.gifpython3 generate_image.py -t png -w 100 -H 100 -o static/100x100.pngpython3 generate_image.py -t txt -s 131072 -o static/128k.txt

local counter = 1local threads = {}
function setup(thread)    thread:set("id", counter)    table.insert(threads, thread)    counter = counter + 1end
-- 在 init 中预生成所有请求，在 request 中顺序选择function init(args)    current_index = 0
    -- 保存所有预生成请求的表    pregenerated_requests = {}    table.insert(        pregenerated_requests,        wrk.format(            "GET",            "/path/1",            {                ["X-Predefined-Strategy"] = "png,100x100.png"            }        )    )    table.insert(        pregenerated_requests,        wrk.format(            "POST",            "/path/2",            {                ["X-Predefined-Strategy"] = "gif,100x100.gif",                ["Content-Type"] = "application/x-www-form-urlencoded"            },            "foo=bar&baz=quux"        )    )    table.insert(        pregenerated_requests,        wrk.format(            "PUT",            "/path/3/arbitrary/here",            {                ["X-Predefined-Strategy"] = "text,128k.txt",                ["Content-Type"] = "application/x-www-form-urlencoded"            },            "foo=bar&baz=quux"        )    )    table.insert(        pregenerated_requests,        wrk.format(            "GET",            "/path/4",            {                ["Content-Type"] = "application/json;charset=utf8"            },            [[                {                    "headers": {"x-header-a": "a", "content-type": "text/plain"},                    "status_code": 200,                    "body": "this is a very very simple text body, but it maybe meet some rules."                }]]        )    )end
function request()    current_index = current_index + 1    return pregenerated_requests[current_index%#pregenerated_requests+1]end

worker_processes  auto;
error_log  logs/error.log;error_log  logs/error.log  notice;error_log  logs/error.log  info;
pid        logs/nginx.pid;
events {    worker_connections  4096;}
http {    include       mime.types;    default_type  application/octet-stream;
    log_format  main  '$remote_addr - $remote_user [$time_local] "$request" '                      '$status $body_bytes_sent "$http_referer" '                      '"$http_user_agent" "$http_x_forwarded_for"';    access_log  logs/access.log  main;
    sendfile          on;    tcp_nodelay       on;    keepalive_timeout 15;
    open_file_cache          max=2048 inactive=30s;    open_file_cache_valid    10s;    open_file_cache_min_uses 1;    open_file_cache_errors   on;
    # 请求体不能超过该设置。    # 如果客户端需要传递更大的请求体，那么调整该设置    client_max_body_size 10m;
    # 预置策略。    # 预置策略通过名称进行唯一标识。    # 预置策略封装响应码、响应头和响应体（可选），可避免客户端重复传输这些内容。    # 客户端通过请求头指定使用的策略。    # 如此一来，客户端可以同时自由地定制请求和响应，只多出用于指定所用策略的请求头。    # 服务端使用紧跟在策略名称后面的磁盘文件名称，获取响应体。    # 因此，需要提前将生成的响应体存储到磁盘文件。    # 如果未提供文件名称，并且策略中存在响应体，那么使用策略中的响应体。    # 如果都未提供，那么返回空响应体。    # 请求头类似 X-Predefined-Strategy: png,100x100.png 或 X-Predefined-Strategy: forbidden。    # 注意：    # 1. 预置的响应体不能太大，否则将占用太多的共享内存    # 2. 预置策略名称不能包含 ","    # 3. 预置策略中指定的响应头名称不会被规范化，因此设置时，需要注意    lua_shared_dict predefined_stategies 128m;
    init_by_lua_block {        -- 保存策略        local function save_strategy(name, status_code, headers, body)            if name == nil then return ngx.log(ngx.ERR, "no strategy name provided") end            local t = {                status_code = status_code or ngx.HTTP_OK,                headers = {},                body = body or ""            }            for name, value in pairs(headers or {}) do t.headers[name] = value end
            local cjson = require("cjson.safe")            local j = cjson.encode(t)            if j == nil then return end            local s = ngx.shared.predefined_stategies            local suc, err = s:set(name, j)            if suc then ngx.log(ngx.INFO, "setting strategy " .. name .. " succeeded")            else ngx.log(ngx.ERR, "setting strategy " .. name .. " failed with " .. err) end        end
        -- 按需添加策略        save_strategy("png", ngx.HTTP_OK, {["Content-Type"]="image/png"})        save_strategy("jpeg", ngx.HTTP_OK, {["Content-Type"]="image/jpeg"})        save_strategy("gif", ngx.HTTP_OK, {["Content-Type"]="image/gif"})        save_strategy("ico", ngx.HTTP_OK, {["Content-Type"]="image/x-icon"})        save_strategy("text", ngx.HTTP_OK, {["Content-Type"]="text/plain"})        save_strategy("json", ngx.HTTP_OK, {["Content-Type"]="application/json;charset=utf8"})        save_strategy(            "forbidden_default",            ngx.HTTP_FORBIDDEN,            {["Content-Type"]="text/plain"},            "forbidden"        )        save_strategy(            "notfound_default",            ngx.HTTP_NOT_FOUND,            {["Content-Type"]="text/plain"},            "not found"        )    }
    server {        listen       80;        server_name  localhost;
        location /static {            alias static/;        }
        location / {            content_by_lua_block {                local cjson = require "cjson.safe"
                -- 规范化响应名称                local normalize_header = function(name)                    -- 1. _ 替换成 -                    name = string.gsub(name, "_", "-")                    -- 2. 第一个字符大写                    name = string.gsub(name, "%w", function(m) return string.upper(m) end, 1)                    -- 3. - 后面的字符大写                    return string.gsub(name, "-%w", function(m) return string.upper(m) end)                end
                -- 生成指定长度的随机字符串                local function generate_random_string(length)                    local chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"                    -- 使用当前时间作为随机种子                    math.randomseed(os.time())                    local r = ""                    for _ = 1, length do                        local idx = math.random(1, #chars)                        r = r .. string.sub(chars, idx, idx)                    end                    return r                end
                -- 当指定的策略或路径不存在时，返回该响应码                local strategy_not_found_status = ngx.HTTP_NOT_FOUND
                local strategy                local path                local predefined_strategy_hdr = ngx.var.http_x_predefined_strategy                -- 如果客户端指定策略                if predefined_strategy_hdr ~= nil then                    local start_pos, end_pos = string.find(predefined_strategy_hdr, ",", 1, true)                    if start_pos == nil then                        strategy = predefined_strategy_hdr                    else                        strategy = string.sub(predefined_strategy_hdr, 1, start_pos-1)                        path = string.sub(predefined_strategy_hdr, end_pos+1)                    end                    -- 获取策略                    local s = ngx.shared.predefined_stategies                    strategy = s:get(strategy)                    -- 如果策略不存在，那么返回错误                    if strategy == nil then                        ngx.status = strategy_not_found_status                        ngx.header.content_type = "text/plain"                        ngx.say("the provided stategy not found")                        return ngx.exit(ngx.HTTP_OK)                    end                    -- 根据策略设置响应                    strategy = cjson.decode(strategy)                    ngx.status = strategy.status_code                    for name, value in pairs(strategy.headers or {}) do                        ngx.header[name] = value                    end                    -- 设置响应体                    if path == nil then                        ngx.print(strategy.body)                        return ngx.exit(ngx.HTTP_OK)                    end                    -- 发起子请求，从磁盘获取响应体                    local res = ngx.location.capture("/static".."/"..path)                    if res.status == ngx.HTTP_OK then                        ngx.print(res.body)                    else                        ngx.status = strategy_not_found_status                        ngx.header.content_type = "text/plain"                        ngx.say("the provided path not found")                    end                    return ngx.exit(ngx.HTTP_OK)                end
                -- 非策略模式用于构建期望的响应
                -- 当请求体不合法时，返回该响应码                local bad_request_status_code = ngx.HTTP_BAD_REQUEST
                -- 读取请求体                ngx.req.read_body()                local body_raw = ngx.req.get_body_data()                -- 解析请求体                local body, err = cjson.decode(body_raw)                -- 解析请求体失败                if err ~= nil then                    ngx.status = bad_request_status_code                    ngx.say(err)                    return ngx.exit(ngx.HTTP_OK)                end                -- 设置响应码                local status_code = body["status_code"]                if status_code == nil then status_code = ngx.HTTP_OK end                ngx.status = status_code                -- 设置响应头                for name, value in pairs(body["headers"] or {}) do                    ngx.header[normalize_header(name)] = value                end                -- 设置响应体                if body["body"] then                    -- 使用客户端指定的响应体                    ngx.print(body["body"])                else                     -- 使用客户端指定的随机响应体长度生成响应体。                    -- 指定的长度越长，生成随机字符串的效率越低，请设置合理的长度                    if body["random_body_length"] then                        ngx.print(generate_random_string(tonumber(body["random_body_length"])))                    else                        -- 返回空响应体                        ngx.print("")                    end                end                return ngx.exit(ngx.HTTP_OK)            }        }    }}