import java.security._
import java.net.URL
import javax.net.ssl.{KeyManagerFactory, TrustManagerFactory}
import org.apache.http.conn.ssl.SSLSocketFactory
import org.apache.http.conn.scheme.Scheme
import dispatch._

// Gist trait for doing SSL with mutual certificate exchange using dispatch.
// This works without having to set up global JDK-wide keystore and truststore files.
//
// Assumes the following:
//   You have the server certificate from the site you are calling and you have created a JKS
//   format keystore containing this file, which will act as your truststore:
//     > keytool -import -alias foo -file server.crt -keystore my.truststore
//
//   You have been issued with a client certificate and private key and these have been issued
//   as a pkcs12 format keystore.
//
trait SSLCommunications {

  // Implement these methods to define the keystore and truststore locations and the passwords for each
  protected val keystoreLocation: String
  protected val keystorePassword: Array[Char]
  protected val truststoreLocation: String
  protected val truststorePassword: Array[Char]

  // Usage: 
  //   val targetUrl = new URL("https://www.example.com/someResource")
  //   val req = url(targetUrl.toString)
  //   Https(targetUrl)(req as_str))
  //
  def Https(targetUrl: URL) = new Http {
    schemeFor(targetUrl, sslSocketFactory) foreach { scheme =>
      client.getConnectionManager.getSchemeRegistry.register(scheme)
    }
  }

  private val SSLPort = 443

  private def schemeFor(url: URL, sslSocketFactory: SSLSocketFactory) = {
    val port = if ( url.getPort == -1 ) SSLPort else url.getPort

    if ( url.getProtocol == "https" ) Some(new Scheme("https", port, sslSocketFactory)) else None
  }

  private lazy val sslSocketFactory = createSSLSocketFactory

  private def createSSLSocketFactory = {
    val sslContext = javax.net.ssl.SSLContext.getInstance("TLS")
    sslContext.init(keyManagers, trustManagers, new SecureRandom())
    new SSLSocketFactory(sslContext)
  }

  private def keyManagers = {
    val factory = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm)
    factory.init(populateKeystore("PKCS12", keystoreLocation, keystorePassword), keystorePassword)
    factory.getKeyManagers
  }

  private def trustManagers = {
    val factory = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm)
    factory.init(populateKeystore("JKS", truststoreLocation, truststorePassword))
    factory.getTrustManagers
  }

  private def populateKeystore(keystoreType: String, location: String, password: Array[Char]) = {
    val keyStore = KeyStore.getInstance(keystoreType)
    val is = getClass.getResourceAsStream(location) 
    if ( is == null ) throw new IllegalArgumentException("Unable to load the keystore at the given location: " + location)
    try {
      keyStore.load(is, password)
    } finally {
      is.close()
    }
    keyStore
  }
}