package com.atlassian.asap.core.keys;
   
   
   import com.atlassian.asap.api.exception.CannotRetrieveKeyException;
   import com.atlassian.asap.core.SecurityProvider;
   import org.apache.commons.io.IOUtils;
   import org.bouncycastle.asn1.pkcs.PrivateKeyInfo;
   import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
   import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo;
  import org.bouncycastle.openssl.PEMParser;
  import org.bouncycastle.openssl.jcajce.JcaPEMKeyConverter;
  
  import java.io.IOException;
  import java.io.Reader;
  import java.io.StringReader;
  import java.nio.charset.StandardCharsets;
  import java.security.KeyFactory;
  import java.security.NoSuchAlgorithmException;
  import java.security.PrivateKey;
  import java.security.Provider;
  import java.security.PublicKey;
  import java.security.spec.InvalidKeySpecException;
  import java.security.spec.PKCS8EncodedKeySpec;
  
  /**
   * Reads keys encoded in a data uri. Only Base64 DER encoded key data is supported.
   * The format of the data uri must be:
   * <pre>
   *  data:{key-mediatype};kid={key-identifier};base64,{encoded-key-data}
   * </pre>
   * In addition to the media type of the key, the key identifier must be specified using
   * url escaping as the value of the 'kid' parameter.
   *
   * <ul>
   *    <li> For private keys, only {@code application/pkcs8} media type is supported.
   *         PKCS#8 object in the data section MUST be DER-ASN.1-encoded {@code PrivateKeyInfo}.</li>
   *    <li> For public keys, only {@code application/x-pem-file} media type is supported.
   *         The public key in the PEM must be DER-ASN.1-encoded {@code SubjectPublicKeyInfo}.</li>
   * </ul>
   *
   * @see <a href="https://tools.ietf.org/html/rfc2397">RFC 2397 The "data" uri scheme</a>
   * @see DataUriUtil
   */
  public class DataUriKeyReader implements KeyReader {
      public static final String DATA_URI_PKCS8_HEADER = "data:application/pkcs8;";
      public static final String DATA_URI_PEM_HEADER = "data:application/x-pem-file;";
  
      private final Provider provider;
  
      public DataUriKeyReader() {
          this(SecurityProvider.getProvider());
      }
  
      public DataUriKeyReader(Provider provider) {
          this.provider = provider;
      }
  
      @Override
      public PrivateKey readPrivateKey(Reader reader) throws CannotRetrieveKeyException {
          try {
              String dataUri = IOUtils.toString(reader);
              if (dataUri.startsWith(DATA_URI_PKCS8_HEADER)) {
                  byte[] keyData = DataUriUtil.getKeyData(dataUri);
                  AlgorithmIdentifier algorithmIdentifier = PrivateKeyInfo.getInstance(keyData).getPrivateKeyAlgorithm();
                  KeyFactory keyFactory = KeyFactory.getInstance(algorithmIdentifier.getAlgorithm().getId(), provider);
                  return keyFactory.generatePrivate(new PKCS8EncodedKeySpec(keyData));
              } else {
                  throw new CannotRetrieveKeyException("Data uri could not be parsed due to unexpected prefix");
              }
          } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
              throw new CannotRetrieveKeyException("Error reading private key, unknown key type", e);
          } catch (IOException | IllegalArgumentException e) {
              throw new CannotRetrieveKeyException("Error reading private key", e);
          }
      }
  
      @Override
      public PublicKey readPublicKey(Reader reader) throws CannotRetrieveKeyException {
          try {
              String dataUri = IOUtils.toString(reader);
              if (dataUri.startsWith(DATA_URI_PEM_HEADER)) {
                  byte[] keyData = DataUriUtil.getKeyData(dataUri);
                  PEMParser pemParser = new PEMParser(new StringReader(new String(keyData, StandardCharsets.US_ASCII)));
                  SubjectPublicKeyInfo pub = SubjectPublicKeyInfo.getInstance(pemParser.readObject());
                  JcaPEMKeyConverter converter = new JcaPEMKeyConverter().setProvider(provider);
                  return converter.getPublicKey(pub);
              } else {
                  throw new CannotRetrieveKeyException("Data uri could not be parsed due to unexpected prefix");
              }
          } catch (IOException | IllegalArgumentException e) {
              throw new CannotRetrieveKeyException("Error reading Public key", e);
          }
      }
  }