aftl_integration_test.py

#!/usr/bin/env python

# Copyright 2019, The Android Open Source Project
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation
# files (the "Software"), to deal in the Software without
# restriction, including without limitation the rights to use, copy,
# modify, merge, publish, distribute, sublicense, and/or sell copies
# of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
# BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
"""Integration tests for the avbtool with an actual AFTL.

The test cases directly interact with a transparency log. However,
before using this script the following environment variables
need to be set:

  AFTL_HOST: host:port of the transparency log to test with.
  AFTL_PUBKEY: Transparency log public key in PEM format.
  AFTL_VBMETA_IMAGE: VBMeta image that should be used for submission to AFTL.
  AFTL_MANUFACTURER_KEY: Manufacturer signing key used to sign submissions
      to the transparency log in PEM format.
"""

import io
import os
import unittest

import aftltool


class AftlIntegrationTest(unittest.TestCase):
  """Test suite for integration testing aftltool with an actual AFTL."""

  def setUp(self):
    """Sets up the test bed for the integration tests."""
    super(AftlIntegrationTest, self).setUp()
    self.aftltool = aftltool.Aftl()
    self.output_filename = 'vbmeta_icp.img'

    self.aftl_host = os.environ.get('AFTL_HOST')
    self.aftl_pubkey = os.environ.get('AFTL_PUBKEY')
    self.vbmeta_image = os.environ.get('AFTL_VBMETA_IMAGE')
    self.manufacturer_key = os.environ.get('AFTL_MANUFACTURER_KEY')

    if (not self.aftl_host or not self.aftl_pubkey or not self.vbmeta_image
        or not self.manufacturer_key):
      self.fail('Environment variables not correctly set up. See description of'
                ' this test case for details')

    self.make_icp_default_params = {
        'vbmeta_image_path': self.vbmeta_image,
        'output': None,
        'signing_helper': None,
        'signing_helper_with_files': None,
        'version_incremental': '1',
        'transparency_log_servers': [self.aftl_host],
        'transparency_log_pub_keys': [self.aftl_pubkey],
        'manufacturer_key': self.manufacturer_key,
        'padding_size': 0,
        'timeout': None
    }

    self.info_icp_default_params = {
        'vbmeta_image_path': self.output_filename,
        'output': io.BytesIO()
    }

    self.verify_icp_default_params = {
        'vbmeta_image_path': self.output_filename,
        'transparency_log_pub_keys': [self.aftl_pubkey],
        'output': io.BytesIO()
    }

    self.load_test_aftl_default_params = {
        'vbmeta_image_path': self.vbmeta_image,
        'output': io.BytesIO(),
        'transparency_log_server': self.aftl_host,
        'transparency_log_pub_key': self.aftl_pubkey,
        'manufacturer_key': self.manufacturer_key,
        'process_count': 1,
        'submission_count': 1,
        'stats_filename': None,
        'preserve_icp_images': False,
        'timeout': None
    }
    self.load_test_stats_file_p1_s1 = 'load_test_p1_s1.csv'
    self.load_test_stats_file_p2_p2 = 'load_test_p2_s2.csv'

    self.files_to_cleanup = [
        self.output_filename,
        self.load_test_stats_file_p1_s1,
        self.load_test_stats_file_p2_p2
    ]

  def tearDown(self):
    """Tears down the test bed for the unit tests."""
    for filename in self.files_to_cleanup:
      try:
        os.remove(filename)
      except OSError:
        pass
    super(AftlIntegrationTest, self).tearDown()

  def test_make_and_verify_icp_with_1_log(self):
    """Tests integration of aftltool with one AFTL."""
    # Make a VBmeta image with ICP.
    with open(self.output_filename, 'wb') as output_file:
      self.make_icp_default_params['output'] = output_file
      result = self.aftltool.make_icp_from_vbmeta(
          **self.make_icp_default_params)
      self.assertTrue(result)

    # Checks that there is 1 ICP.
    aftl_descriptor = self.aftltool.get_aftl_descriptor(self.output_filename)
    self.assertEqual(aftl_descriptor.icp_header.icp_count, 1)

    # Verifies the generated image.
    result = self.aftltool.verify_image_icp(**self.verify_icp_default_params)
    self.assertTrue(result)

    # Prints the image details.
    result = self.aftltool.info_image_icp(**self.info_icp_default_params)
    self.assertTrue(result)

  def test_make_and_verify_icp_with_2_logs(self):
    # Reconfigures default parameters with two transparency logs.
    self.make_icp_default_params['transparency_log_servers'] = [
        self.aftl_host, self.aftl_host]
    self.make_icp_default_params['transparency_log_pub_keys'] = [
        self.aftl_pubkey, self.aftl_pubkey]

    # Make a VBmeta image with ICP.
    with open(self.output_filename, 'wb') as output_file:
      self.make_icp_default_params['output'] = output_file
      result = self.aftltool.make_icp_from_vbmeta(
          **self.make_icp_default_params)
      self.assertTrue(result)

    # Checks that there are 2 ICPs.
    aftl_descriptor = self.aftltool.get_aftl_descriptor(self.output_filename)
    self.assertEqual(aftl_descriptor.icp_header.icp_count, 2)

    # Verifies the generated image.
    result = self.aftltool.verify_image_icp(**self.verify_icp_default_params)
    self.assertTrue(result)

    # Prints the image details.
    result = self.aftltool.info_image_icp(**self.info_icp_default_params)
    self.assertTrue(result)

  def test_make_icp_with_invalid_grpc_service(self):
    """Tests make_icp_from_vbmeta command with a host that does not support GRPC."""
    self.make_icp_default_params[
        'transparency_log_servers'] = ['www.google.com:80']
    with open(self.output_filename, 'wb') as output_file:
      self.make_icp_default_params['output'] = output_file
      result = self.aftltool.make_icp_from_vbmeta(
          **self.make_icp_default_params)
      self.assertFalse(result)

  def test_make_icp_grpc_timeout(self):
    """Tests make_icp_from_vbmeta command when running into GRPC timeout."""
    # The timeout is set to 1 second which is way below the minimum processing
    # time of the transparency log per load test results in b/139407814#2 where
    # it was 3.43 seconds.
    self.make_icp_default_params['timeout'] = 1
    with open(self.output_filename, 'wb') as output_file:
      self.make_icp_default_params['output'] = output_file
      result = self.aftltool.make_icp_from_vbmeta(
          **self.make_icp_default_params)
      self.assertFalse(result)

  def test_load_test_single_process_single_submission(self):
    """Tests load_test_aftl command with 1 process which does 1 submission."""
    result = self.aftltool.load_test_aftl(**self.load_test_aftl_default_params)
    self.assertTrue(result)

    output = self.load_test_aftl_default_params['output'].getvalue()
    self.assertRegexpMatches(output, 'Succeeded:.+?1\n')
    self.assertRegexpMatches(output, 'Failed:.+?0\n')

    self.assertTrue(os.path.exists(self.load_test_stats_file_p1_s1))

  def test_load_test_multi_procces_multi_submission(self):
    """Tests load_test_aftl command with 2 processes and 2 submissions each."""
    self.load_test_aftl_default_params['process_count'] = 2
    self.load_test_aftl_default_params['submission_count'] = 2
    result = self.aftltool.load_test_aftl(**self.load_test_aftl_default_params)
    self.assertTrue(result)

    output = self.load_test_aftl_default_params['output'].getvalue()
    self.assertRegexpMatches(output, 'Succeeded:.+?4\n')
    self.assertRegexpMatches(output, 'Failed:.+?0\n')

    self.assertTrue(os.path.exists(self.load_test_stats_file_p2_p2))

  def test_load_test_invalid_grpc_service(self):
    """Tests load_test_aftl command with a host that does not support GRPC."""
    self.load_test_aftl_default_params[
        'transparency_log_server'] = 'www.google.com:80'
    result = self.aftltool.load_test_aftl(**self.load_test_aftl_default_params)
    self.assertFalse(result)

    output = self.load_test_aftl_default_params['output'].getvalue()
    self.assertRegexpMatches(output, 'Succeeded:.+?0\n')
    self.assertRegexpMatches(output, 'Failed:.+?1\n')

  def test_load_test_grpc_timeout(self):
    """Tests load_test_aftl command when running into timeout."""
    self.load_test_aftl_default_params['timeout'] = 1
    result = self.aftltool.load_test_aftl(**self.load_test_aftl_default_params)
    self.assertFalse(result)

    output = self.load_test_aftl_default_params['output'].getvalue()
    self.assertRegexpMatches(output, 'Succeeded:.+?0\n')
    self.assertRegexpMatches(output, 'Failed:.+?1\n')


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