mm_os/middleware/waf_ip/index.js

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

const {
	exec
} = require('child_process');
const platform = require('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) => {
			try {
				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 json;
				try {
					var str = await $.cache.get("ip_" + ip);
					if (str) {
						if (typeof(str) === "string") {
							try {
								json = JSON.parse(str);
							} catch (jsonError) {
								$.log.error('WAF IP中间件JSON解析错误:', jsonError);
								json = null;
							}
						} else {
							json = str;
						}
						if (json) {
							try {
								if (json.date !== date) {
									num = 1;
								} else {
									// 判断时间间隔是否在范围外
									// 安全地处理json.time，无论它是字符串还是对象
									if (typeof json.time === 'string') {
										// 如果json.time是字符串，尝试解析它
										const savedTime = new Date(json.time);
										if (!isNaN(savedTime.getTime()) && (now - savedTime) > duration) {
											num = 1;
										} else {
											// 如果是在周期内，访问次数+1，并判断是否超出上限
											num = json.num + 1;
											if (num > limit) {
												// 超出上限禁止访问，并加入黑名单
												pass = false;
												if (block) {
													setting_blacklist(ip);
												}
											}
										}
									} else if (json.time && json.time.toTime && typeof json.time.toTime().interval === 'function') {
										// 原有逻辑，处理对象类型的时间
										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);
												}
											}
										}
									} else {
										// 默认情况：如果时间格式不正确，重置计数
										num = 1;
									}
								}
							} catch (timeError) {
								$.log.error('WAF IP中间件时间处理错误:', timeError);
								// 出错时重置计数以保证安全
								num = 1;
							}
						}
					}
				} catch (cacheError) {
					$.log.error('WAF IP中间件缓存操作错误:', cacheError);
					// 缓存出错时，默认允许请求通过
				}
				if (!time) {
					time = now.toStr('yyyy-MM-dd hh:mm:ss');
				}
				if (pass) {
					try {
						await $.cache.set("ip_" + ip, JSON.stringify({
							date,
							time,
							num
						}), duration);
					} catch (setCacheError) {
						$.log.error('WAF IP中间件缓存设置错误:', setCacheError);
						// 缓存设置失败不影响请求处理
					}
					ctx.request.ip = ip;
					ctx.ip = ip;
					await next();
				} else {
					ctx.status = 429;
					ctx.body = '请求频率过高，请稍后再试。';
				}
			} catch (error) {
				$.log.error('WAF IP中间件错误:', error);
				// 出错时默认允许请求通过
				ctx.ip = getClientIP(ctx.req);
				await next();
			}
		});
	}
	return server;
};