#!/usr/bin/env python3
#
# Copyright 2019-2022 Canonical Ltd.
# Authors:
# - dann frazier <dann.frazier@canonical.com>
#
# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License version 3, as published
# by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranties of MERCHANTABILITY,
# SATISFACTORY QUALITY, or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License along with
# this program.  If not, see <http://www.gnu.org/licenses/>.
#

import os
import pexpect
import subprocess
import sys
import time
import unittest

from UEFI.Filesystems import GrubShellBootableIsoImage
from UEFI.SignedBinary import SignedBinary
from UEFI.Qemu import QemuEfiMachine, QemuEfiVariant, QemuEfiFlashSize
from UEFI import Qemu

DPKG_ARCH = subprocess.check_output(
    ['dpkg', '--print-architecture']
).decode().rstrip()

EfiArchToGrubArch = {
    'X64': "x86_64",
    'AA64': "arm64",
}

TEST_TIMEOUT = 120


def get_local_grub_path(efi_arch, signed=False):
    grub_subdir = "%s-efi" % EfiArchToGrubArch[efi_arch.upper()]
    ext = "efi"
    if signed:
        grub_subdir = f"{grub_subdir}-signed"
        ext = f"{ext}.signed"

    grub_path = os.path.join(
        os.path.sep, 'usr', 'lib', 'grub',
        '%s' % (grub_subdir),
        "" if signed else "monolithic",
        'grub%s.%s' % (efi_arch.lower(), ext)
    )
    return grub_path


def get_local_shim_path(efi_arch, signed=False):
    ext = 'efi'
    if signed:
        ext = f"{ext}.signed"
    shim_path = os.path.join(
        os.path.sep, 'usr', 'lib', 'shim',
        'shim%s.%s' % (efi_arch.lower(), ext)
    )
    return shim_path


class BootToShellTest(unittest.TestCase):
    debug = True

    def setUp(self):
        self.startTime = time.time()

    def tearDown(self):
        t = time.time() - self.startTime
        sys.stdout.write("%s runtime: %.3fs\n" % (self.id(), t))

    def run_cmd_check_shell(self, cmd):
        child = pexpect.spawn(' '.join(cmd), encoding='UTF-8')

        if self.debug:
            child.logfile = sys.stdout
        try:
            while True:
                i = child.expect(
                    [
                        'Press .* or any other key to continue',
                        'Shell> '
                    ],
                    timeout=TEST_TIMEOUT,
                )
                if i == 0:
                    child.sendline('\x1b')
                    continue
                if i == 1:
                    child.sendline('reset -s\r')
                    continue
        except pexpect.EOF:
            child.close()
            if child.exitstatus != 0:
                self.fail("ERROR: exit code %d\n" % (child.exitstatus))
        except pexpect.TIMEOUT as err:
            self.fail("%s\n" % (err))

    def run_cmd_check_secure_boot(self, cmd, efiarch, should_verify):
        child = pexpect.spawn(' '.join(cmd), encoding='UTF-8')

        if self.debug:
            child.logfile = sys.stdout
        try:
            while True:
                i = child.expect(
                    [
                        'grub> ',
                        'BdsDxe: failed to load.*: Access Denied',
                    ],
                    timeout=TEST_TIMEOUT,
                )
                if i == 0:
                    child.sendline('halt\r')
                    verified = True
                    continue
                if i == 1:
                    child.close()
                    verified = False
                    break
        except pexpect.EOF:
            child.close()
            if child.exitstatus != 0:
                self.fail("ERROR: exit code %d\n" % (child.exitstatus))
        except pexpect.TIMEOUT as err:
            self.fail("%s\n" % (err))
        self.assertEqual(should_verify, verified)

    def test_aavmf(self):
        q = Qemu.QemuCommand(QemuEfiMachine.AAVMF)
        self.run_cmd_check_shell(q.command)

    @unittest.skipUnless(DPKG_ARCH == 'arm64', "Requires grub-efi-arm64")
    def test_aavmf_ms_secure_boot_signed(self):
        q = Qemu.QemuCommand(
            QemuEfiMachine.AAVMF,
            variant=QemuEfiVariant.MS,
        )
        grub = get_local_grub_path('AA64', signed=True)
        shim = get_local_shim_path('AA64', signed=True)
        iso = GrubShellBootableIsoImage('AA64', shim, grub)
        q.add_disk(iso.path)
        self.run_cmd_check_secure_boot(q.command, 'aa64', True)

    @unittest.skipUnless(DPKG_ARCH == 'arm64', "Requires grub-efi-arm64")
    def test_aavmf_ms_secure_boot_unsigned(self):
        q = Qemu.QemuCommand(
            QemuEfiMachine.AAVMF,
            variant=QemuEfiVariant.MS,
        )
        grub = get_local_grub_path('AA64', signed=False)
        shim = get_local_shim_path('AA64', signed=False)
        iso = GrubShellBootableIsoImage('AA64', shim, grub)
        q.add_disk(iso.path)
        self.run_cmd_check_secure_boot(q.command, 'aa64', False)

    @unittest.skipUnless(DPKG_ARCH == 'arm64', "arm64-only")
    def test_aavmf_snakeoil_secure_boot_signed(self):
        q = Qemu.QemuCommand(
            QemuEfiMachine.AAVMF,
            variant=QemuEfiVariant.SNAKEOIL,
        )
        shim = SignedBinary(
            get_local_shim_path('AA64', signed=False),
            "/usr/share/qemu-efi-aarch64/PkKek-1-snakeoil.key",
            "/usr/share/qemu-efi-aarch64/PkKek-1-snakeoil.pem",
            "snakeoil",
        )
        grub = SignedBinary(
            get_local_grub_path('AA64', signed=False),
            "/usr/share/qemu-efi-aarch64/PkKek-1-snakeoil.key",
            "/usr/share/qemu-efi-aarch64/PkKek-1-snakeoil.pem",
            "snakeoil",
        )
        iso = GrubShellBootableIsoImage('AA64', shim.path, grub.path)
        q.add_disk(iso.path)
        self.run_cmd_check_secure_boot(q.command, 'aa64', True)

    @unittest.skipUnless(DPKG_ARCH == 'arm64', "arm64-only")
    def test_aavmf_snakeoil_secure_boot_unsigned(self):
        q = Qemu.QemuCommand(
            QemuEfiMachine.AAVMF,
            variant=QemuEfiVariant.SNAKEOIL,
            flash_size=QemuEfiFlashSize.DEFAULT,
        )
        grub = get_local_grub_path('AA64', signed=False)
        shim = get_local_shim_path('AA64', signed=False)
        iso = GrubShellBootableIsoImage('AA64', shim, grub)
        q.add_disk(iso.path)
        self.run_cmd_check_secure_boot(q.command, 'aa64', False)

    def test_aavmf32(self):
        q = Qemu.QemuCommand(QemuEfiMachine.AAVMF32)
        self.run_cmd_check_shell(q.command)

    def test_loongarch64(self):
        q = Qemu.QemuCommand(QemuEfiMachine.LOONGARCH64)
        self.run_cmd_check_shell(q.command)

    def test_ovmf_4m_pc(self):
        q = Qemu.QemuCommand(
            QemuEfiMachine.OVMF_PC,
            flash_size=QemuEfiFlashSize.SIZE_4MB,
        )
        self.run_cmd_check_shell(q.command)

    def test_ovmf_4m_q35(self):
        q = Qemu.QemuCommand(
            QemuEfiMachine.OVMF_Q35,
            flash_size=QemuEfiFlashSize.SIZE_4MB,
        )
        self.run_cmd_check_shell(q.command)

    @unittest.skipUnless(DPKG_ARCH == 'amd64', "amd64-only")
    def test_ovmf_4m_ms_secure_boot_signed(self):
        q = Qemu.QemuCommand(
            QemuEfiMachine.OVMF_Q35,
            variant=QemuEfiVariant.MS,
            flash_size=QemuEfiFlashSize.SIZE_4MB,
        )
        grub = get_local_grub_path('X64', signed=True)
        shim = get_local_shim_path('X64', signed=True)
        iso = GrubShellBootableIsoImage('X64', shim, grub)
        q.add_disk(iso.path)
        self.run_cmd_check_secure_boot(q.command, 'x64', True)

    @unittest.skipUnless(DPKG_ARCH == 'amd64', "amd64-only")
    def test_ovmf_4m_ms_secure_boot_unsigned(self):
        q = Qemu.QemuCommand(
            QemuEfiMachine.OVMF_Q35,
            variant=QemuEfiVariant.MS,
            flash_size=QemuEfiFlashSize.SIZE_4MB,
        )
        grub = get_local_grub_path('X64', signed=False)
        shim = get_local_shim_path('X64', signed=False)
        iso = GrubShellBootableIsoImage('X64', shim, grub)
        q.add_disk(iso.path)
        self.run_cmd_check_secure_boot(q.command, 'x64', False)

    @unittest.skipUnless(DPKG_ARCH == 'amd64', "amd64-only")
    def test_ovmf_snakeoil_secure_boot_signed(self):
        q = Qemu.QemuCommand(
            QemuEfiMachine.OVMF_Q35,
            variant=QemuEfiVariant.SNAKEOIL,
        )
        shim = SignedBinary(
            get_local_shim_path('X64', signed=False),
            "/usr/share/ovmf/PkKek-1-snakeoil.key",
            "/usr/share/ovmf/PkKek-1-snakeoil.pem",
            "snakeoil",
        )
        grub = SignedBinary(
            get_local_grub_path('X64', signed=False),
            "/usr/share/ovmf/PkKek-1-snakeoil.key",
            "/usr/share/ovmf/PkKek-1-snakeoil.pem",
            "snakeoil",
        )
        iso = GrubShellBootableIsoImage('X64', shim.path, grub.path)
        q.add_disk(iso.path)
        self.run_cmd_check_secure_boot(q.command, 'x64', True)

    @unittest.skipUnless(DPKG_ARCH == 'amd64', "amd64-only")
    def test_ovmf_snakeoil_secure_boot_unsigned(self):
        q = Qemu.QemuCommand(
            QemuEfiMachine.OVMF_Q35,
            variant=QemuEfiVariant.SNAKEOIL,
            flash_size=QemuEfiFlashSize.DEFAULT,
        )
        grub = get_local_grub_path('X64', signed=False)
        shim = get_local_shim_path('X64', signed=False)
        iso = GrubShellBootableIsoImage('X64', shim, grub)
        q.add_disk(iso.path)
        self.run_cmd_check_secure_boot(q.command, 'x64', False)

    def test_ovmf32_4m_pc(self):
        q = Qemu.QemuCommand(
            QemuEfiMachine.OVMF32_PC,
            flash_size=QemuEfiFlashSize.SIZE_4MB,
        )
        self.run_cmd_check_shell(q.command)

    def test_ovmf32_4m_q35(self):
        q = Qemu.QemuCommand(
            QemuEfiMachine.OVMF32_Q35,
            flash_size=QemuEfiFlashSize.SIZE_4MB,
        )
        self.run_cmd_check_shell(q.command)

    def test_riscv64(self):
        q = Qemu.QemuCommand(QemuEfiMachine.RISCV64)
        self.run_cmd_check_shell(q.command)

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