#!/usr/bin/env python3

# Copyright (C) 2025 Alberto Ruiz <aruiz@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 tempfile
import subprocess
import os
import struct
import zlib
import unittest

# --- Configuration ---
UKIBOOTCTL_BINARY = os.environ.get("UKIBOOTCTL_PATH", default = "./ukibootctl")

# --- Constants from C Source ---
# UKI_BOOT_CTRL_MAGIC as defined in the C code
UKI_BOOT_CTRL_MAGIC = 1420550408 # This is 0x544C4342 in hex, representing 'BTCL' (little-endian)

# --- Structure Definitions (Python equivalent using struct) ---
# UkiBootSlotInfo: priority (uint8_t), tries_remaining (uint8_t), successful_boot (uint8_t), reserved (uint8_t)
# Format: <BBBB (little-endian, 4 unsigned chars)
SLOT_INFO_FORMAT = "<BBBB"

# UkiBootCtrl: magic (uint32_t), slots[2] (2 * 4 bytes), reserved[8] (8 bytes), CRC32 (uint32_t)
# Format: <I (magic) + SLOT_INFO_FORMAT * 2 + 8s (reserved array) + <I (CRC32)
BOOT_CTRL_FORMAT = "<I" + SLOT_INFO_FORMAT[1:] * 2 + "8s" + "I"

BOOT_CTRL_SIZE = struct.calcsize(BOOT_CTRL_FORMAT)

# The size of the data block over which CRC32 is calculated (full struct minus the CRC32 field itself)
BOOT_CTRL_CRC_DATA_SIZE = BOOT_CTRL_SIZE - struct.calcsize("<I")

def calculate_crc32(data_bytes):
    """Calculates CRC32 for the given bytes, ensuring it's a 32-bit unsigned integer."""
    return zlib.crc32(data_bytes) & 0xFFFFFFFF

def create_boot_ctrl_data(slot0_prio=0, slot0_tries=0, slot0_success=0,
                          slot1_prio=0, slot1_tries=0, slot1_success=0):
    """
    Constructs a byte string representing the UkiBootCtrl structure.
    Initializes slot data and calculates the correct CRC32.
    """
    slot0_data = struct.pack(SLOT_INFO_FORMAT, slot0_prio, slot0_tries, slot0_success, 0)
    slot1_data = struct.pack(SLOT_INFO_FORMAT, slot1_prio, slot1_tries, slot1_success, 0)

    reserved_global_bytes = b'\x00' * 8

    initial_struct_bytes = struct.pack("<I", UKI_BOOT_CTRL_MAGIC) + \
                           slot0_data + \
                           slot1_data + \
                           reserved_global_bytes

    crc_value = calculate_crc32(initial_struct_bytes)

    full_struct_bytes = initial_struct_bytes + struct.pack("<I", crc_value)

    return full_struct_bytes

def read_boot_ctrl_data(file_path):
    """
    Reads the raw bytes from the simulated partition file and unpacks them
    into a human-readable dictionary, also validating the CRC32.
    """
    with open(file_path, "rb") as f:
        full_data_bytes = f.read()

    if len(full_data_bytes) != BOOT_CTRL_SIZE:
        raise ValueError(f"File size mismatch: Expected {BOOT_CTRL_SIZE} bytes, got {len(full_data_bytes)}")

    magic, \
    s0_prio, s0_tries, s0_success, s0_reserved, \
    s1_prio, s1_tries, s1_success, s1_reserved, \
    reserved_global_bytes, \
    crc_from_file = struct.unpack(BOOT_CTRL_FORMAT, full_data_bytes)

    data_for_crc_check = full_data_bytes[:BOOT_CTRL_CRC_DATA_SIZE]
    calculated_crc = calculate_crc32(data_for_crc_check)

    return {
        "magic": magic,
        "slot0": {
            "priority": s0_prio,
            "tries_remaining": s0_tries,
            "successful_boot": s0_success,
            "reserved": s0_reserved
        },
        "slot1": {
            "priority": s1_prio,
            "tries_remaining": s1_tries,
            "successful_boot": s1_success,
            "reserved": s1_reserved
        },
        "reserved_global_bytes": reserved_global_bytes,
        "CRC32_stored": crc_from_file,
        "CRC32_calculated": calculated_crc,
        "CRC32_valid": (crc_from_file == calculated_crc)
    }

class TestUkiBootCtl(unittest.TestCase):

    def setUp(self):
        """Set up for each test: create temporary files and check binary."""
        self.temp_partition_file = tempfile.NamedTemporaryFile(delete=False)
        self.temp_partition_file_path = self.temp_partition_file.name
        self.temp_partition_file.close() # Close to allow other processes (like ukibootctl) to open it

        self.temp_cmdline_file = tempfile.NamedTemporaryFile(delete=False)
        self.temp_cmdline_file_path = self.temp_cmdline_file.name
        self.temp_cmdline_file.close()

        if not os.path.exists(UKIBOOTCTL_BINARY):
            self.fail(f"Error: ukibootctl binary not found at '{UKIBOOTCTL_BINARY}'. "
                      "Please compile it and place it in the same directory as this script.")
        if not os.access(UKIBOOTCTL_BINARY, os.X_OK):
            try:
                os.chmod(UKIBOOTCTL_BINARY, 0o755)
            except OSError as e:
                self.fail(f"Error: Could not make '{UKIBOOTCTL_BINARY}' executable. "
                          f"Please check permissions. {e}")

    def tearDown(self):
        """Clean up after each test: delete temporary files."""
        os.unlink(self.temp_partition_file_path)
        os.unlink(self.temp_cmdline_file_path)

    def run_ukibootctl_command(self, command_args, expected_exit_code=0):
        """
        Helper function to execute the ukibootctl binary with given arguments
        and capture its output.
        """
        cmd = [UKIBOOTCTL_BINARY, "--partition", self.temp_partition_file_path]

        cmd.extend(command_args) # Add the specific command and its arguments

        result = subprocess.run(cmd, capture_output=True, text=True, check=False)
        
        self.assertEqual(result.returncode, expected_exit_code,
                         f"Command '{' '.join(cmd)}' failed with unexpected exit code {result.returncode}.\n"
                         f"STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}")
        
        return result.stdout.strip(), result.stderr.strip()

    def test_dump_command(self):
        """Verify the 'dump' command correctly prints slot information."""
        initial_data_bytes = create_boot_ctrl_data(
            slot0_prio=15, slot0_tries=7, slot0_success=0,  # Slot 0: active-like state
            slot1_prio=14, slot1_tries=0, slot1_success=1   # Slot 1: successful-like state
        )
        with open(self.temp_partition_file_path, "wb") as f:
            f.write(initial_data_bytes)

        stdout, _ = self.run_ukibootctl_command(["dump"])

        self.assertIn("Slot 0: prio: 15, tries_remaining: 7, successful_boot: 0", stdout)
        self.assertIn("Slot 1: prio: 14, tries_remaining: 0, successful_boot: 1", stdout)

    def test_get_active_command_slot0_active(self):
        """Verify 'get-active' identifies slot 0 when it's the highest priority valid slot."""
        initial_data_bytes = create_boot_ctrl_data(
            slot0_prio=15, slot0_tries=7, slot0_success=0,  # Slot 0 is valid and high priority
            slot1_prio=14, slot1_tries=0, slot1_success=1   # Slot 1 is valid but lower priority
        )
        with open(self.temp_partition_file_path, "wb") as f:
            f.write(initial_data_bytes)

        stdout, _ = self.run_ukibootctl_command(["get-active"])
        self.assertEqual(stdout, "0")

    def test_get_active_command_slot1_active(self):
        """Verify 'get-active' identifies slot 1 when it's the highest priority valid slot."""
        initial_data_bytes = create_boot_ctrl_data(
            slot0_prio=14, slot0_tries=0, slot0_success=1,  # Slot 0 is valid but lower priority
            slot1_prio=15, slot1_tries=7, slot1_success=0   # Slot 1 is valid and high priority
        )
        with open(self.temp_partition_file_path, "wb") as f:
            f.write(initial_data_bytes)

        stdout, _ = self.run_ukibootctl_command(["get-active"])
        self.assertEqual(stdout, "1")

    def test_get_active_command_no_active_slot_found(self):
        """Verify 'get-active' handles cases where no valid active slot exists."""
        initial_data_bytes = create_boot_ctrl_data(
            slot0_prio=0, slot0_tries=0, slot0_success=0, # Slot 0 is not valid (success=0, tries=0)
            slot1_prio=0, slot1_tries=0, slot1_success=0  # Slot 1 is not valid
        )
        with open(self.temp_partition_file_path, "wb") as f:
            f.write(initial_data_bytes)

        stdout, stderr = self.run_ukibootctl_command(["get-active"], expected_exit_code=1)
        self.assertEqual(stdout, "")
        self.assertIn("No active slot", stderr)

    def test_prepare_switch_command(self):
        """Test 'prepare-switch' command modifies slot states correctly before an update."""
        initial_data_bytes = create_boot_ctrl_data(
            slot0_prio=15, slot0_tries=5, slot0_success=0,  # Slot 0 is current active
            slot1_prio=14, slot1_tries=0, slot1_success=1   # Slot 1 is a successful old slot
        )
        with open(self.temp_partition_file_path, "wb") as f:
            f.write(initial_data_bytes)

        self.run_ukibootctl_command(["prepare-switch", "1"])

        updated_data = read_boot_ctrl_data(self.temp_partition_file_path)

        self.assertEqual(updated_data["slot0"]["successful_boot"], 1)
        self.assertEqual(updated_data["slot0"]["tries_remaining"], 0)
        self.assertEqual(updated_data["slot1"]["successful_boot"], 0)
        self.assertEqual(updated_data["slot1"]["tries_remaining"], 0)
        self.assertEqual(updated_data["slot1"]["priority"], 0)
        self.assertTrue(updated_data["CRC32_valid"])

    def test_finalize_switch_command(self):
        """Test 'finalize-switch' command modifies slot states correctly after an update."""
        initial_data_bytes = create_boot_ctrl_data(
            slot0_prio=15, slot0_tries=0, slot0_success=1, # Slot 0: marked successful by prepare-switch
            slot1_prio=0, slot1_tries=0, slot1_success=0   # Slot 1: marked non-bootable by prepare-switch
        )
        with open(self.temp_partition_file_path, "wb") as f:
            f.write(initial_data_bytes)

        self.run_ukibootctl_command(["finalize-switch", "1"])

        updated_data = read_boot_ctrl_data(self.temp_partition_file_path)

        self.assertEqual(updated_data["slot0"]["priority"], 14)
        self.assertEqual(updated_data["slot1"]["successful_boot"], 0)
        self.assertEqual(updated_data["slot1"]["tries_remaining"], 7)
        self.assertEqual(updated_data["slot1"]["priority"], 15)
        self.assertTrue(updated_data["CRC32_valid"])

    def test_mark_successful_command_explicit_slot(self):
        """Test 'mark-successful' command with an explicit slot number."""
        initial_data_bytes = create_boot_ctrl_data(
            slot0_prio=15, slot0_tries=3, slot0_success=0,
            slot1_prio=14, slot1_tries=0, slot1_success=1
        )
        with open(self.temp_partition_file_path, "wb") as f:
            f.write(initial_data_bytes)

        self.run_ukibootctl_command(["mark-successful", "0"])

        updated_data = read_boot_ctrl_data(self.temp_partition_file_path)

        self.assertEqual(updated_data["slot0"]["successful_boot"], 1)
        self.assertEqual(updated_data["slot0"]["tries_remaining"], 0)
        self.assertEqual(updated_data["slot0"]["priority"], 15) # Priority should remain unchanged
        self.assertTrue(updated_data["CRC32_valid"])

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