package com.sun.identity.saml.xmlsig;

import com.sun.identity.common.SystemConfigurationUtil;
import com.sun.identity.saml.common.SAMLException;
import com.sun.identity.saml.common.SAMLUtilsCommon;
import com.sun.identity.shared.encode.Base64;
import jakarta.annotation.Nullable;
import java.io.BufferedReader;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.security.KeyStore;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.cert.Certificate;
import java.security.cert.X509Certificate;
import java.util.Enumeration;
import java.util.HashMap;

public class JKSKeyProvider implements KeyProvider {
   private static final String DEFAULT_KEYSTORE_FILE_PROP = "com.sun.identity.saml.xmlsig.keystore";
   private static final String DEFAULT_KEYSTORE_PASS_FILE_PROP = "com.sun.identity.saml.xmlsig.storepass";
   private static final String DEFAULT_KEYSTORE_TYPE_PROP = "com.sun.identity.saml.xmlsig.storetype";
   private static final String DEFAULT_PRIVATE_KEY_PASS_FILE_PROP = "com.sun.identity.saml.xmlsig.keypass";
   private final HashMap<String, Certificate> keyTable = new HashMap<>();
   private KeyStore ks;
   private String privateKeyPass;
   private String keystorePass = "";
   private String keystoreFile = "";
   private String keystoreType = "JKS";

   private void initialize(String keyStoreFilePropName, String keyStorePassFilePropName, String keyStoreTypePropName, String privateKeyPassFilePropName) {
      this.keystoreFile = SystemConfigurationUtil.getProperty(keyStoreFilePropName);
      if (this.keystoreFile == null || this.keystoreFile.isEmpty()) {
         SAMLUtilsCommon.DEBUG.error("JKSKeyProvider: keystore file does not exist");
      }

      String kspfile = SystemConfigurationUtil.getProperty(keyStorePassFilePropName);
      String tmpKsType = SystemConfigurationUtil.getProperty(keyStoreTypePropName);
      if (null != tmpKsType) {
         this.keystoreType = tmpKsType.trim();
      }

      if (kspfile != null) {
         try (
            FileInputStream fis = new FileInputStream(kspfile);
            InputStreamReader isr = new InputStreamReader(fis, StandardCharsets.UTF_8);
            BufferedReader br = new BufferedReader(isr);
         ) {
            this.keystorePass = SAMLUtilsCommon.decodePassword(br.readLine());
         } catch (Exception ex) {
            SAMLUtilsCommon.DEBUG.error("JKSKeyProvider.initialize: Unable to read keystore password file " + kspfile + ": " + ex.getMessage());
         }
      } else {
         SAMLUtilsCommon.DEBUG.error("JKSKeyProvider: keystore password is null");
      }

      String pkpfile = SystemConfigurationUtil.getProperty(privateKeyPassFilePropName);
      if (pkpfile != null) {
         try (
            FileInputStream fis = new FileInputStream(pkpfile);
            InputStreamReader isr = new InputStreamReader(fis, StandardCharsets.UTF_8);
            BufferedReader br = new BufferedReader(isr);
         ) {
            this.privateKeyPass = SAMLUtilsCommon.decodePassword(br.readLine());
         } catch (Exception ex) {
            SAMLUtilsCommon.DEBUG.error("JKSKeyProvider.initialize: Unable to read privatekey password file " + kspfile + ": " + ex.getMessage());
         }
      }
   }

   private void mapPk2Cert() {
      try (FileInputStream fis = new FileInputStream(this.keystoreFile)) {
         this.ks = KeyStore.getInstance(this.keystoreType);
         this.ks.load(fis, this.keystorePass.toCharArray());
         Enumeration<String> e = this.ks.aliases();

         while (e.hasMoreElements()) {
            String alias = e.nextElement();
            Certificate cert = this.getCertificate(alias);
            PublicKey pk = this.getPublicKey(alias);
            String key = Base64.encode(pk.getEncoded());
            this.keyTable.put(key, cert);
         }

         SAMLUtilsCommon.DEBUG.message("KeyTable size = " + this.keyTable.size());
      } catch (Exception e) {
         SAMLUtilsCommon.DEBUG.error("mapPk2Cert.JKSKeyProvider:", e);
      }
   }

   public JKSKeyProvider() {
      this(
         "com.sun.identity.saml.xmlsig.keystore",
         "com.sun.identity.saml.xmlsig.storepass",
         "com.sun.identity.saml.xmlsig.storetype",
         "com.sun.identity.saml.xmlsig.keypass"
      );
   }

   public JKSKeyProvider(String keyStoreFilePropName, String keyStorePassFilePropName, String keyStoreTypePropName, String privateKeyPassFilePropName) {
      this.initialize(keyStoreFilePropName, keyStorePassFilePropName, keyStoreTypePropName, privateKeyPassFilePropName);
      this.mapPk2Cert();
   }

   @Override
   public X509Certificate getX509Certificate(String certAlias) {
      if (certAlias != null && !certAlias.isEmpty()) {
         X509Certificate cert = null;

         try {
            cert = (X509Certificate)this.ks.getCertificate(certAlias);
            if (cert == null) {
               SAMLUtilsCommon.DEBUG.warning("Unable to find certificate alias \"" + certAlias + "\" in keystore " + this.keystoreFile);
            }
         } catch (Exception e) {
            SAMLUtilsCommon.DEBUG.error("Unable to get cert alias:" + certAlias, e);
         }

         return cert;
      } else {
         return null;
      }
   }

   @Override
   public PublicKey getPublicKey(String keyAlias) {
      if (keyAlias != null && !keyAlias.isEmpty()) {
         try {
            X509Certificate cert = (X509Certificate)this.ks.getCertificate(keyAlias);
            if (cert != null) {
               return cert.getPublicKey();
            }

            SAMLUtilsCommon.DEBUG.warning("Unable to find certificate alias \"" + keyAlias + "\" in keystore " + this.keystoreFile);
         } catch (Exception e) {
            SAMLUtilsCommon.DEBUG.error("Unable to get public key:" + keyAlias, e);
         }

         return null;
      } else {
         return null;
      }
   }

   @Override
   public PrivateKey getPrivateKey(String certAlias) {
      PrivateKey key = null;

      try {
         key = (PrivateKey)this.ks.getKey(certAlias, this.privateKeyPass.toCharArray());
         if (key == null) {
            SAMLUtilsCommon.DEBUG.warning("Unable to find key alias \"" + certAlias + "\" in keystore " + this.keystoreFile);
         }
      } catch (Exception e) {
         SAMLUtilsCommon.DEBUG.error(e.getMessage());
      }

      return key;
   }

   public String getPrivateKeyPass() {
      return this.privateKeyPass;
   }

   @Override
   public KeyStore getKeyStore() {
      return this.ks;
   }

   @Nullable
   public Certificate getCertificate(String certAlias) {
      try {
         Certificate cert = this.ks.getCertificate(certAlias);
         if (cert == null) {
            SAMLUtilsCommon.DEBUG.warning("Unable to find certificate alias \"" + certAlias + "\" in keystore " + this.keystoreFile);
         }

         return cert;
      } catch (Exception e) {
         SAMLUtilsCommon.DEBUG.error(e.getMessage());
         return null;
      }
   }

   public void store() throws SAMLException {
      try (OutputStream keyStoreOStream = new FileOutputStream(this.keystoreFile)) {
         this.ks.store(keyStoreOStream, this.keystorePass.toCharArray());
         if (SAMLUtilsCommon.DEBUG.messageEnabled()) {
            SAMLUtilsCommon.DEBUG.message("Keystore saved in " + this.keystoreFile);
         }
      } catch (Exception e) {
         SAMLUtilsCommon.DEBUG.error(e.getMessage());
         throw new SAMLException(e.getMessage());
      }
   }

   @Override
   public Certificate getCertificate(PublicKey publicKey) {
      String key = Base64.encode(publicKey.getEncoded());
      return this.keyTable.get(key);
   }
}