package it.acxent.api.amz;

import java.math.BigInteger;
import java.net.URLEncoder;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.text.DateFormat;
import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.HashMap;
import java.util.Map;
import java.util.TimeZone;
import java.util.TreeMap;
import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;

public class AWSV4Auth {
  private String accessKeyID;
  
  private String secretAccessKey;
  
  private String regionName;
  
  private String serviceName;
  
  private String httpMethodName;
  
  private String canonicalURI;
  
  private TreeMap<String, String> queryParametes;
  
  private TreeMap<String, String> awsHeaders;
  
  private String payload;
  
  private AWSV4Auth() {}
  
  public static void main(String[] args) {
    String url = "xxxxx-yyyyy-r6nvlhpscgdwms5.ap-northeast-1.es.amazonaws.com/inventory/simple/123";
    TreeMap<String, String> awsHeaders = new TreeMap<>();
    awsHeaders.put("host", "xxxxx-yyyyy-r6nvlhpscgdwms5.ap-northeast-1.es.amazonaws.com");
    AWSV4Auth aWSV4Auth = new Builder("exampleKey", "exampleSecret").regionName("xx-yy-zzz").serviceName("es")

      
      .httpMethodName("GET")
      .canonicalURI("/inventory/simple/123")
      .queryParametes(null)
      .awsHeaders(awsHeaders)
      .payload(null)
      .debug()
      .build();
    Map<String, String> header = aWSV4Auth.getHeaders();
    for (Map.Entry<String, String> entrySet : header.entrySet()) {
      String key = entrySet.getKey();
      String value = entrySet.getValue();
      System.out.println(key + ":" + key);
    } 
  }
  
  public static class Builder {
    private String accessKeyID;
    
    private String secretAccessKey;
    
    private String regionName;
    
    private String serviceName;
    
    private String httpMethodName;
    
    private String canonicalURI;
    
    private TreeMap<String, String> queryParametes;
    
    private TreeMap<String, String> awsHeaders;
    
    private String payload;
    
    private boolean debug = false;
    
    public Builder(String accessKeyID, String secretAccessKey) {
      this.accessKeyID = accessKeyID;
      this.secretAccessKey = secretAccessKey;
    }
    
    public Builder regionName(String regionName) {
      this.regionName = regionName;
      return this;
    }
    
    public Builder serviceName(String serviceName) {
      this.serviceName = serviceName;
      return this;
    }
    
    public Builder httpMethodName(String httpMethodName) {
      this.httpMethodName = httpMethodName;
      return this;
    }
    
    public Builder canonicalURI(String canonicalURI) {
      this.canonicalURI = canonicalURI;
      return this;
    }
    
    public Builder queryParametes(TreeMap<String, String> queryParametes) {
      this.queryParametes = queryParametes;
      return this;
    }
    
    public Builder awsHeaders(TreeMap<String, String> awsHeaders) {
      this.awsHeaders = awsHeaders;
      return this;
    }
    
    public Builder payload(String payload) {
      this.payload = payload;
      return this;
    }
    
    public Builder debug() {
      this.debug = true;
      return this;
    }
    
    public AWSV4Auth build() {
      return new AWSV4Auth(this);
    }
  }
  
  private boolean debug = false;
  
  private final String HMACAlgorithm = "AWS4-HMAC-SHA256";
  
  private final String aws4Request = "aws4_request";
  
  private String strSignedHeader;
  
  private String xAmzDate;
  
  private String currentDate;
  
  private AWSV4Auth(Builder builder) {
    this.accessKeyID = builder.accessKeyID;
    this.secretAccessKey = builder.secretAccessKey;
    this.regionName = builder.regionName;
    this.serviceName = builder.serviceName;
    this.httpMethodName = builder.httpMethodName;
    this.canonicalURI = builder.canonicalURI;
    this.queryParametes = builder.queryParametes;
    this.awsHeaders = builder.awsHeaders;
    this.payload = builder.payload;
    this.debug = builder.debug;
    this.xAmzDate = getTimeStamp();
    this.currentDate = getDate();
  }
  
  private String prepareCanonicalRequest() {
    StringBuilder canonicalURL = new StringBuilder("");
    canonicalURL.append(this.httpMethodName).append("\n");
    this.canonicalURI = (this.canonicalURI == null || this.canonicalURI.trim().isEmpty()) ? "/" : this.canonicalURI;
    canonicalURL.append(this.canonicalURI).append("\n");
    StringBuilder queryString = new StringBuilder("");
    if (this.queryParametes != null && !this.queryParametes.isEmpty()) {
      for (Map.Entry<String, String> entrySet : this.queryParametes.entrySet()) {
        String key = entrySet.getKey();
        String value = entrySet.getValue();
        queryString.append(key).append("=").append(encodeParameter(value)).append("&");
      } 
      queryString.deleteCharAt(queryString.lastIndexOf("&"));
      queryString.append("\n");
    } else {
      queryString.append("\n");
    } 
    canonicalURL.append((CharSequence)queryString);
    StringBuilder signedHeaders = new StringBuilder("");
    if (this.awsHeaders != null && !this.awsHeaders.isEmpty()) {
      for (Map.Entry<String, String> entrySet : this.awsHeaders.entrySet()) {
        String key = entrySet.getKey();
        String value = entrySet.getValue();
        signedHeaders.append(key).append(";");
        canonicalURL.append(key).append(":").append(value).append("\n");
      } 
      canonicalURL.append("\n");
    } else {
      canonicalURL.append("\n");
    } 
    this.strSignedHeader = signedHeaders.substring(0, signedHeaders.length() - 1);
    canonicalURL.append(this.strSignedHeader).append("\n");
    this.payload = (this.payload == null) ? "" : this.payload;
    canonicalURL.append(generateHex(this.payload));
    if (this.debug)
      System.out.println("##Canonical Request:\n" + canonicalURL.toString()); 
    return canonicalURL.toString();
  }
  
  private String prepareStringToSign(String canonicalURL) {
    String stringToSign = "";
    stringToSign = "AWS4-HMAC-SHA256\n";
    stringToSign = stringToSign + stringToSign + "\n";
    stringToSign = stringToSign + stringToSign + "/" + this.currentDate + "/" + this.regionName + "/aws4_request\n";
    stringToSign = stringToSign + stringToSign;
    if (this.debug)
      System.out.println("##String to sign:\n" + stringToSign); 
    return stringToSign;
  }
  
  private String calculateSignature(String stringToSign) {
    try {
      byte[] signatureKey = getSignatureKey(this.secretAccessKey, this.currentDate, this.regionName, this.serviceName);
      byte[] signature = HmacSHA256(signatureKey, stringToSign);
      String strHexSignature = bytesToHex(signature);
      return strHexSignature;
    } catch (Exception ex) {
      ex.printStackTrace();
      return null;
    } 
  }
  
  public Map<String, String> getHeaders() {
    this.awsHeaders.put("x-amz-date", this.xAmzDate);
    String canonicalURL = prepareCanonicalRequest();
    String stringToSign = prepareStringToSign(canonicalURL);
    String signature = calculateSignature(stringToSign);
    if (signature != null) {
      Map<String, String> header = new HashMap<>(0);
      header.put("Authorization", buildAuthorizationString(signature));
      if (this.debug) {
        System.out.println("##Signature:\n" + signature);
        System.out.println("##Header:");
        for (Map.Entry<String, String> entrySet : header.entrySet())
          System.out.println((String)entrySet.getKey() + " = " + (String)entrySet.getKey()); 
        System.out.println("================================");
      } 
      return header;
    } 
    if (this.debug)
      System.out.println("##Signature:\n" + signature); 
    return null;
  }
  
  private String buildAuthorizationString(String strSignature) {
    return "AWS4-HMAC-SHA256 Credential=" + this.accessKeyID + "/" + getDate() + "/" + this.regionName + "/" + this.serviceName + "/aws4_request,SignedHeaders=" + this.strSignedHeader + ",Signature=" + strSignature;
  }
  
  public static String generateHex(String data) {
    try {
      MessageDigest messageDigest = MessageDigest.getInstance("SHA-256");
      messageDigest.update(data.getBytes("UTF-8"));
      byte[] digest = messageDigest.digest();
      return String.format("%064x", new BigInteger(1, digest));
    } catch (NoSuchAlgorithmException|java.io.UnsupportedEncodingException e) {
      e.printStackTrace();
      return null;
    } 
  }
  
  private byte[] HmacSHA256(byte[] key, String data) throws Exception {
    String algorithm = "HmacSHA256";
    Mac mac = Mac.getInstance(algorithm);
    mac.init(new SecretKeySpec(key, algorithm));
    return mac.doFinal(data.getBytes("UTF8"));
  }
  
  private byte[] getSignatureKey(String key, String date, String regionName, String serviceName) throws Exception {
    byte[] kSecret = ("AWS4" + key).getBytes("UTF8");
    byte[] kDate = HmacSHA256(kSecret, date);
    byte[] kRegion = HmacSHA256(kDate, regionName);
    byte[] kService = HmacSHA256(kRegion, serviceName);
    byte[] kSigning = HmacSHA256(kService, "aws4_request");
    return kSigning;
  }
  
  protected static final char[] hexArray = "0123456789ABCDEF".toCharArray();
  
  private String bytesToHex(byte[] bytes) {
    char[] hexChars = new char[bytes.length * 2];
    for (int j = 0; j < bytes.length; j++) {
      int v = bytes[j] & 0xFF;
      hexChars[j * 2] = hexArray[v >>> 4];
      hexChars[j * 2 + 1] = hexArray[v & 0xF];
    } 
    return new String(hexChars).toLowerCase();
  }
  
  private String getTimeStamp() {
    DateFormat dateFormat = new SimpleDateFormat("yyyyMMdd'T'HHmmss'Z'");
    dateFormat.setTimeZone(TimeZone.getTimeZone("UTC"));
    return dateFormat.format(new Date());
  }
  
  private String getDate() {
    DateFormat dateFormat = new SimpleDateFormat("yyyyMMdd");
    dateFormat.setTimeZone(TimeZone.getTimeZone("UTC"));
    return dateFormat.format(new Date());
  }
  
  private String encodeParameter(String param) {
    try {
      return URLEncoder.encode(param, "UTF-8");
    } catch (Exception e) {
      return URLEncoder.encode(param);
    } 
  }
}