mm_os/middleware/waf_ip/index.js

这是超级美眉服务端框架，用于快速构建应用程序。

const os = require('os');
const {
	exec
} = require('child_process');
const platform = os.platform();

/**
 * 获取客户端IP
 * @param {Object} req 请求对象
 * @returns {String} 返回真实IP
 */
function getClientIP(req) {
	var ip = req.headers['x-forwarded-for'] || req.headers['X-Forwarded-For'] || req.headers['x-real-ip'] ||
		req.connection.remoteAddress ||
		req.socket.remoteAddress ||
		req.connection.socket.remoteAddress;
	if (ip && ip.split(',').length > 0) {
		ip = ip.split(',')[0]; // 取第一个IP地址
	}
	return ip;
};

/**
 * 设置黑名单
 * @param {String} ip IP地址
 */
function setting_blacklist(ip) {
	var cmd;
	if (platform == "win32") {
		// window 系统
		cmd =
			`netsh advfirewall firewall add rule name="Blacklist ${ip}" dir=in action=block remoteip="${ip}" protocol=any`
	} else {
		// linux 系统
		cmd = `sudo iptables -A INPUT -s ${ip} -j DROP`;
	}

	exec(cmd, (error, stdout, stderr) => {
		if (error) {
			console.error(`执行的错误: ${error}`);
			return;
		}
		$.log.info(`加入黑名单: ${ip}`);
		if (stderr) {
			console.error(`标准错误输出: ${stderr}`);
		}
	});
}

/**
 * IP防火墙
 * @param {Object} server 服务
 * @param {Object} config 配置参数
 */
module.exports = function(server, config) {
	var limit = config.request_limit || 0;
	var duration = config.request_duration || 0;
	var block = config.request_block || false;

	if (limit && duration) {
		/* WAF（web防火墙） */
		server.use(async (ctx, next) => {
			var pass = true;
			// 获取IP
			var ip = getClientIP(ctx.req);
			var num = 1;
			var now = new Date();
			var date = now.toStr('yyyy-MM-dd');
			var time;
			var str = await $.cache.get("ip_" + ip);
			var json;
			if (str) {
				if (typeof(str) === "string") {
					json = JSON.parse(str);
				} else {
					json = str;
				}
				if (json.date !== date) {
					num = 1;
				} else {
					// 判断时间间隔是否在范围外
					if (json.time.toTime().interval(now) > duration) {
						num = 1;
					} else {
						// 如果是在周期内，访问次数+1，并判断是否超出上限
						num = json.num + 1;
						if (num > limit) {
							// 超出上限禁止访问，并加入黑名单
							pass = false;
							if (block) {
								setting_blacklist(ip);
							}
						}
					}
				}
			}
			if (!time) {
				time = now.toStr('yyyy-MM-dd hh:mm:ss');
			}
			if (pass) {
				await $.cache.set("ip_" + ip, JSON.stringify({
					date,
					time,
					num
				}), duration);
				ctx.request.ip = ip;
				ctx.ip = ip;
				await next();
			} else {
				ctx.status = 429;
				ctx.body = '请求频率过高，请稍后再试。';
			}
		});
	}
	return server;
};