#!/usr/bin/env python3

# Copyright (C) 2026 Alexander Larsson <alexl@redhat.com>
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA

import os
import shutil
import struct
import subprocess
import tempfile
import unittest


UKIBOOTIMG_BINARY = os.environ.get("UKIBOOTIMG_PATH", "./ukibootimg")
UKIBOOTCTL_BINARY = os.environ.get("UKIBOOTCTL_PATH", "./ukibootctl")
HAS_SBSIGN = shutil.which('sbsign') is not None and shutil.which('openssl') is not None


def create_pe(sections):
    """Create a minimal PE32+ file with the given sections.

    sections: list of (name, content_bytes) tuples
    """
    nt_offset = 64
    opt_header_size = 240
    num_sections = len(sections)
    headers_size = nt_offset + 4 + 20 + opt_header_size + num_sections * 40
    aligned_headers = (headers_size + 511) & ~511

    # DOS header
    dos = bytearray(64)
    struct.pack_into('<H', dos, 0, 0x5A4D)
    struct.pack_into('<i', dos, 60, nt_offset)

    # NT signature
    sig = struct.pack('<I', 0x00004550)

    # File header
    file_hdr = struct.pack('<HHIIIHH',
        0x8664,          # Machine: AMD64
        num_sections,
        0, 0, 0,         # TimeDateStamp, PointerToSymbolTable, NumberOfSymbols
        opt_header_size,
        0x22,            # Characteristics
    )

    # Optional header (PE32+)
    opt = bytearray(opt_header_size)
    struct.pack_into('<H', opt, 0, 0x20b)       # Magic: PE32+
    struct.pack_into('<I', opt, 56, aligned_headers)  # SizeOfImage (approx)
    struct.pack_into('<I', opt, 60, aligned_headers)  # SizeOfHeaders
    struct.pack_into('<I', opt, 108, 16)         # NumberOfRvaAndSizes

    # Build section headers and data
    current_offset = aligned_headers
    sec_hdrs = bytearray()
    sec_data = bytearray()

    for name, content in sections:
        name_bytes = name.encode('ascii')[:8].ljust(8, b'\x00')
        raw_size = (len(content) + 511) & ~511

        shdr = bytearray(40)
        shdr[0:8] = name_bytes
        struct.pack_into('<I', shdr, 8, len(content))    # VirtualSize
        struct.pack_into('<I', shdr, 12, current_offset)  # VirtualAddress
        struct.pack_into('<I', shdr, 16, raw_size)        # SizeOfRawData
        struct.pack_into('<I', shdr, 20, current_offset)  # PointerToRawData

        sec_hdrs += shdr
        sec_data += content.ljust(raw_size, b'\x00')
        current_offset += raw_size

    pe = dos + sig + file_hdr + bytes(opt) + bytes(sec_hdrs)
    pe = pe.ljust(aligned_headers, b'\x00')
    pe += bytes(sec_data)
    return bytes(pe)


SLOT_A_SECTIONS = [
    ('.osrel', b'NAME="TestOS"\nVERSION="1.0"\n'),
    ('.uname', b'6.0.0-slota'),
    ('.cmdline', b'root=/dev/sda1'),
    ('.linux', b'KERNEL_A_DATA'),
    ('.initrd', b'INITRD_A_DATA'),
]

SLOT_B_SECTIONS = [
    ('.osrel', b'NAME="TestOS"\nVERSION="2.0"\n'),
    ('.uname', b'6.0.0-slotb'),
    ('.cmdline', b'root=/dev/sda2'),
    ('.linux', b'KERNEL_B_DATA'),
    ('.initrd', b'INITRD_B_DATA'),
]


class TestUkiBootImg(unittest.TestCase):

    def setUp(self):
        self.tmpdir = tempfile.mkdtemp()
        self.pe_a_path = os.path.join(self.tmpdir, 'slot_a.efi')
        self.pe_b_path = os.path.join(self.tmpdir, 'slot_b.efi')
        self.bootctl_path = os.path.join(self.tmpdir, 'bootctl.bin')
        self.extract_dir = os.path.join(self.tmpdir, 'extract')
        os.makedirs(self.extract_dir)

        with open(self.pe_a_path, 'wb') as f:
            f.write(create_pe(SLOT_A_SECTIONS))
        with open(self.pe_b_path, 'wb') as f:
            f.write(create_pe(SLOT_B_SECTIONS))

        # Initialize bootctl (defaults to slot 0/a active)
        open(self.bootctl_path, 'wb').close()
        subprocess.run([
            UKIBOOTCTL_BINARY,
            '--partition=' + self.bootctl_path,
            'init'
        ], check=True, capture_output=True)

    def tearDown(self):
        shutil.rmtree(self.tmpdir)

    def run_ukibootimg(self, args, expected_exit_code=0):
        cmd = [UKIBOOTIMG_BINARY,
               '--boot-a', self.pe_a_path,
               '--boot-b', self.pe_b_path,
               '--bootctl', self.bootctl_path] + args
        result = subprocess.run(cmd, capture_output=True, text=True, check=False)
        self.assertEqual(result.returncode, expected_exit_code,
                         f"Command failed (exit {result.returncode}):\n"
                         f"STDOUT: {result.stdout}\nSTDERR: {result.stderr}")
        return result.stdout, result.stderr

    def switch_to_slot_b(self):
        ctl = UKIBOOTCTL_BINARY
        part = '--partition=' + self.bootctl_path
        subprocess.run([ctl, part, 'prepare-switch', '1'], check=True, capture_output=True)
        subprocess.run([ctl, part, 'finalize-switch', '1'], check=True, capture_output=True)

    # --- list command ---

    def test_list(self):
        stdout, _ = self.run_ukibootimg(['list'])
        for name, content in SLOT_A_SECTIONS:
            self.assertIn(name, stdout)
            self.assertIn(str(len(content)), stdout)

    # --- status command ---

    def test_status_slot_a(self):
        stdout, _ = self.run_ukibootimg(['status'])
        self.assertIn('Active slot: a', stdout)
        self.assertIn('TestOS 1.0', stdout)
        self.assertIn('6.0.0-slota', stdout)
        self.assertIn('root=/dev/sda1 androidboot.slot_suffix=_a', stdout)

    def test_status_slot_b(self):
        self.switch_to_slot_b()
        stdout, _ = self.run_ukibootimg(['status'])
        self.assertIn('Active slot: b', stdout)
        self.assertIn('TestOS 2.0', stdout)
        self.assertIn('6.0.0-slotb', stdout)
        self.assertIn('root=/dev/sda2 androidboot.slot_suffix=_b', stdout)

    # --- extract command ---

    def test_extract_all(self):
        self.run_ukibootimg(['--output-dir', self.extract_dir, 'extract'])
        for name, content in SLOT_A_SECTIONS:
            fname = name.lstrip('.')
            path = os.path.join(self.extract_dir, fname)
            self.assertTrue(os.path.exists(path), f'{fname} not extracted')
            with open(path, 'rb') as f:
                data = f.read()
            if name == '.cmdline':
                self.assertEqual(data, content + b' androidboot.slot_suffix=_a')
            else:
                self.assertEqual(data, content)

    def test_extract_selective(self):
        self.run_ukibootimg(['--output-dir', self.extract_dir, 'extract', '.uname'])
        self.assertTrue(os.path.exists(os.path.join(self.extract_dir, 'uname')))
        self.assertFalse(os.path.exists(os.path.join(self.extract_dir, 'linux')))

    def test_extract_without_dot(self):
        self.run_ukibootimg(['--output-dir', self.extract_dir, 'extract', 'cmdline'])
        path = os.path.join(self.extract_dir, 'cmdline')
        self.assertTrue(os.path.exists(path))

    def test_extract_slot_b(self):
        self.switch_to_slot_b()
        self.run_ukibootimg(['--output-dir', self.extract_dir, 'extract'])
        with open(os.path.join(self.extract_dir, 'uname'), 'rb') as f:
            self.assertEqual(f.read(), b'6.0.0-slotb')
        with open(os.path.join(self.extract_dir, 'cmdline'), 'rb') as f:
            self.assertEqual(f.read(), b'root=/dev/sda2 androidboot.slot_suffix=_b')

    def test_slot_override(self):
        # bootctl says slot a, but --slot b overrides
        cmd = [UKIBOOTIMG_BINARY,
               '--boot-a', self.pe_a_path,
               '--boot-b', self.pe_b_path,
               '--bootctl', self.bootctl_path,
               '--slot', 'b',
               '--output-dir', self.extract_dir,
               'extract']
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        with open(os.path.join(self.extract_dir, 'uname'), 'rb') as f:
            self.assertEqual(f.read(), b'6.0.0-slotb')

    # --- verify command ---

    def test_verify_no_cert(self):
        _, stderr = self.run_ukibootimg(['verify'], expected_exit_code=1)
        self.assertIn('--cert', stderr)

    def test_verify_unsigned(self):
        # Create a dummy cert file (content doesn't matter, PE has no signature)
        cert_path = os.path.join(self.tmpdir, 'dummy.crt')
        with open(cert_path, 'w') as f:
            f.write('not a real cert')
        _, stderr = self.run_ukibootimg([
            '--cert', cert_path, 'verify'
        ], expected_exit_code=1)
        self.assertIn('no embedded signature', stderr)

    # --- error cases ---

    def test_missing_input(self):
        result = subprocess.run([UKIBOOTIMG_BINARY, 'list'],
                                capture_output=True, text=True, check=False)
        self.assertNotEqual(result.returncode, 0)

    def test_unknown_command(self):
        result = subprocess.run([
            UKIBOOTIMG_BINARY,
            '--boot-a', self.pe_a_path,
            'nonexistent'
        ], capture_output=True, text=True, check=False)
        self.assertNotEqual(result.returncode, 0)
        self.assertIn('Unknown command', result.stderr)


@unittest.skipUnless(HAS_SBSIGN, "sbsign and openssl not available")
class TestUkiBootImgVerify(unittest.TestCase):

    def setUp(self):
        self.tmpdir = tempfile.mkdtemp()
        self.extract_dir = os.path.join(self.tmpdir, 'extract')
        os.makedirs(self.extract_dir)

        self.key_path = os.path.join(self.tmpdir, 'test.key')
        self.cert_path = os.path.join(self.tmpdir, 'test.crt')
        self.other_key_path = os.path.join(self.tmpdir, 'other.key')
        self.other_cert_path = os.path.join(self.tmpdir, 'other.crt')

        # Generate two key pairs: one for signing, one for wrong-cert test
        for key, cert, cn in [
            (self.key_path, self.cert_path, 'Test Signer'),
            (self.other_key_path, self.other_cert_path, 'Other Signer'),
        ]:
            subprocess.run([
                'openssl', 'req', '-new', '-x509', '-newkey', 'rsa:2048',
                '-keyout', key, '-out', cert,
                '-days', '1', '-nodes', '-subj', f'/CN={cn}/'
            ], check=True, capture_output=True)

        # Create and sign PE files
        self.pe_a_path = os.path.join(self.tmpdir, 'slot_a.efi')
        self.pe_b_path = os.path.join(self.tmpdir, 'slot_b.efi')
        self.bootctl_path = os.path.join(self.tmpdir, 'bootctl.bin')

        for pe_path, sections in [
            (self.pe_a_path, SLOT_A_SECTIONS),
            (self.pe_b_path, SLOT_B_SECTIONS),
        ]:
            unsigned = pe_path + '.unsigned'
            with open(unsigned, 'wb') as f:
                f.write(create_pe(sections))
            subprocess.run([
                'sbsign', '--key', self.key_path, '--cert', self.cert_path,
                '--output', pe_path, unsigned
            ], check=True, capture_output=True)

        open(self.bootctl_path, 'wb').close()
        subprocess.run([
            UKIBOOTCTL_BINARY,
            '--partition=' + self.bootctl_path,
            'init'
        ], check=True, capture_output=True)

    def tearDown(self):
        shutil.rmtree(self.tmpdir)

    def run_ukibootimg(self, args, expected_exit_code=0):
        cmd = [UKIBOOTIMG_BINARY,
               '--boot-a', self.pe_a_path,
               '--boot-b', self.pe_b_path,
               '--bootctl', self.bootctl_path] + args
        result = subprocess.run(cmd, capture_output=True, text=True, check=False)
        self.assertEqual(result.returncode, expected_exit_code,
                         f"Command failed (exit {result.returncode}):\n"
                         f"STDOUT: {result.stdout}\nSTDERR: {result.stderr}")
        return result.stdout, result.stderr

    def switch_to_slot_b(self):
        ctl = UKIBOOTCTL_BINARY
        part = '--partition=' + self.bootctl_path
        subprocess.run([ctl, part, 'prepare-switch', '1'], check=True, capture_output=True)
        subprocess.run([ctl, part, 'finalize-switch', '1'], check=True, capture_output=True)

    def test_verify_slot_a(self):
        stdout, _ = self.run_ukibootimg(['--cert', self.cert_path, 'verify'])
        self.assertIn('Signature verification successful', stdout)

    def test_verify_slot_b(self):
        self.switch_to_slot_b()
        stdout, _ = self.run_ukibootimg(['--cert', self.cert_path, 'verify'])
        self.assertIn('Signature verification successful', stdout)

    def test_verify_wrong_cert(self):
        _, stderr = self.run_ukibootimg([
            '--cert', self.other_cert_path, 'verify'
        ], expected_exit_code=1)
        self.assertIn('verification failed', stderr.lower())

    def test_extract_with_verify(self):
        self.run_ukibootimg([
            '--cert', self.cert_path,
            '--output-dir', self.extract_dir,
            'extract', '--verify'
        ])
        with open(os.path.join(self.extract_dir, 'uname'), 'rb') as f:
            self.assertEqual(f.read(), b'6.0.0-slota')

    def test_extract_verify_wrong_cert(self):
        _, stderr = self.run_ukibootimg([
            '--cert', self.other_cert_path,
            '--output-dir', self.extract_dir,
            'extract', '--verify'
        ], expected_exit_code=1)
        # Should fail before extracting anything
        self.assertFalse(os.path.exists(os.path.join(self.extract_dir, 'uname')))

    def test_signed_pe_sections_still_readable(self):
        stdout, _ = self.run_ukibootimg(['status'])
        self.assertIn('TestOS 1.0', stdout)
        self.assertIn('6.0.0-slota', stdout)


if __name__ == '__main__':
    unittest.main(argv=['first-arg-is-ignored'], exit=False)
