################################################################################
# Copyright 2016 Google Inc.
#
# 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.
################################################################################

"""Utility class for Provisioning SDK testing."""

import os

import pywrapcertificate_type
import pywrapprovisioning_status
from protos.public import certificate_provisioning_pb2

TEST_DATA_FOLDER = os.path.join('example', 'example_data')


def GetTestData(filename):
  current_dir = os.path.realpath(os.path.dirname(__file__))
  while not os.path.isdir(os.path.join(current_dir, TEST_DATA_FOLDER)):
    current_dir = os.path.dirname(current_dir)
  filename = os.path.join(current_dir, TEST_DATA_FOLDER, filename)
  with open(filename, 'rb') as data_file:
    data = data_file.read()
  return data


SERVICE_DRM_CERT = GetTestData('service.cert')
SERVICE_PUBLIC_KEY = GetTestData('service.public')
SERVICE_PRIVATE_KEY = GetTestData('service.encrypted.private')
SERVICE_PRIVATE_KEY_PASS = GetTestData('service.passphrase')
PROVISIONER_DRM_CERT = GetTestData('provisioner.cert')
PROVISIONER_PRIVATE_KEY = GetTestData('provisioner.encrypted.private')
PROVISIONER_PRIVATE_KEY_PASS = GetTestData('provisioner.passphrase')
PROVISIONER_SPOID_SECRET = GetTestData('provisioner.spoid_secret')
CA_PUBLIC_KEY = GetTestData('intermediate.public')
DEVICE_PUBLIC_KEY = GetTestData('user.public')
DEVICE_PRIVATE_KEY = GetTestData('user.private')
MESSAGE = GetTestData('message')


def InitProvisionEngineWithTestData(engine, verify_success=False):
  """Initialize the provisioning engine with sample credentials.

  Args:
    engine: a pywrapprovisioning_engine.ProvisioningEngine instance
    verify_success: whether to verify that resulting status code equals OK

  Returns:
    OK on success, or an appropriate error status code otherwise.
  """
  status = engine.Initialize(pywrapcertificate_type.kCertTesting,
                             SERVICE_DRM_CERT, SERVICE_PRIVATE_KEY,
                             SERVICE_PRIVATE_KEY_PASS, PROVISIONER_DRM_CERT,
                             PROVISIONER_PRIVATE_KEY,
                             PROVISIONER_PRIVATE_KEY_PASS,
                             PROVISIONER_SPOID_SECRET)
  if verify_success:
    AssertSuccess(status, 'Failed to initialize.')
  return status


def SetCertificateStatusListWithTestData(engine,
                                         expiration_period_seconds,
                                         verify_success=False):
  """Set the certificate status list with sample certificate status list.

  Args:
    engine: a pywrapprovisioning_engine.ProvisioningEngine instance
    expiration_period_seconds: number of seconds until certificate_status_list
      expires after its creation time
    verify_success: whether to verify that resulting status code equals OK

  Returns:
    OK on success, or an appropriate error status code otherwise.
  """
  certificate_status_list = GetTestData('certificate_list')

  status = engine.SetCertificateStatusList(certificate_status_list,
                                           expiration_period_seconds)

  if verify_success:
    AssertSuccess(status, 'Failed to set certificate status list.')

  return status


def AddDrmIntermediateCertificateWithTestData(engine,
                                              system_id,
                                              verify_success=False):
  """Generate an intermediate DRM cert and add it to provisioning engine.

  The intermediate DRM certificate is generated with sample public key and
  is added to the provisioning engine with sample certificate private key and
  passphrase.

  Args:
    engine: a pywrapprovisioning_engine.ProvisioningEngine instance
    system_id: Widevine system ID for the type of device
    verify_success: whether to verify that resulting status code equals OK

  Returns:
    OK on success, or an appropriate error status code otherwise.
  """
  ca_private_key = GetTestData('intermediate.encrypted.private')
  ca_private_key_passphrase = GetTestData('intermediate.passphrase')

  gen_status, ca_certificate = engine.GenerateDrmIntermediateCertificate(
      system_id, CA_PUBLIC_KEY)
  AssertSuccess(gen_status, 'Failed to generate intermediate certificate.')

  add_ca_status = engine.AddDrmIntermediateCertificate(
      ca_certificate, ca_private_key, ca_private_key_passphrase)

  if verify_success:
    AssertSuccess(add_ca_status, 'Failed to add intermediate certificate.')

  return add_ca_status


def NewProvisioningSessionWithTestData(engine, verify_success=False):
  """Create a provisioning session with sample device public and private keys.

  Args:
    engine: a pywrapprovisioning_engine.ProvisioningEngine instance
    verify_success: whether to verify that resulting status code equals OK

  Returns:
      status: OK on success, or an appropriate error status code otherwise.
      new_session: A new provisioning_session.
  """
  status, new_session = engine.NewProvisioningSession(DEVICE_PUBLIC_KEY,
                                                      DEVICE_PRIVATE_KEY)

  if verify_success:
    AssertSuccess(status, 'Failed to create session.')

  return (status, new_session)


def AssertSuccess(status, message=None):
  """Assert status equals OK."""
  assert pywrapprovisioning_status.OK == status, message


def ConvertToSignedProvisioningMessage(serialized_message):
  signed_message = certificate_provisioning_pb2.SignedProvisioningMessage()
  signed_message.ParseFromString(serialized_message)
  return signed_message