jmeter_aws_signature.groovy

import java.text.SimpleDateFormat
import java.security.MessageDigest
import javax.crypto.Mac
import javax.crypto.spec.SecretKeySpec
import org.apache.jmeter.protocol.http.control.Header

def region = vars.get("aws_region")
def serviceName = vars.get("aws_service_name")
def accessKeyId = vars.get("aws_access_key_id")
def secretAccessKey = vars.get("aws_secret_access_key")
def sessionToken = vars.get("aws_session_token")

def now = new Date()
def amzFormat = new SimpleDateFormat("yyyyMMdd'T'HHmmss'Z'")
def stampFormat = new SimpleDateFormat("yyyyMMdd")
amzFormat.setTimeZone(TimeZone.getTimeZone("UTC"))
stampFormat.setTimeZone(TimeZone.getTimeZone("UTC"))
def amzDate = amzFormat.format(now)
def dateStamp = stampFormat.format(now)

def signedHeaders = "host;x-amz-date"
def canonicalHeaders = "" +
	"host:${sampler.getDomain()}\n" +
	"x-amz-date:${amzDate}\n"
if (sessionToken) {
	signedHeaders += ";x-amz-security-token"
	canonicalHeaders += "x-amz-security-token:${sessionToken}\n"
}

def method = sampler.getMethod()
def payload = ""
if (method in ["POST", "PUT", "PATCH"]) {
	payload = sampler.getArguments()
	.collect { it.getStringValue().substring(1) }
	.join("")
}

def canonicalRequest = "" +
	"${method}\n" +
	"${sampler.getPath()}\n" +
	"${canonicalQuery(sampler.getQueryString())}\n" +
	"${canonicalHeaders}\n" +
	"${signedHeaders}\n" +
	hex(sha256(payload))

def algorithm = "AWS4-HMAC-SHA256"
def aws4request = "aws4_request"
def credentialScope = "${dateStamp}/${region}/${serviceName}/${aws4request}"
def stringToSign = "" +
	"${algorithm}\n" +
	"${amzDate}\n" +
	"${credentialScope}\n" +
	hex(sha256(canonicalRequest))

def signingKey = calculateKey(secretAccessKey, dateStamp, region, serviceName, aws4request)
def signature = hex(hmacSha256(signingKey, stringToSign))

def authorizationHeader = "" +
	algorithm + " " +
	"Credential=${accessKeyId}/${dateStamp}/${region}/${serviceName}/${aws4request}, " +
	"SignedHeaders=${signedHeaders}, " +
	"Signature=${signature}"

def manager = sampler.getHeaderManager()
manager.removeHeaderNamed("Authorization")
manager.add(new Header("Authorization", authorizationHeader))
manager.removeHeaderNamed("X-Amz-Date")
manager.add(new Header("X-Amz-Date", amzDate))
if (sessionToken) {
	manager.removeHeaderNamed("X-Amz-Security-Token")
	sampler.getHeaderManager().add(new Header("X-Amz-Security-Token", sessionToken))
}

def urlEncode(value) {
	URLEncoder.encode(value, "UTF-8")
		.replace("=", "%20")
		.replace("*", "%2A")
		.replace("%7E", "~")
}

def canonicalQuery(query) {
	if (!query)
		return ""
	query.split("&")
		.collect { it.split("=", 2) }
		.collect {
			def key = URLDecoder.decode(it[0], "UTF-8")
			def value = URLDecoder.decode(it[1], "UTF-8")
			[urlEncode(key), urlEncode(value)]
		}
		.sort { a, b -> a[0] <=> b[0] }
		.collect { "${it[0]}=${it[1]}" }
		.join("&")
}

def sha256(value) {
	def md = MessageDigest.getInstance("SHA-256")
	md.digest(value.getBytes("UTF-8"))
}

def hmacSha256(key, value) {
	def mac = Mac.getInstance("HmacSHA256")
	def secretKeySpec = new SecretKeySpec(key, "HmacSHA256")
	mac.init(secretKeySpec)
	mac.doFinal(value.getBytes())
}

def hex(value) {
	value.encodeHex()
}

def calculateKey(key, date, region, service, word) {
	hmacSha256(hmacSha256(hmacSha256(hmacSha256(("AWS4${key}").getBytes(), date), region), service), word)
}