package de.ivu.elect.encrypt;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.math.BigInteger;
import java.security.GeneralSecurityException;
import java.security.KeyPair;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.UnrecoverableKeyException;
import java.security.cert.Certificate;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.util.Date;

import javax.annotation.Nullable;

import sun.security.x509.AlgorithmId;
import sun.security.x509.CertificateAlgorithmId;
import sun.security.x509.CertificateSerialNumber;
import sun.security.x509.CertificateValidity;
import sun.security.x509.CertificateVersion;
import sun.security.x509.CertificateX509Key;
import sun.security.x509.X500Name;
import sun.security.x509.X509CertImpl;
import sun.security.x509.X509CertInfo;

public class KeystoreHandler {

  private final KeystoreHandlerConfig keystoreHandlerConfig;

  public KeystoreHandler(KeystoreHandlerConfig keystoreHandlerConfig) {
    this.keystoreHandlerConfig = keystoreHandlerConfig;
  }

  public KeyStore loadOrCreateKeyStore(@Nullable InputStream is, char[] password) throws KeyStoreException, CertificateException, NoSuchAlgorithmException, IOException {
    KeyStore keyStore = KeyStore.getInstance(keystoreHandlerConfig.getKeystoreType());
    keyStore.load(is, password);
    return keyStore;
  }

  public void saveKeyStore(KeyStore keyStore, char[] password, OutputStream os) throws CertificateException, NoSuchAlgorithmException, KeyStoreException,
      IOException {
    keyStore.store(os, password);
  }

  public KeyPair readKeyPairFromKeyStore(KeyStore keyStore, char[] password, String alias) throws KeyStoreException, UnrecoverableKeyException, NoSuchAlgorithmException {
    PrivateKey privateKey = (PrivateKey) keyStore.getKey(alias, password);
    Certificate cert = keyStore.getCertificate(alias);
    PublicKey publicKey = cert.getPublicKey();
    return new KeyPair(publicKey, privateKey);
  }

  public void saveKeyPairToKeyStore(KeyPair keyPair, KeyStore keyStore, char[] password, String alias) throws GeneralSecurityException, IOException {
    Certificate[] chain = {generateCertificate(keyPair)};
    keyStore.setKeyEntry(alias, keyPair.getPrivate(), password, chain);
  }

  private X509Certificate generateCertificate(KeyPair keyPair) throws GeneralSecurityException, IOException {
    PrivateKey privateKey = keyPair.getPrivate();

    X509CertInfo info = new X509CertInfo();

    Date from = new Date();
    Date to = new Date(from.getTime() + keystoreHandlerConfig.getCertificateValidityInDays() * 1000L * 24L * 60L * 60L);

    CertificateValidity interval = new CertificateValidity(from, to);
    BigInteger serialNumber = new BigInteger(64, SecureRandom.getInstanceStrong());
    X500Name owner = new X500Name(keystoreHandlerConfig.getCertificateOwner());
    AlgorithmId sigAlgId = new AlgorithmId(AlgorithmId.md5WithRSAEncryption_oid);

    info.set(X509CertInfo.VALIDITY, interval);
    info.set(X509CertInfo.SERIAL_NUMBER, new CertificateSerialNumber(serialNumber));
    info.set(X509CertInfo.SUBJECT, owner);
    info.set(X509CertInfo.ISSUER, owner);
    info.set(X509CertInfo.KEY, new CertificateX509Key(keyPair.getPublic()));
    info.set(X509CertInfo.VERSION, new CertificateVersion(CertificateVersion.V3));
    info.set(X509CertInfo.ALGORITHM_ID, new CertificateAlgorithmId(sigAlgId));

    // Sign the cert to identify the algorithm that's used.
    X509CertImpl certificate = new X509CertImpl(info);
    certificate.sign(privateKey, keystoreHandlerConfig.getCertificateSigAlgName());

    // Update the algorith, and resign.
    sigAlgId = (AlgorithmId) certificate.get(X509CertImpl.SIG_ALG);
    info.set(CertificateAlgorithmId.NAME + "." + CertificateAlgorithmId.ALGORITHM, sigAlgId);
    certificate = new X509CertImpl(info);
    certificate.sign(privateKey, keystoreHandlerConfig.getCertificateSigAlgName());

    return certificate;
  }

}
