//
// Copyright 2021 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine License
// Agreement.
//

#ifndef WV_DRM_PLUGIN_H_
#define WV_DRM_PLUGIN_H_

#include <aidl/android/hardware/drm/BnDrmPlugin.h>
#include <aidl/android/hardware/drm/IDrmPluginListener.h>
#include <aidl/android/hardware/drm/Status.h>
#include <utils/Mutex.h>

#include <list>
#include <map>

#include "OEMCryptoCENC.h"
#include "WVGenericCryptoInterface.h"
#include "WVTypes.h"
#include "cdm_client_property_set.h"
#include "cdm_identifier.h"
#include "log.h"
#include "wv_cdm_event_listener.h"
#include "wv_content_decryption_module.h"

namespace wvdrm {
namespace hardware {
namespace drm {
namespace widevine {

using ::aidl::android::hardware::drm::Status;

using wvcdm::CdmIdentifier;
using wvcdm::CdmKeyStatusMap;
using wvcdm::CdmResponseType;
using wvcdm::CdmSessionId;
using wvcdm::WvContentDecryptionModule;

const OEMCrypto_Algorithm kInvalidCryptoAlgorithm =
    static_cast<OEMCrypto_Algorithm>(-1);

class WVDrmPlugin : public ::aidl::android::hardware::drm::BnDrmPlugin,
                     wvcdm::WvCdmEventListener {
 public:
  WVDrmPlugin(const android::sp<WvContentDecryptionModule>& cdm,
              const std::string& appPackageName,
              WVGenericCryptoInterface* crypto, bool useSpoid);

  virtual ~WVDrmPlugin();

  void Close();

  ::ndk::ScopedAStatus closeSession(
      const std::vector<uint8_t>& in_sessionId) override;
  ::ndk::ScopedAStatus decrypt(const std::vector<uint8_t>& in_sessionId,
                               const std::vector<uint8_t>& in_keyId,
                               const std::vector<uint8_t>& in_input,
                               const std::vector<uint8_t>& in_iv,
                               std::vector<uint8_t>* _aidl_return) override;
  ::ndk::ScopedAStatus encrypt(const std::vector<uint8_t>& in_sessionId,
                               const std::vector<uint8_t>& in_keyId,
                               const std::vector<uint8_t>& in_input,
                               const std::vector<uint8_t>& in_iv,
                               std::vector<uint8_t>* _aidl_return) override;
  ::ndk::ScopedAStatus getHdcpLevels(
      ::aidl::android::hardware::drm::HdcpLevels* _aidl_return) override;
  ::ndk::ScopedAStatus getKeyRequest(
      const std::vector<uint8_t>& in_scope,
      const std::vector<uint8_t>& in_initData, const std::string& in_mimeType,
      ::aidl::android::hardware::drm::KeyType in_keyType,
      const std::vector<::aidl::android::hardware::drm::KeyValue>&
          in_optionalParameters,
      ::aidl::android::hardware::drm::KeyRequest* _aidl_return) override;
  ::ndk::ScopedAStatus getLogMessages(
      std::vector<::aidl::android::hardware::drm::LogMessage>* _aidl_return)
      override;
  ::ndk::ScopedAStatus getMetrics(
      std::vector<::aidl::android::hardware::drm::DrmMetricGroup>* _aidl_return)
      override;
  ::ndk::ScopedAStatus getNumberOfSessions(
      ::aidl::android::hardware::drm::NumberOfSessions* _aidl_return) override;
  ::ndk::ScopedAStatus getOfflineLicenseKeySetIds(
      std::vector<::aidl::android::hardware::drm::KeySetId>* _aidl_return)
      override;
  ::ndk::ScopedAStatus getOfflineLicenseState(
      const ::aidl::android::hardware::drm::KeySetId& in_keySetId,
      ::aidl::android::hardware::drm::OfflineLicenseState* _aidl_return)
      override;
  ::ndk::ScopedAStatus getPropertyByteArray(
      const std::string& in_propertyName,
      std::vector<uint8_t>* _aidl_return) override;
  ::ndk::ScopedAStatus getPropertyString(const std::string& in_propertyName,
                                         std::string* _aidl_return) override;
  ::ndk::ScopedAStatus getProvisionRequest(
      const std::string& in_certificateType,
      const std::string& in_certificateAuthority,
      ::aidl::android::hardware::drm::ProvisionRequest* _aidl_return) override;
  ::ndk::ScopedAStatus getSecureStop(
      const ::aidl::android::hardware::drm::SecureStopId& in_secureStopId,
      ::aidl::android::hardware::drm::SecureStop* _aidl_return) override;
  ::ndk::ScopedAStatus getSecureStopIds(
      std::vector<::aidl::android::hardware::drm::SecureStopId>* _aidl_return)
      override;
  ::ndk::ScopedAStatus getSecureStops(
      std::vector<::aidl::android::hardware::drm::SecureStop>* _aidl_return)
      override;
  ::ndk::ScopedAStatus getSecurityLevel(
      const std::vector<uint8_t>& in_sessionId,
      ::aidl::android::hardware::drm::SecurityLevel* _aidl_return) override;
  ::ndk::ScopedAStatus openSession(
      ::aidl::android::hardware::drm::SecurityLevel in_securityLevel,
      std::vector<uint8_t>* _aidl_return) override;
  ::ndk::ScopedAStatus provideKeyResponse(
      const std::vector<uint8_t>& in_scope,
      const std::vector<uint8_t>& in_response,
      ::aidl::android::hardware::drm::KeySetId* _aidl_return) override;
  ::ndk::ScopedAStatus provideProvisionResponse(
      const std::vector<uint8_t>& in_response,
      ::aidl::android::hardware::drm::ProvideProvisionResponseResult*
          _aidl_return) override;
  ::ndk::ScopedAStatus queryKeyStatus(
      const std::vector<uint8_t>& in_sessionId,
      std::vector<::aidl::android::hardware::drm::KeyValue>* _aidl_return)
      override;
  ::ndk::ScopedAStatus releaseAllSecureStops() override;
  ::ndk::ScopedAStatus releaseSecureStop(
      const ::aidl::android::hardware::drm::SecureStopId& in_secureStopId)
      override;
  ::ndk::ScopedAStatus releaseSecureStops(
      const ::aidl::android::hardware::drm::OpaqueData& in_ssRelease) override;
  ::ndk::ScopedAStatus removeAllSecureStops() override;
  ::ndk::ScopedAStatus removeKeys(
      const std::vector<uint8_t>& in_sessionId) override;
  ::ndk::ScopedAStatus removeOfflineLicense(
      const ::aidl::android::hardware::drm::KeySetId& in_keySetId) override;
  ::ndk::ScopedAStatus removeSecureStop(
      const ::aidl::android::hardware::drm::SecureStopId& in_secureStopId)
      override;
  ::ndk::ScopedAStatus requiresSecureDecoder(
      const std::string& in_mime,
      ::aidl::android::hardware::drm::SecurityLevel in_level,
      bool* _aidl_return) override;
  ::ndk::ScopedAStatus restoreKeys(
      const std::vector<uint8_t>& in_sessionId,
      const ::aidl::android::hardware::drm::KeySetId& in_keySetId) override;
  ::ndk::ScopedAStatus setCipherAlgorithm(
      const std::vector<uint8_t>& in_sessionId,
      const std::string& in_algorithm) override;
  ::ndk::ScopedAStatus setListener(
      const std::shared_ptr<::aidl::android::hardware::drm::IDrmPluginListener>&
          in_listener) override;
  ::ndk::ScopedAStatus setMacAlgorithm(
      const std::vector<uint8_t>& in_sessionId,
      const std::string& in_algorithm) override;
  ::ndk::ScopedAStatus setPlaybackId(const std::vector<uint8_t>& in_sessionId,
                                     const std::string& in_playbackId) override;
  ::ndk::ScopedAStatus setPropertyByteArray(
      const std::string& in_propertyName,
      const std::vector<uint8_t>& in_value) override;
  ::ndk::ScopedAStatus setPropertyString(const std::string& in_propertyName,
                                         const std::string& in_value) override;
  ::ndk::ScopedAStatus sign(const std::vector<uint8_t>& in_sessionId,
                            const std::vector<uint8_t>& in_keyId,
                            const std::vector<uint8_t>& in_message,
                            std::vector<uint8_t>* _aidl_return) override;
  ::ndk::ScopedAStatus signRSA(const std::vector<uint8_t>& in_sessionId,
                               const std::string& in_algorithm,
                               const std::vector<uint8_t>& in_message,
                               const std::vector<uint8_t>& in_wrappedkey,
                               std::vector<uint8_t>* _aidl_return) override;
  ::ndk::ScopedAStatus verify(const std::vector<uint8_t>& in_sessionId,
                              const std::vector<uint8_t>& in_keyId,
                              const std::vector<uint8_t>& in_message,
                              const std::vector<uint8_t>& in_signature,
                              bool* _aidl_return) override;

  // The following methods do not use HAL interface, it is used internally.
  virtual Status unprovisionDevice();

  virtual void OnSessionRenewalNeeded(const CdmSessionId& cdmSessionId);
  virtual void OnSessionKeysChange(const CdmSessionId& cdmSessionId,
                                   const CdmKeyStatusMap& cdmKeysStatus,
                                   bool hasNewUsableKey);
  virtual void OnExpirationUpdate(const CdmSessionId& cdmSessionId,
                                  int64_t newExpiryTimeSeconds);
  virtual void OnSessionLostState(const CdmSessionId& cdmSessionId);

 private:
  WVDRM_DISALLOW_COPY_AND_ASSIGN_AND_NEW(WVDrmPlugin);

  // List this field first so it is destructed last; ensure logging uid
  // is cleared right before plugin is destructed.
  wvutil::LoggingUidSetter mLoggingUidSetter;

  struct CryptoSession {
   public:
    CryptoSession()
        : mOecSessionId(-1),
          mCipherAlgorithm(kInvalidCryptoAlgorithm),
          mMacAlgorithm(kInvalidCryptoAlgorithm) {}

    CryptoSession(OEMCrypto_SESSION sessionId)
        : mOecSessionId(sessionId),
          mCipherAlgorithm(kInvalidCryptoAlgorithm),
          mMacAlgorithm(kInvalidCryptoAlgorithm) {}

    OEMCrypto_SESSION oecSessionId() const { return mOecSessionId; }

    OEMCrypto_Algorithm cipherAlgorithm() const { return mCipherAlgorithm; }

    void setCipherAlgorithm(OEMCrypto_Algorithm newAlgorithm) {
      mCipherAlgorithm = newAlgorithm;
    }

    OEMCrypto_Algorithm macAlgorithm() const { return mMacAlgorithm; }

    void setMacAlgorithm(OEMCrypto_Algorithm newAlgorithm) {
      mMacAlgorithm = newAlgorithm;
    }

   private:
    OEMCrypto_SESSION mOecSessionId;
    OEMCrypto_Algorithm mCipherAlgorithm;
    OEMCrypto_Algorithm mMacAlgorithm;
  };

  class WVClientPropertySet : public wvcdm::CdmClientPropertySet {
   public:
    WVClientPropertySet()
        : mUsePrivacyMode(false),
          mShareKeys(false),
          mSessionSharingId(0),
          mUseAtscMode(false) {}

    virtual ~WVClientPropertySet() {}

    virtual const std::string& security_level() const { return mSecurityLevel; }

    void set_security_level(const std::string& securityLevel) {
      mSecurityLevel = securityLevel;
    }

    virtual bool use_privacy_mode() const { return mUsePrivacyMode; }

    void set_use_privacy_mode(bool usePrivacyMode) {
      mUsePrivacyMode = usePrivacyMode;
    }

    virtual const std::string& service_certificate() const {
      return mServiceCertificate;
    }

    virtual void set_service_certificate(
        const std::string& serviceCertificate) {
      mServiceCertificate = serviceCertificate;
    }

    virtual const std::string& device_provisioning_service_certificate() const {
      // Android does not support service certificates for provisioning.
      return mEmptyString;
    }

    virtual void set_device_provisioning_service_certificate(
        const std::string&) {
      // Ignore. Android does not support service certificates for provisioning
      // TODO(b/69562876): Android SHOULD support service cert for provisioning
    }

    virtual bool is_session_sharing_enabled() const { return mShareKeys; }

    void set_is_session_sharing_enabled(bool shareKeys) {
      mShareKeys = shareKeys;
    }

    virtual uint32_t session_sharing_id() const { return mSessionSharingId; }

    virtual void set_session_sharing_id(uint32_t id) { mSessionSharingId = id; }

    virtual const std::string& app_id() const { return mAppId; }

    void set_app_id(const std::string& appId) { mAppId = appId; }

    virtual bool use_atsc_mode() const { return mUseAtscMode; }

    void set_use_atsc_mode(bool useAtscMode) { mUseAtscMode = useAtscMode; }

   private:
    DISALLOW_EVIL_CONSTRUCTORS(WVClientPropertySet);

    std::string mSecurityLevel;
    bool mUsePrivacyMode;
    std::string mServiceCertificate;
    bool mShareKeys;
    uint32_t mSessionSharingId;
    std::string mAppId;
    bool mUseAtscMode;
    const std::string mEmptyString;
  } mPropertySet;

  class CdmIdentifierBuilder {
   public:
    CdmIdentifierBuilder(bool useSpoid, const WVDrmPlugin& parent,
                         const std::string& appPackageName);

    // Fills in the passed-in struct with the CDM Identifier for the current
    // combination of Origin, Application, and Device. This is needed by some
    // calls into the CDM in order to identify which CDM instance should receive
    // the call. Calling this will seal the CDM Identifier Builder, thus making
    // it an error to change the origin.
    WvStatus getCdmIdentifier(CdmIdentifier* identifier);

    // Gets the application-safe device-unique ID. On non-SPOID devices, this is
    // the device-unique ID from OEMCrypto. On SPOID devices, this is the SPOID.
    // On SPOID devices, calling this will seal the CDM Identifier Builder, thus
    // making it an error to change the origin.
    WvStatus getDeviceUniqueId(std::string* id);
    WvStatus getProvisioningUniqueId(std::string* id);

    const std::string& origin() const { return mCdmIdentifier.origin; }
    bool set_origin(const std::string& id);

    // This sets the app package name to allow apps to access ATSC licenses
    bool set_use_atsc_mode(bool enable);

    // Indicates whether the builder can still be modified. This returns false
    // until a call to getCdmIdentifier.
    bool is_sealed() { return mIsIdentifierSealed; }

    uint32_t user_id() const { return mCdmIdentifier.user_id; }

   private:
    WVDRM_DISALLOW_COPY_AND_ASSIGN(CdmIdentifierBuilder);

    CdmIdentifier mCdmIdentifier;
    bool mIsIdentifierSealed;

    bool mUseSpoid;
    std::string mAppPackageName;
    const WVDrmPlugin& mParent;

    WvStatus calculateSpoid();
    Status calculateSpoid(const std::string& deviceID, std::string* spoid);

    // Gets the device-unique ID from OEMCrypto. This must be private, since
    // this value must not be exposed to applications on SPOID devices. Code
    // outside this class should use getDeviceUniqueId() to get the
    // application-safe device-unique ID.
    WvStatus getOemcryptoDeviceId(std::string* id);
    WvStatus getOemcryptoDeviceId(wvcdm::RequestedSecurityLevel securityLevel,
                                  std::string* id);

    // The unique identifier is meant to ensure that two clients with the
    // same spoid, origin and app package name still get different cdm engine
    // instances. This is a stepping stone to simplifying the implementation.
    // Note that we do not have a lock or mutex around this object. We assume
    // that locking is handled external to this object.
    uint32_t getNextUniqueId();
  } mCdmIdentifierBuilder;

  // Properly close plugins on SIGTERM then exit
  class Terminator {
   public:
    static void Register(WVDrmPlugin* plugin) { instance().DoRegister(plugin); }
    static void Forget(WVDrmPlugin* plugin) { instance().DoForget(plugin); }
    static void Terminate(int /*signal*/) { instance().DoTerminate(); }

   private:
    WVDRM_DISALLOW_COPY_AND_ASSIGN(Terminator);

    Terminator() { signal(SIGTERM, Terminate); }

    static Terminator& instance() {
      static Terminator instance;
      return instance;
    }

    void DoRegister(WVDrmPlugin* plugin) {
      android::Mutex::Autolock lock(mLock);
      mPlugins.push_back(plugin);
    }

    void DoForget(WVDrmPlugin* plugin) {
      android::Mutex::Autolock lock(mLock);
      mPlugins.remove(plugin);
    }

    void DoTerminate() {
      android::Mutex::Autolock lock(mLock);
      for_each(mPlugins.begin(), mPlugins.end(),
               [](WVDrmPlugin* plugin) { plugin->Close(); });
      exit(0);
    }

    std::list<WVDrmPlugin*> mPlugins;
    android::Mutex mLock;
  };

  android::sp<wvcdm::WvContentDecryptionModule> const mCDM;
  WVGenericCryptoInterface* mCrypto;
  std::map<CdmSessionId, CryptoSession> mCryptoSessions;
  std::shared_ptr<::aidl::android::hardware::drm::IDrmPluginListener> mListener;

  std::string mProvisioningServiceCertificate;

  wvcdm::CdmSessionId mDecryptHashSessionId;

  std::string mAppPackageName;

  ::ndk::SpAIBinder createBinder() override;

  WvStatus queryProperty(const std::string& property,
                         std::string& stringValue) const;

  WvStatus queryProperty(wvcdm::RequestedSecurityLevel securityLevel,
                         const std::string& property,
                         std::string& stringValue) const;

  WvStatus queryProperty(const std::string& property,
                         std::vector<uint8_t>& vector_value) const;

  bool isProvisioned(wvcdm::CdmSecurityLevel securityLevel,
                     const std::string& origin, const std::string& spoid,
                     bool atsc_mode_enabled) const;

  WvStatus mapAndNotifyOfCdmResponseType(const std::vector<uint8_t>& sessionId,
                                         CdmResponseType res);

  void notifyOfCdmResponseType(const std::vector<uint8_t>& sessionId,
                               CdmResponseType res);

  Status mapAndNotifyOfOEMCryptoResult(const std::vector<uint8_t>& sessionId,
                                       OEMCryptoResult res);

  Status mapOEMCryptoResult(OEMCryptoResult res);

  ::aidl::android::hardware::drm::SecurityLevel mapSecurityLevel(
      const std::string& level);

  wvcdm::RequestedSecurityLevel getRequestedSecurityLevel() const;

  WvStatus openSessionCommon(std::vector<uint8_t>& sessionId);

  bool initDataResemblesPSSH(const std::vector<uint8_t>& initData);

  WvStatus unprovision(const CdmIdentifier& identifier);

  void sendEvent(::aidl::android::hardware::drm::EventType in_eventType,
                 const std::vector<uint8_t>& in_sessionId,
                 const std::vector<uint8_t>& in_data);
  void sendExpirationUpdate(const std::vector<uint8_t>& in_sessionId,
                            int64_t in_expiryTimeInMS);
  void sendKeysChange(
      const std::vector<uint8_t>& in_sessionId,
      const std::vector<::aidl::android::hardware::drm::KeyStatus>&
          in_keyStatusList,
      bool in_hasNewUsableKey);
  void sendSessionLostState(const std::vector<uint8_t>& in_sessionId);
};

}  // namespace widevine
}  // namespace drm
}  // namespace hardware
}  // namespace wvdrm

#endif  // WV_DRM_PLUGIN_H_