<commit_before>package com.blinkbox.books.auth.server.sso

import java.io.IOException
import java.nio.file.{Files, Paths}
import java.security.spec.{InvalidKeySpecException, X509EncodedKeySpec}
import java.security.{InvalidKeyException, KeyFactory, NoSuchAlgorithmException, PublicKey}

import com.blinkbox.security.jwt.signatures.{ES256, SignatureAlgorithm}
import com.blinkbox.security.jwt.{InvalidTokenException, TokenDecoder, UnsupportedTokenException}

import scala.collection.JavaConverters._
import scala.collection.concurrent.TrieMap
import scala.util.Try

trait KeyStore {
  def verificationKey: PublicKey
}

class FileKeyStore(keyFile: String) extends KeyStore {
  private val DefaultKeyId = ""  // see SSO-369; i'm assuming we'll have more than one key in future
  private val verificationKeys = new TrieMap[String, PublicKey]()

  def verificationKey: PublicKey =
    // not using getOrElseUpdate as this is inherited from MapLike and isn't thread-safe
    // see this bug report for more info: https://issues.scala-lang.org/browse/SI-7943
    // also note that this method may read the file more than once, but that's alright
    verificationKeys.getOrElse(DefaultKeyId, try {
      val keyData = Files.readAllBytes(Paths.get(keyFile))
      val key = KeyFactory.getInstance("EC").generatePublic(new X509EncodedKeySpec(keyData))
      verificationKeys.putIfAbsent(DefaultKeyId, key).getOrElse(key)
    } catch {
      case e: IOException => throw new InvalidTokenException(s"Failed to load key.")
      case e: NoSuchAlgorithmException => throw new UnsupportedTokenException("The ECDSA signature algorithm is not supported.", e)
      case e: InvalidKeySpecException => throw new InvalidKeyException("The public key is invalid.", e)
    })
}

class SsoAccessTokenDecoder(keyStore: KeyStore) extends TokenDecoder {
  override def getVerifier(header: java.util.Map[String, AnyRef]): SignatureAlgorithm =
    if (header.get("alg") == ES256.NAME) new ES256(keyStore.verificationKey)
    else super.getVerifier(header)
}

case class SsoAccessToken(token: String, claims: Map[String, AnyRef]) {
  val subject: String = claims.get("sub") match {
    case Some(s: String) => s
    case _ => throw new InvalidTokenException("The 'sub' claim is missing or invalid.")
  }
}

object SsoAccessToken {
  def decode(token: String, decoder: TokenDecoder): Try[SsoAccessToken] = Try {
    decoder.decode(token) match {
      case claims: java.util.Map[_, _] => SsoAccessToken(token, claims.asInstanceOf[java.util.Map[String, AnyRef]].asScala.toMap)
      case _ => throw new InvalidTokenException("The token does not contain valid claims.")
    }
  }
}<commit_msg>Avoid casting by using ClassTags<commit_after>package com.blinkbox.books.auth.server.sso

import java.io.IOException
import java.nio.file.{Files, Paths}
import java.security.spec.{InvalidKeySpecException, X509EncodedKeySpec}
import java.security.{InvalidKeyException, KeyFactory, NoSuchAlgorithmException, PublicKey}

import com.blinkbox.security.jwt.signatures.{ES256, SignatureAlgorithm}
import com.blinkbox.security.jwt.{InvalidTokenException, TokenDecoder, UnsupportedTokenException}

import scala.collection.JavaConverters._
import scala.collection.concurrent.TrieMap
import scala.reflect.ClassTag
import scala.util.Try

trait KeyStore {
  def verificationKey: PublicKey
}

class FileKeyStore(keyFile: String) extends KeyStore {
  private val DefaultKeyId = ""  // see SSO-369; i'm assuming we'll have more than one key in future
  private val verificationKeys = new TrieMap[String, PublicKey]()

  def verificationKey: PublicKey =
    // not using getOrElseUpdate as this is inherited from MapLike and isn't thread-safe
    // see this bug report for more info: https://issues.scala-lang.org/browse/SI-7943
    // also note that this method may read the file more than once, but that's alright
    verificationKeys.getOrElse(DefaultKeyId, try {
      val keyData = Files.readAllBytes(Paths.get(keyFile))
      val key = KeyFactory.getInstance("EC").generatePublic(new X509EncodedKeySpec(keyData))
      verificationKeys.putIfAbsent(DefaultKeyId, key).getOrElse(key)
    } catch {
      case e: IOException => throw new InvalidTokenException(s"Failed to load key.")
      case e: NoSuchAlgorithmException => throw new UnsupportedTokenException("The ECDSA signature algorithm is not supported.", e)
      case e: InvalidKeySpecException => throw new InvalidKeyException("The public key is invalid.", e)
    })
}

class SsoAccessTokenDecoder(keyStore: KeyStore) extends TokenDecoder {
  override def getVerifier(header: java.util.Map[String, AnyRef]): SignatureAlgorithm =
    if (header.get("alg") == ES256.NAME) new ES256(keyStore.verificationKey)
    else super.getVerifier(header)
}

case class SsoAccessToken(token: String, claims: Map[String, AnyRef]) {
  val subject: String = claims.get("sub") match {
    case Some(s: String) => s
    case _ => throw new InvalidTokenException("The 'sub' claim is missing or invalid.")
  }
}

object SsoAccessToken {
  def decode(token: String, decoder: TokenDecoder): Try[SsoAccessToken] = Try {
    val claims = extractClaims[String, AnyRef](decoder.decode(token))
    SsoAccessToken(token, claims)
  }
  private def extractClaims[K: ClassTag, V: ClassTag](payload: Any): Map[K, V] = payload match {
    case claims: java.util.Map[K, V] => claims.asScala.toMap
    case _ => throw new InvalidTokenException("The token does not contain valid claims.")
  }
}
{'fraction_non_alphanumeric': 0.10012813472451035, 'fraction_numerical': 0.00805418268350723, 'mean_word_length': 4.843850267379679, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 17, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '17746', 'n_tokens_mistral': 1642, 'n_tokens_neox': 1586, 'n_words': 446}