################################################################################
# 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 logging

import pywrapcertificate_type
import pywrapprovisioning_status
import test_data_provider
from protos.public import certificate_provisioning_pb2

logging.basicConfig(level=logging.DEBUG)


def InitProvisionEngineWithTestData(
    engine, verify_success=False,
    cert_type=pywrapcertificate_type.kCertTesting):
  """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
    cert_type: The type of certificate to use for initializing SDK -
      {kCertTesting/kCertDevelopment}

  Returns:
    OK on success, or an appropriate error status code otherwise.
  """
  logging.info('Initializing provisioning engine with test data.')
  data_provider = test_data_provider.TestDataProvider(cert_type)
  status = engine.Initialize(cert_type,
                             data_provider.service_drm_cert,
                             data_provider.service_private_key,
                             data_provider.service_private_key_passphrase,
                             data_provider.provisioner_drm_cert,
                             data_provider.provisioner_private_key,
                             data_provider.provisioner_private_key_passphrase,
                             data_provider.provisioner_spoid_secret)
  if verify_success:
    AssertSuccess(status, 'Failed to initialize.')
  return status


def SetCertificateStatusListWithTestData(
    engine, expiration_period_seconds, verify_success=False,
    cert_type=pywrapcertificate_type.kCertTesting):
  """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
    cert_type: The type of certificate to use for initializing SDK -
      {kCertTesting/kCertDevelopment}

  Returns:
    OK on success, or an appropriate error status code otherwise.
  """
  logging.info('Setting certificate status list with test data.')
  data_provider = test_data_provider.TestDataProvider(cert_type)
  certificate_status_list = data_provider.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,
    cert_type=pywrapcertificate_type.kCertTesting):
  """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
    cert_type: The type of certificate to use for initializing SDK -
      {kCertTesting/kCertDevelopment}

  Returns:
    OK on success, or an appropriate error status code otherwise.
  """
  logging.info(
      'Generating DRM intermediate certificate for system_id <%d>.', system_id)
  data_provider = test_data_provider.TestDataProvider(cert_type)
  gen_status, ca_certificate = engine.GenerateDrmIntermediateCertificate(
      system_id, data_provider.ca_public_key)
  AssertSuccess(gen_status, 'Failed to generate intermediate certificate.')

  logging.info('Adding DRM intermediate certificate.')
  add_ca_status = engine.AddDrmIntermediateCertificate(
      ca_certificate, data_provider.ca_private_key,
      data_provider.ca_private_key_passphrase)

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

  return add_ca_status


def GenerateDeviceDrmCertificate(engine, system_id, serial_number,
                                 verify_success=False,
                                 cert_type=pywrapcertificate_type.kCertTesting):
  """Generate a device DRM certificate.

  Args:
    engine: a pywrapprovisioning_engine.ProvisioningEngine instance
    system_id: Widevine system ID for the type of device
    serial_number: The serial number for the device DRM certificate.
    verify_success: whether to verify that resulting status code equals OK
    cert_type: The type of certificate to use for initializing SDK -
      {kCertTesting/kCertDevelopment}

  Returns:
    OK on success, or an appropriate error status code otherwise.
  """
  logging.info(
      'Generating Device cert for system_id <%d> and serial_number <%s>.',
      system_id, serial_number)
  data_provider = test_data_provider.TestDataProvider(cert_type)
  gen_status, ca_certificate = engine.GenerateDeviceDrmCertificate(
      system_id, data_provider.device_public_key, serial_number)
  if verify_success:
    AssertSuccess(gen_status, 'Failed to generate device DRM certificate.')
  return ca_certificate


def NewProvisioningSessionWithTestData(
    engine, verify_success=False,
    cert_type=pywrapcertificate_type.kCertTesting):
  """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
    cert_type: The type of certificate to use for initializing SDK -
      {kCertTesting/kCertDevelopment}

  Returns:
      status: OK on success, or an appropriate error status code otherwise.
      new_session: A new provisioning_session.
  """
  logging.info('Starting a new provisioning session with'
               'sample device public and private keys.')
  data_provider = test_data_provider.TestDataProvider(cert_type)
  status, new_session = engine.NewProvisioningSession(
      data_provider.device_public_key, data_provider.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