前置p wang 验签脚本
2024-08-19 本文已影响0人
Leoguo小哥
import org.apache.jmeter.protocol.http.sampler.HTTPSamplerProxy;
import org.apache.jmeter.protocol.http.control.HeaderManager;
import org.apache.jmeter.protocol.http.util.HTTPArgument;
import org.apache.jmeter.protocol.http.control.Header;
import java.net.URLEncoder;
import java.util.ArrayList;
import java.util.Collections;
import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import org.apache.commons.codec.binary.Base64;
import java.net.URL;
class ParamUtils {
ArrayList<String> values = new ArrayList<String>();
long timestamp;
ParamUtils(long timestamp) {
this.timestamp = timestamp;
}
void put(String k, String v) throws Exception {
if (k.isEmpty()) {
return;
}
String value = URLEncoder.encode(v, "UTF-8");
values.add(k + "=" + value);
}
ArrayList<String> sortedValues() {
Collections.sort(values);
return values;
}
void addGetParams(ArrayList<String> params) throws Exception {
for (String param : params) {
String[] keyValue = param.split("=", 2);
put(keyValue[0], keyValue.length > 1 ? keyValue[1] : "");
}
sortedValues();
}
String getParams(String requestMethod, String requestBody, String queryString) throws Exception {
if ("GET".equals(requestMethod)) {
ArrayList<String> param = new ArrayList<String>();
if (queryString != null && !queryString.isEmpty()) {
String[] pairs = queryString.split("&");
for (String pair : pairs) {
int idx = pair.indexOf("=");
put(URLDecoder.decode(pair.substring(0, idx), "UTF-8"), idx > 0 ? URLDecoder.decode(pair.substring(idx + 1), "UTF-8") : "");
}
}
put("signTimestamp", String.valueOf(timestamp));
addGetParams(param);
}else if (Arrays.asList("POST", "PUT", "DELETE").contains(requestMethod)) {
if (!requestBody.isEmpty()) {
values.add("requestBody=" + requestBody);
}
put("signTimestamp", String.valueOf(timestamp));
}
return String.join("&", values);
}
}
if (sampler instanceof HTTPSamplerProxy) {
HTTPSamplerProxy currentSampler = (HTTPSamplerProxy) sampler;
String url = currentSampler.getUrl().toString();
log.info("获取的完整请求: " + url);
URL all_url_Path = new URL(url);
String urlPath = all_url_Path.getPath();
//String urlPath = currentSampler.getPath();
log.info("获取的 url部分: " + urlPath);
String apiKey = vars.get("api_key");
log.info("获取的api_key: " + apiKey);
String secretKey = vars.get("api_secret");
log.info("获取的api_secret: " + secretKey);
String requestMethod = currentSampler.getMethod();
log.info("获取的 requestMethod: " + requestMethod);
long timestamp = System.currentTimeMillis();
String actualRequestBody = "";
if (!"GET".equals(requestMethod)) {
if (currentSampler.getPostBodyRaw()) {
actualRequestBody = currentSampler.getArguments().getArgument(0).getValue();
} else {
StringBuilder requestBodyBuilder = new StringBuilder();
for (HTTPArgument arg : (ArrayList<HTTPArgument>) currentSampler.getArguments().getArguments()) {
if (requestBodyBuilder.length() > 0) {
requestBodyBuilder.append("&");
}
requestBodyBuilder.append(URLEncoder.encode(arg.getName(), "UTF-8"));
requestBodyBuilder.append("=");
requestBodyBuilder.append(URLEncoder.encode(arg.getValue(), "UTF-8"));
}
actualRequestBody = requestBodyBuilder.toString();
}
}
ParamUtils paramUtils = new ParamUtils(timestamp);
//String paramValue = paramUtils.getParams(requestMethod, actualRequestBody);
// 获取查询字符串
String queryString = all_url_Path.getQuery();
// 修改getParams方法的调用,添加queryString参数
String paramValue = paramUtils.getParams(requestMethod, actualRequestBody, queryString);
String payload = requestMethod.toUpperCase() + "\n" + urlPath + "\n" + paramValue;
try {
Mac sha256_HMAC = Mac.getInstance("HmacSHA256");
SecretKeySpec secret_key = new SecretKeySpec(secretKey.getBytes("UTF-8"), "HmacSHA256");
sha256_HMAC.init(secret_key);
String signature = Base64.encodeBase64String(sha256_HMAC.doFinal(payload.getBytes("UTF-8")));
vars.put("signature", signature);
log.info("payload: " + payload);
log.info("signature: " + signature);
// 确保sampler具有非空的HeaderManager实例
if (sampler.getHeaderManager() == null) {
sampler.setHeaderManager(new HeaderManager());
}
sampler.getHeaderManager().clear();
Header signatureHeader = new Header("signature", signature);
sampler.getHeaderManager().add(signatureHeader);
Header timeStampHeader = new Header("signTimestamp", String.valueOf(timestamp));
sampler.getHeaderManager().add(timeStampHeader);
Header contentHeader = new Header("Content-Type", "application/json");
sampler.getHeaderManager().add(contentHeader);
Header apiKeyHeader = new Header("key", apiKey);
sampler.getHeaderManager().add(apiKeyHeader);
log.info("==========添加后的请求头" + sampler.getHeaderManager().toString());
} catch (Exception e) {
log.error("Error generating signature: " + e.getMessage());
}
} else {
log.warn("Sampler is not an instance of HTTPSamplerProxy.");
}
image.png