#!/usr/bin/env python3.8
# Copyright 2020 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Unit test for verify_zbi_kernel_cmdline.py.

Need to have SCRUTINY and ZBI environmental variables set.

To manually run this test:

  SCRUTINY=~/fuchsia/out/default/host_x64/scrutiny \
  ZBI=~/fuchsia/out/default/host_x64/zbi python3 \
  verify_zbi_kernel_cmdline_test.py
"""
import os
import sys
import subprocess
import tempfile
import unittest
import unittest.mock as mock

import verify_zbi_kernel_cmdline


def verify_kernel_cmdline(golden, actual):
  with tempfile.TemporaryDirectory() as test_folder:
    golden_file = os.path.join(test_folder, 'golden')
    stamp_file = os.path.join(test_folder, 'stamp')
    fuchsia_folder = os.path.join(test_folder, 'fuchsia')
    test_zbi = os.path.join(test_folder, 'test.zbi')
    cmdline_file = os.path.join(test_folder, 'cmdline')
    scrutiny = os.environ['SCRUTINY']
    with open(golden_file, 'w+') as f:
      f.write(golden)
    with open(cmdline_file, 'wb+') as f:
      f.write(actual)

    # Use ZBI to create a test.zbi that only contains cmdline.
    subprocess.check_call(
        [os.environ['ZBI'], '-o', test_zbi, '-T', 'CMDLINE', cmdline_file])
    os.mkdir(fuchsia_folder)

    args = [
        '--zbi-file', test_zbi, '--scrutiny', scrutiny, '--fuchsia-dir',
        fuchsia_folder, '--kernel-cmdline-golden-file', golden_file, '--stamp',
        stamp_file
    ]
    # Verify the cmdline in the generated ZBI.
    return verify_zbi_kernel_cmdline.main(args)


class RunVerifyZbiKernelCmdlineTest(unittest.TestCase):

  def test_verify_kernel_cmdline_sucess_normal_case(self):
    self.assertEqual(
        0,
        verify_kernel_cmdline('key1=v1\nkey2=v2\nkey3=v3',
                              b'key1=v1 key2=v2 key3=v3'))

  def test_verify_kernel_cmdline_success_order_diff(self):
    self.assertEqual(
        0,
        verify_kernel_cmdline('key1=v1\nkey2=v2\nkey3=v3',
                              b'key2=v2 key1=v1 key3=v3'))

  def test_verify_kernel_cmdline_success_no_value_option(self):
    self.assertEqual(
        0, verify_kernel_cmdline('option1\noption2', b'option1 option2'))

  def test_verify_kernel_cmdline_fail_golden_empty(self):
    self.assertEqual(-1, verify_kernel_cmdline('', b'key2=v2 key1=v1 key3=v3'))

  def test_verify_kernel_cmdline_fail_missing_key2(self):
    self.assertEqual(-1, verify_kernel_cmdline('key1=v1\nkey2=v2', b'key1=v1'))

  def test_verify_kernel_cmdline_fail_key1_mismatch(self):
    self.assertEqual(
        -1, verify_kernel_cmdline('key1=v1\nkey2=v2', b'key1=v2 key2=v2'))

  def test_verify_kernel_cmdline_fail_key2_mismatch(self):
    self.assertEqual(
        -1, verify_kernel_cmdline('key1=v1\nkey2=v2', b'key1=v1 key2=v1'))

  def test_verify_kernel_cmdline_fail_additional_key3(self):
    self.assertEqual(
        -1, verify_kernel_cmdline('key1=v1\nkey2=v2',
                                  b'key1=v1 key2=v2 key3=v3'))

  def test_verify_kernel_cmdline_fail_invalid_format(self):
    self.assertEqual(
        -1, verify_kernel_cmdline('key1=v1\nkey2=v2', b'invalid=format=1'))

  def test_verify_kernel_cmdline_fail_option1_missing(self):
    self.assertEqual(-1, verify_kernel_cmdline('option1\noption2', b'option2'))

  def test_verify_kernel_cmdline_fail_additional_option3(self):
    self.assertEqual(
        -1, verify_kernel_cmdline('option1\noption2',
                                  b'option1 option2 option3'))

  def test_verify_kernel_cmdline_zbi_not_found(self):
    with tempfile.TemporaryDirectory() as test_folder:
      golden_file = os.path.join(test_folder, 'golden')
      stamp_file = os.path.join(test_folder, 'stamp')
      fuchsia_folder = os.path.join(test_folder, 'fuchsia')
      test_zbi = os.path.join(test_folder, 'test.zbi')
      scrutiny = os.environ['SCRUTINY']
      with open(golden_file, 'w+') as f:
        f.write('option1')

      # Do not create test_zbi

      os.mkdir(fuchsia_folder)

      args = [
          '--zbi-file', test_zbi, '--scrutiny', scrutiny, '--fuchsia-dir',
          fuchsia_folder, '--kernel-cmdline-golden-file', golden_file,
          '--stamp', stamp_file
      ]
      self.assertEqual(-1, verify_zbi_kernel_cmdline.main(args))

  def test_verify_kernel_cmdline_success_no_cmdline_found(self):
    with tempfile.TemporaryDirectory() as test_folder:
      golden_file = os.path.join(test_folder, 'golden')
      stamp_file = os.path.join(test_folder, 'stamp')
      fuchsia_folder = os.path.join(test_folder, 'fuchsia')
      test_zbi = os.path.join(test_folder, 'test.zbi')
      scrutiny = os.environ['SCRUTINY']

      # Create an empty golden file
      with open(golden_file, 'w+') as f:
        f.write('')

      # Use ZBI to create a test.zbi with no cmdline.
      subprocess.check_call([os.environ['ZBI'], '-o', test_zbi])

      os.mkdir(fuchsia_folder)

      args = [
          '--zbi-file', test_zbi, '--scrutiny', scrutiny, '--fuchsia-dir',
          fuchsia_folder, '--kernel-cmdline-golden-file', golden_file,
          '--stamp', stamp_file
      ]
      self.assertEqual(0, verify_zbi_kernel_cmdline.main(args))

  def test_verify_kernel_cmdline_fail_golden_empty_cmdline_found(self):
    self.assertEqual(-1, verify_kernel_cmdline('', b'option2'))

  def test_verify_kernel_cmdline_fail_golden_not_empty_cmdline_not_found(self):
    with tempfile.TemporaryDirectory() as test_folder:
      golden_file = os.path.join(test_folder, 'golden')
      stamp_file = os.path.join(test_folder, 'stamp')
      fuchsia_folder = os.path.join(test_folder, 'fuchsia')
      test_zbi = os.path.join(test_folder, 'test.zbi')
      scrutiny = os.environ['SCRUTINY']

      # Create an empty golden file
      with open(golden_file, 'w+') as f:
        f.write('option1')

      # Use ZBI to create a test.zbi with no cmdline.
      subprocess.check_call([os.environ['ZBI'], '-o', test_zbi])

      os.mkdir(fuchsia_folder)

      args = [
          '--zbi-file', test_zbi, '--scrutiny', scrutiny, '--fuchsia-dir',
          fuchsia_folder, '--kernel-cmdline-golden-file', golden_file,
          '--stamp', stamp_file
      ]
      self.assertEqual(-1, verify_zbi_kernel_cmdline.main(args))


if __name__ == '__main__':
  if 'SCRUTINY' not in os.environ or 'ZBI' not in os.environ:
    print('Please set SCRUTINY and ZBI environmental path')
    sys.exit(-1)
  unittest.main()