////////////////////////////////////////////////////////////////////////////////
// Copyright 2016 Google LLC.
//
// This software is licensed under the terms defined in the Widevine Master
// License Agreement. For a copy of this agreement, please contact
// widevine-licensing@google.com.
////////////////////////////////////////////////////////////////////////////////

//
// Description:
//   Declaration of classes representing RSA private and public keys used
//   for message signing, signature verification, encryption and decryption.

#ifndef COMMON_RSA_KEY_H_
#define COMMON_RSA_KEY_H_

#include <cstdint>
#include <memory>
#include <string>

#include <cstdint>
#include "absl/base/macros.h"
#include "openssl/rsa.h"
#include "common/hash_algorithm.h"

namespace widevine {

class RsaPublicKey;

class RsaPrivateKey {
 public:
  explicit RsaPrivateKey(RSA* key);
  RsaPrivateKey(const RsaPrivateKey&);
  virtual ~RsaPrivateKey();

  // Create an RsaPrivateKey object using a DER encoded PKCS#1 RSAPrivateKey.
  // Returns NULL on failure.
  static RsaPrivateKey* Create(const std::string& serialized_key);

  // Decrypt a message using RSA-OAEP. Caller retains ownership of all
  // parameters. Returns true if successful, false otherwise.
  virtual bool Decrypt(const std::string& encrypted_message,
                       std::string* decrypted_message) const;


  // Generate RSSASSA-PSS signature. Caller retains ownership of all parameters.
  // |hash_algorithm| indicates the hash algorithm used. Returns true if
  // successful, false otherwise.
  virtual bool GenerateSignature(const std::string& message,
                                 HashAlgorithm hash_algorithm,
                                 std::string* signature) const;

  // Generate SHA256 digest, PKCS#7 padded signature. Caller retains ownership
  // of all parameters. Returns true if successful, false otherwise.
  virtual bool GenerateSignatureSha256Pkcs7(const std::string& message,
                                            std::string* signature) const;

  // Return true if the underlying key matches with |private_key|.
  virtual bool MatchesPrivateKey(const RsaPrivateKey& private_key) const;

  // Return true if the underlying key is a public-private key pair with
  // |public_key|.
  virtual bool MatchesPublicKey(const RsaPublicKey& public_key) const;

  // Returns the RSA key size (modulus) in bytes.
  virtual uint32_t KeySize() const;

 private:
  friend class RsaPublicKey;
  friend class X509CertificateBuilder;  // TODO(user): Get rid of this.

  const RSA* key() const { return key_; }

  RSA* key_;

// SWIG appears to think this declaration is a syntax error.  Excluding it for
// python SWIG wrapping.
#ifndef SWIG
  // Disallow assignment operator.
  RsaPrivateKey& operator=(const RsaPrivateKey&) = delete;
#endif  // SWIG
};

class RsaPublicKey {
 public:
  explicit RsaPublicKey(RSA* key);

  // Copy constructor.
  RsaPublicKey(const RsaPublicKey& rsa_key);

  // Construct RsaPublicKey object from RsaPrivateKey.
  explicit RsaPublicKey(const RsaPrivateKey& rsa_key);

  virtual ~RsaPublicKey();

  // Create an RsaPublicKey object using a DER encoded PKCS#1 RSAPublicKey.
  // Returns NULL on failure.
  static RsaPublicKey* Create(const std::string& serialized_key);

  // Encrypt a message using RSA-OAEP. Caller retains ownership of all
  // parameters. Returns true if successful, false otherwise.
  virtual bool Encrypt(const std::string& clear_message,
                       std::string* encrypted_message) const;


  // Verify RSSASSA-PSS signature. Caller retains ownership of all parameters.
  // |hash_algorithm| indicates the hash algorithm used. Returns true if
  // validation succeeds, false otherwise.
  virtual bool VerifySignature(const std::string& message,
                               HashAlgorithm hash_algorithm,
                               const std::string& signature) const;

  // Verify a signature. This method takes two parameters: |message| which is a
  // std::string containing the data which was signed, and |signature| which is a
  // std::string containing the message SHA256 digest signature with PKCS#7
  // padding. Returns true if verification succeeds, false otherwise.
  virtual bool VerifySignatureSha256Pkcs7(const std::string& message,
                                          const std::string& signature) const;

  // Return true if the underlying key is a public-private key pair with
  // |private_key|.
  virtual bool MatchesPrivateKey(const RsaPrivateKey& private_key) const;

  // Return true if the underlying key matches with |public_key|.
  virtual bool MatchesPublicKey(const RsaPublicKey& public_key) const;

  // Returns the RSA key size (modulus) in bytes.
  virtual uint32_t KeySize() const;

 private:
  friend class RsaPrivateKey;
  friend class X509CertificateBuilder;  // TODO(user): Get rid of this.

  const RSA* key() const { return key_; }

  RSA* key_;

// SWIG appears to think this declaration is a syntax error.  Excluding it for
// python SWIG wrapping.
#ifndef SWIG
  // Disallow assignment operator.
  RsaPublicKey& operator=(const RsaPublicKey&) = delete;
#endif  // SWIG
};

class RsaKeyFactory {
 public:
  RsaKeyFactory();

  RsaKeyFactory(const RsaKeyFactory&) = delete;
  RsaKeyFactory& operator=(const RsaKeyFactory&) = delete;

  virtual ~RsaKeyFactory();

  // Create an RsaPrivateKey object using a DER encoded PKCS#1 RSAPrivateKey.
  virtual std::unique_ptr<RsaPrivateKey> CreateFromPkcs1PrivateKey(
      const std::string& private_key) const;

  // Create a PKCS#1 RsaPrivateKey object using an PKCS#8 PrivateKeyInfo or
  // EncryptedPrivateKeyInfo (if |private_key_passprhase| is not empty).
  virtual std::unique_ptr<RsaPrivateKey> CreateFromPkcs8PrivateKey(
      const std::string& private_key,
      const std::string& private_key_passphrase) const;

  // Create an RsaPublicKey object using a DER encoded PKCS#1 RSAPublicKey.
  virtual std::unique_ptr<RsaPublicKey> CreateFromPkcs1PublicKey(
      const std::string& public_key) const;
};

}  // namespace widevine

#endif  // COMMON_RSA_KEY_H_