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

import unittest

import crypto_utility
import pywrapprovisioning_engine
import pywrapprovisioning_status
import test_data_utility
from protos.public import certificate_provisioning_pb2
from protos.public import signed_device_certificate_pb2


class NewSessionTest(unittest.TestCase):

  def setUp(self):
    self._engine = pywrapprovisioning_engine.ProvisioningEngine()
    test_data_utility.InitProvisionEngineWithTestData(
        self._engine, verify_success=True)
    test_data_utility.SetCertificateStatusListWithTestData(
        self._engine, 0, verify_success=True)

  def testNewSessionSuccess(self):
    test_data_utility.AddDrmIntermediateCertificateWithTestData(
        self._engine, 2001, verify_success=True)

    (_, new_session) = test_data_utility.NewProvisioningSessionWithTestData(
        self._engine, verify_success=True)
    (status, raw_response,
     _) = new_session.ProcessMessage(test_data_utility.MESSAGE)
    test_data_utility.AssertSuccess(status, 'Failed to create session.')

    signed_request = test_data_utility.ConvertToSignedProvisioningMessage(
        test_data_utility.MESSAGE)

    unsigned_request = certificate_provisioning_pb2.ProvisioningRequest()
    unsigned_request.ParseFromString(signed_request.message)

    signed_response = test_data_utility.ConvertToSignedProvisioningMessage(
        raw_response)

    self._VerifyMessageSignature(test_data_utility.SERVICE_PUBLIC_KEY,
                                 signed_response)

    unsigned_response = certificate_provisioning_pb2.ProvisioningResponse()
    unsigned_response.ParseFromString(signed_response.message)

    self._VerifyProvisioningResponse(unsigned_request, unsigned_response)

  def testProcessInvalidMessage(self):
    test_data_utility.AddDrmIntermediateCertificateWithTestData(
        self._engine, 2001, verify_success=True)
    (_, new_session
    ) = test_data_utility.NewProvisioningSessionWithTestData(self._engine)
    (status, _, _) = new_session.ProcessMessage('INVALID_MESSAGE')
    self.assertEqual(pywrapprovisioning_status.INVALID_REQUEST_MESSAGE, status)

  def testNewSessionWithoutIntermediateCert(self):
    (_, new_session) = test_data_utility.NewProvisioningSessionWithTestData(
        self._engine, verify_success=True)
    (status, _, _) = new_session.ProcessMessage(test_data_utility.MESSAGE)
    self.assertEqual(pywrapprovisioning_status.MISSING_DRM_INTERMEDIATE_CERT,
                     status)

  def testNewSessionInvalidDevicePublicKey(self):
    test_data_utility.AddDrmIntermediateCertificateWithTestData(
        self._engine, 2001, verify_success=True)
    (session_status, _) = self._engine.NewProvisioningSession(
        'INVALID_PUBLIC_KEY', test_data_utility.DEVICE_PRIVATE_KEY)
    self.assertEqual(pywrapprovisioning_status.INVALID_DEVICE_PUBLIC_KEY,
                     session_status)

  def testNewSessionInvalidDevicePrivateKey(self):
    test_data_utility.AddDrmIntermediateCertificateWithTestData(
        self._engine, 2001, verify_success=True)
    (session_status, _) = self._engine.NewProvisioningSession(
        test_data_utility.DEVICE_PUBLIC_KEY, 'INVALID_PRIVATE_KEY')
    self.assertEqual(pywrapprovisioning_status.INVALID_DEVICE_PRIVATE_KEY,
                     session_status)

  def _VerifyMessageSignature(self, public_key, signed_response):
    crypto_utility.VerifySignature(public_key, signed_response.signature,
                                   signed_response.message)

  def _VerifyCertSignature(self, public_key, signed_cert):
    crypto_utility.VerifySignature(public_key, signed_cert.signature,
                                   signed_cert.drm_certificate)

  def _VerifyProvisioningResponse(self, request, response):
    self.assertEqual(request.nonce, response.nonce)

    signed_cert = signed_device_certificate_pb2.SignedDrmDeviceCertificate()
    signed_cert.ParseFromString(response.device_certificate)

    self._VerifyCertSignature(test_data_utility.CA_PUBLIC_KEY, signed_cert)

if __name__ == '__main__':
  unittest.main()