################################################################################
# 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.
################################################################################
"""Utility class for Provisioning SDK testing."""

import logging

from certificate_type import CertificateType
import test_data_provider
from provisioning_status import ProvisioningStatus
from protos.public import certificate_provisioning_pb2

logging.basicConfig(level=logging.DEBUG)


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

  Args:
    engine: a ProvisioningEngine instance
    verify_success: whether to verify that resulting status code equals OK
    cert_type: The type of certificate to use for initializing SDK -
      {kCertificateTypeTesting/kCertificateTypeDevelopment}

  Returns:
    OK on success, or an appropriate error status code otherwise.
  """
  data_provider = test_data_provider.TestDataProvider(cert_type)
  logging.info('Adding service certificate.')

  logging.info('Initializing provisioning engine with test data.')
  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:
    assert ProvisioningStatus.OK == status

  return status


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

  Args:
    engine: a 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 -
      {kCertificateTypeTesting/kCertificateTypeDevelopment}

  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:
    assert ProvisioningStatus.OK == status

  return status


def AddDrmIntermediateCertificateWithTestData(
    engine,
    system_id,
    verify_success=False,
    cert_type=CertificateType.kCertificateTypeTesting):
  """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 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 -
      {kCertificateTypeTesting/kCertificateTypeDevelopment}

  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)
  assert ProvisioningStatus.OK == gen_status

  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:
    assert ProvisioningStatus.OK == add_ca_status

  return add_ca_status


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

  Args:
    engine: a 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 -
      {kCertificateTypeTesting/kCertificateTypeDevelopment}

  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:
    assert ProvisioningStatus.OK == gen_status
  return ca_certificate


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

  Args:
    engine: a ProvisioningEngine instance
    verify_success: whether to verify that resulting status code equals OK
    cert_type: The type of certificate to use for initializing SDK -
      {kCertificateTypeTesting/kCertificateTypeDevelopment}

  Returns:
      new_session: A new provisioning_session.
      status: OK on success, or an appropriate error status code otherwise.
  """
  logging.info('Starting a new provisioning session with'
               'sample device public and private keys.')
  data_provider = test_data_provider.TestDataProvider(cert_type)
  new_session, status = engine.NewProvisioningSession(
      certificate_provisioning_pb2.SignedProvisioningMessage.PROVISIONING_30,
      data_provider.device_public_key, data_provider.device_private_key)
  if verify_success:
    assert (ProvisioningStatus.OK == status), 'status = %r' % status

  return new_session, status


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