package org.jboss.resteasy.security.doseta;

import org.jboss.resteasy.logging.Logger;
import org.jboss.resteasy.spi.ResteasyProviderFactory;
import org.jboss.resteasy.util.Base64;
import org.jboss.resteasy.util.ParameterParser;

import javax.naming.directory.Attributes;
import javax.naming.directory.DirContext;
import javax.naming.directory.InitialDirContext;
import javax.ws.rs.core.SecurityContext;
import java.io.IOException;
import java.io.InputStream;
import java.security.KeyFactory;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.spec.X509EncodedKeySpec;
import java.util.Hashtable;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author <a href="mailto:bill@burkecentral.com">Bill Burke</a>
 * @version $Revision: 1 $
 */
public class DosetaKeyRepository implements KeyRepository
{
   private static final Logger log = Logger.getLogger(DosetaKeyRepository.class);

   protected class CacheEntry<T>
   {
      public long time = System.currentTimeMillis();
      public T key;

      protected CacheEntry(T key)
      {
         this.key = key;
      }

      public boolean isStale()
      {
         return time + cacheTimeout >= System.currentTimeMillis();
      }
   }

   protected ConcurrentHashMap<String, CacheEntry<PrivateKey>> privateCache = new ConcurrentHashMap<String, CacheEntry<PrivateKey>>();
   protected ConcurrentHashMap<String, CacheEntry<PublicKey>> publicCache = new ConcurrentHashMap<String, CacheEntry<PublicKey>>();
   protected KeyStoreKeyRepository keyStore;
   protected String defaultPrivateDomain;
   protected boolean useDns = false;
   protected boolean userPrincipalAsPrivateSelector = false;
   protected String dnsUri;
   protected long cacheTimeout = 3600000l; // 1 hour
   protected String keyStorePath;
   protected String keyStoreFile;
   protected String keyStorePassword;

   public void start()
   {
      if (keyStore == null)
      {
         if (keyStoreFile != null)
         {
            try
            {
               keyStore = new KeyStoreKeyRepository(keyStoreFile, keyStorePassword);
            }
            catch (IOException e)
            {
               throw new RuntimeException(e);
            }
         }
         else
         {
            if (keyStorePath != null)
            {
               InputStream is = Thread.currentThread().getContextClassLoader().getResourceAsStream(keyStorePath.trim());
               if (is == null) throw new RuntimeException("Unable to find key store in path: " + keyStorePath);
               keyStore = new KeyStoreKeyRepository(is, keyStorePassword);
            }
         }
      }
   }

   public String getDefaultPrivateSelector()
   {
      if (userPrincipalAsPrivateSelector)
      {
         SecurityContext securityContext = ResteasyProviderFactory.getContextData(SecurityContext.class);
         if (securityContext != null)
         {
            return securityContext.getUserPrincipal().getName();
         }
      }
      return null;
   }

   public String getKeyStorePath()
   {
      return keyStorePath;
   }

   public void setKeyStorePath(String keyStorePath)
   {
      this.keyStorePath = keyStorePath;
   }

   public String getKeyStoreFile()
   {
      return keyStoreFile;
   }

   public void setKeyStoreFile(String keyStoreFile)
   {
      this.keyStoreFile = keyStoreFile;
   }

   public String getKeyStorePassword()
   {
      return keyStorePassword;
   }

   public void setKeyStorePassword(String keyStorePassword)
   {
      this.keyStorePassword = keyStorePassword;
   }

   public KeyStoreKeyRepository getKeyStore()
   {
      return keyStore;
   }

   public void setKeyStore(KeyStoreKeyRepository keyStore)
   {
      this.keyStore = keyStore;
   }

   public String getDefaultPrivateDomain()
   {
      return defaultPrivateDomain;
   }

   public void setDefaultPrivateDomain(String defaultPrivateDomain)
   {
      this.defaultPrivateDomain = defaultPrivateDomain;
   }

   public boolean isUseDns()
   {
      return useDns;
   }

   public void setUseDns(boolean useDns)
   {
      this.useDns = useDns;
   }

   public boolean isUserPrincipalAsPrivateSelector()
   {
      return userPrincipalAsPrivateSelector;
   }

   public void setUserPrincipalAsPrivateSelector(boolean userPrincipalAsPrivateSelector)
   {
      this.userPrincipalAsPrivateSelector = userPrincipalAsPrivateSelector;
   }

   public String getDnsUri()
   {
      return dnsUri;
   }

   public void setDnsUri(String dnsUri)
   {
      this.dnsUri = dnsUri;
   }

   public long getCacheTimeout()
   {
      return cacheTimeout;
   }

   public void setCacheTimeout(long cacheTimeout)
   {
      this.cacheTimeout = cacheTimeout;
   }

   protected void addPrivate(String alias, PrivateKey key)
   {
      privateCache.put(alias, new CacheEntry<PrivateKey>(key));
   }

   protected void addPublic(String alias, PublicKey key)
   {
      publicCache.put(alias, new CacheEntry<PublicKey>(key));
   }

   protected PrivateKey getPrivateCache(String alias)
   {
      CacheEntry<PrivateKey> entry = privateCache.get(alias);
      if (entry == null) return null;
      if (entry.isStale())
      {
         privateCache.remove(entry, entry);
         return null;
      }
      return entry.key;
   }

   protected PublicKey getPublicCache(String alias)
   {
      CacheEntry<PublicKey> entry = publicCache.get(alias);
      if (entry == null) return null;
      if (entry.isStale())
      {
         publicCache.remove(entry, entry);
         return null;
      }
      return entry.key;
   }

   public String getAlias(DKIMSignature header)
   {
      StringBuffer buf = new StringBuffer();
      String selector = header.getSelector();
      if (selector != null) buf.append(selector.trim()).append(".");
      buf.append("_domainKey.");
      String domain = header.getDomain();
      if (domain == null)
      {
         throw new RuntimeException("domain attribute is required in header to find a key");
      }
      buf.append(domain);
      return buf.toString();
   }

   public PrivateKey findPrivateKey(DKIMSignature header)
   {
      String alias = getAlias(header);
      if (alias == null) return null;
      PrivateKey key = getPrivateCache(alias);
      if (key != null) return key;

      if (keyStore != null)
      {
         key = keyStore.getPrivateKey(alias);
         if (key != null) addPrivate(alias, key);
      }

      return key;
   }

   public PublicKey findPublicKey(DKIMSignature header)
   {
      String alias = getAlias(header); // never use principal to find a public key
      if (alias == null) return null;
      PublicKey key = getPublicCache(alias);
      if (key != null) return key;

      if (keyStore != null)
      {
         key = keyStore.getPublicKey(alias);
         if (key != null) addPublic(alias, key);
      }

      if (useDns)
      {
         key = findFromDns(alias);
         addPublic(alias, key);
      }


      return key;
   }

   protected PublicKey findFromDns(String alias)
   {
      if (log.isDebugEnabled()) log.debug(">>>> Check DNS: " + alias);
      PublicKey key;
      try
      {
         Hashtable<String, String> env = new Hashtable<String, String>();
         env.put("java.naming.factory.initial", "com.sun.jndi.dns.DnsContextFactory");
         if (dnsUri != null)
         {
            env.put("java.naming.provider.url", dnsUri);
         }
         DirContext dnsContext = new InitialDirContext(env);

         Attributes attrs1 = dnsContext.getAttributes(alias, new String[]{"TXT"});
         javax.naming.directory.Attribute txtrecord = attrs1.get("txt");
         String record = txtrecord.get().toString();
         if (log.isDebugEnabled()) log.debug(">>>> DNS found record: " + record);
         ParameterParser parser = new ParameterParser();
         parser.setLowerCaseNames(true);
         Map<String, String> keyEntry = parser.parse(record, ';');
         String type = keyEntry.get("k");
         if (type != null && !type.toLowerCase().equals("rsa"))
            throw new RuntimeException("Unsupported key type: " + type);
         String pem = keyEntry.get("p");
         if (pem == null)
         {
            throw new RuntimeException("No p entry in text record.");
         }
         if (log.isDebugEnabled()) log.debug("pem: " + pem);
         byte[] der = Base64.decode(pem);


         X509EncodedKeySpec spec =
                 new X509EncodedKeySpec(der);
         KeyFactory kf = KeyFactory.getInstance("RSA");
         key = kf.generatePublic(spec);

      }
      catch (Exception e)
      {
         throw new RuntimeException("Failed to find public key in DNS " + alias, e);
      }
      return key;
   }


}
