# Copyright 2018 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Chrome OS sign_image unittests"""

from chromite.lib import cros_test_lib
from chromite.signing.bin import sign_image
from chromite.signing.image_signing import imagefile


class TestMain(cros_test_lib.RunCommandTestCase):
    """Test sign_image."""

    def setUp(self) -> None:
        self.rc.SetDefaultCmdResult()
        self.infile = "/path/to/image.bin"
        self.outfile = "/path/to/output.bin"
        self.sign_mock = self.PatchObject(imagefile, "SignImage")

    def testSsd(self) -> None:
        """Signs image_type=SSD."""
        self.assertEqual(
            0,
            sign_image.main(
                [
                    "--type",
                    "ssd",
                    "-i",
                    self.infile,
                    "--keyset-dir",
                    "/keydir",
                    "-o",
                    self.outfile,
                ]
            ),
        )
        self.sign_mock.assert_called_once_with(
            "SSD", self.infile, self.outfile, 2, "/keydir", ""
        )

    def testBase(self) -> None:
        """Signs image_type=base."""
        self.assertEqual(
            0,
            sign_image.main(
                [
                    "--type",
                    "base",
                    "-i",
                    self.infile,
                    "--keyset-dir",
                    "/keydir",
                    "-o",
                    self.outfile,
                ]
            ),
        )
        self.sign_mock.assert_called_once_with(
            "SSD", self.infile, self.outfile, 2, "/keydir", ""
        )

    def testUsb(self) -> None:
        """Signs image_type=usb."""
        self.assertEqual(
            0,
            sign_image.main(
                [
                    "--type",
                    "usb",
                    "-i",
                    self.infile,
                    "--keyset-dir",
                    "/keydir",
                    "-o",
                    self.outfile,
                ]
            ),
        )
        self.sign_mock.assert_called_once_with(
            "USB", self.infile, self.outfile, 2, "/keydir", ""
        )

    def testRecovery(self) -> None:
        """Signs image_type=recovery."""
        self.assertEqual(
            0,
            sign_image.main(
                [
                    "--type",
                    "recovery",
                    "-i",
                    self.infile,
                    "--keyset-dir",
                    "/keydir",
                    "-o",
                    self.outfile,
                ]
            ),
        )
        self.sign_mock.assert_called_once_with(
            "recovery", self.infile, self.outfile, 4, "/keydir", "recovery_"
        )

    def testFactory(self) -> None:
        """Signs image_type=factory."""
        self.assertEqual(
            0,
            sign_image.main(
                [
                    "--type",
                    "factory",
                    "-i",
                    self.infile,
                    "--keyset-dir",
                    "/keydir",
                    "-o",
                    self.outfile,
                ]
            ),
        )
        self.sign_mock.assert_called_once_with(
            "factory_install",
            self.infile,
            self.outfile,
            2,
            "/keydir",
            "installer_",
        )

    def testInstall(self) -> None:
        """Signs image_type=install."""
        self.assertEqual(
            0,
            sign_image.main(
                [
                    "--type",
                    "install",
                    "-i",
                    self.infile,
                    "--keyset-dir",
                    "/keydir",
                    "-o",
                    self.outfile,
                ]
            ),
        )
        self.sign_mock.assert_called_once_with(
            "factory_install",
            self.infile,
            self.outfile,
            2,
            "/keydir",
            "installer_",
        )

    def testBadType(self) -> None:
        """Fails with bad image_type."""
        with self.assertRaises(SystemExit):
            sign_image.main(
                [
                    "--type",
                    "badtype",
                    "-i",
                    self.infile,
                    "--keyset-dir",
                    "/keydir",
                    "-o",
                    self.outfile,
                ]
            )
        self.assertEqual(0, self.sign_mock.call_count)

    def testRejectsVersionFile(self) -> None:
        """Rejects obsolete --version-file."""
        with self.assertRaises(SystemExit):
            sign_image.main(
                [
                    "--type",
                    "badtype",
                    "-i",
                    self.infile,
                    "--keyset-dir",
                    "/keydir",
                    "-o",
                    self.outfile,
                    "--version-file",
                    "key.versions",
                ]
            )
        self.assertEqual(0, self.sign_mock.call_count)