#!/usr/bin/env fuchsia-vendored-python

# Copyright 2022 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 tests for gpt.py."""

import dataclasses
import pathlib
import re
import tempfile
import unittest
from typing import Dict, Tuple

import gpt

_MY_DIR = pathlib.Path(__file__).parent

# A real disk image for testing against, generated via:
#  $ dd if=/dev/zero of=gpt_test_disk.img bs=4M count=1
#  $ sgdisk gpt_test_disk.img \
#      --clear \
#      --new=0:+0:+1MiB \
#      --change-name=1:foo \
#      --typecode=1:9B37FFF6-2E58-466A-983A-F7926D0B04E0 \
#      --partition-guid=1:A271DF36-BAAC-45B7-B806-A636996DAF75 \
#      --new=0:+0:+1MiB \
#      --change-name=2:bar \
#      --typecode=2:421A8BFC-85D9-4D85-ACDA-B64EEC0133E9 \
#      --partition-guid=2:8F1ED054-0475-43FE-AB8D-13C6E709D5A7
_TEST_DISK_PATH = _MY_DIR / "gpt_test_disk.img"

# Expected partitions based on the above disk image.
_TEST_DISK_FOO_PARTITION = gpt.Partition(
    index=1,
    name="foo",
    first_lba=2048,
    last_lba=4095,
    size_lba=2048,
    type_guid="9B37FFF6-2E58-466A-983A-F7926D0B04E0",
    unique_guid="A271DF36-BAAC-45B7-B806-A636996DAF75",
    attr_flags="0000000000000000",
)
_TEST_DISK_BAR_PARTITION = gpt.Partition(
    index=2,
    name="bar",
    first_lba=4096,
    last_lba=6143,
    size_lba=2048,
    type_guid="421A8BFC-85D9-4D85-ACDA-B64EEC0133E9",
    unique_guid="8F1ED054-0475-43FE-AB8D-13C6E709D5A7",
    attr_flags="0000000000000000",
)
_TEST_DISK_PARTITIONS = [_TEST_DISK_FOO_PARTITION, _TEST_DISK_BAR_PARTITION]

# The matching input spec that should generate the above image.
_TEST_DISK_SPEC = {
    "size_mib": 4,
    "block_size": 512,
    "alignment_kib": 1024,
    "partitions": [
        {
            "name": "foo",
            "size_kib": 1024,
            "unique_guid": "A271DF36-BAAC-45B7-B806-A636996DAF75",
            "type_guid": "9B37FFF6-2E58-466A-983A-F7926D0B04E0",
        },
        {
            "name": "bar",
            "size_kib": 1024,
            "unique_guid": "8F1ED054-0475-43FE-AB8D-13C6E709D5A7",
            "type_guid": "421A8BFC-85D9-4D85-ACDA-B64EEC0133E9",
        },
    ],
}


def default_test_spec(
    num_partitions: int, *, disk_size_mib: int = 1, part_size_kib: int = 64
) -> Dict:
    """Returns a partition spec for testing.

    The idea is that this returns a reasonable default, and tests can then
    modify just the pieces that they want to exercise.
    """
    return {
        "size_mib": 1,
        "block_size": 512,
        "alignment_kib": 1,
        "partitions": [
            {
                "name": f"part_{i}",
                "size_kib": part_size_kib,
                "unique_guid": "00000000-1111-2222-3333-444444444444",
                "type_guid": "55555555-6666-7777-8888-999999999999",
            }
            for i in range(num_partitions)
        ],
    }


def create_and_read_gpt(spec: Dict) -> Tuple[Dict, str]:
    """Wraps create_gpt() to also fetch the resulting partitions.

    create_gpt() provides the raw mbr/primary/backup blobs, this
    function re-assembles the blobs into a disk image and calls
    get_partition() so we can verify the blobs were correct.

    Args:
        spec: specification to pass to create_gpt()

    Returns:
        A tuple containing:
            * The result from get_partitions()
            * The create_gpt() README contents
    """
    with tempfile.TemporaryDirectory() as temp_dir:
        temp_dir = pathlib.Path(temp_dir)
        gpt.create_gpt(spec, temp_dir)

        # Re-assemble the disk so we can use get_partitions() to validate.
        mbr = (temp_dir / "mbr.bin").read_bytes()
        primary = (temp_dir / "primary.bin").read_bytes()
        backup = (temp_dir / "backup.bin").read_bytes()
        disk_size = spec["size_mib"] * 1024 * 1024

        image_path = temp_dir / "result.img"
        with open(image_path, "wb") as image_file:
            image_file.write(mbr)
            image_file.write(primary)
            image_file.truncate(disk_size - len(backup))
            image_file.write(backup)

        result = gpt.get_partitions(image_path)
        readme = (temp_dir / "README.md").read_text()
        return (result, readme)


class GptTests(unittest.TestCase):
    def test_get_partitions(self):
        self.assertEqual(
            gpt.get_partitions(_TEST_DISK_PATH), _TEST_DISK_PARTITIONS
        )

    def test_get_partition_by_name(self):
        self.assertEqual(
            gpt.get_partition(_TEST_DISK_PATH, part_name="foo"),
            _TEST_DISK_FOO_PARTITION,
        )

    def test_get_partition_by_index(self):
        self.assertEqual(
            gpt.get_partition(_TEST_DISK_PATH, index=2),
            _TEST_DISK_BAR_PARTITION,
        )

    def test_get_partition_by_name_and_index(self):
        self.assertEqual(
            gpt.get_partition(_TEST_DISK_PATH, part_name="bar", index=2),
            _TEST_DISK_BAR_PARTITION,
        )

    def test_get_partition_by_name_not_found(self):
        with self.assertRaises(ValueError):
            gpt.get_partition(_TEST_DISK_PATH, part_name="not_a_partition")

    def test_get_partition_by_index_not_found(self):
        with self.assertRaises(ValueError):
            gpt.get_partition(_TEST_DISK_PATH, index=5)

    def test_get_partition_by_name_and_index_bad_name(self):
        with self.assertRaises(ValueError):
            gpt.get_partition(_TEST_DISK_PATH, part_name="invalid", index=1)

    def test_get_partition_by_name_and_index_bad_index(self):
        with self.assertRaises(ValueError):
            gpt.get_partition(_TEST_DISK_PATH, part_name="bar", index=4)

    def test_create_gpt(self):
        result, readme = create_and_read_gpt(_TEST_DISK_SPEC)

        self.assertEqual(result, _TEST_DISK_PARTITIONS)

        # Check that our README at least contains a dump of all the partition
        # information by making sure all the values exist somewhere in the text.
        for part in _TEST_DISK_PARTITIONS:
            for value in dataclasses.asdict(part).values():
                self.assertIn(str(value), readme)

    def test_create_gpt_fill(self):
        spec = default_test_spec(
            num_partitions=3, disk_size_mib=1, part_size_kib=64
        )
        spec["partitions"][2]["size_kib"] = "fill"

        result, _ = create_and_read_gpt(spec)

        # The "fill" partition should take up the rest of the partition.
        # 2048 disk - 1 MBR - 33 GPT - 128 part_1 - 128 part_2 - 33 backup GPT.
        self.assertEqual(result[2].size_lba, 1725)

    def test_create_gpt_type_map(self):
        spec = default_test_spec(num_partitions=3)
        spec["partitions"][0]["type_guid"] = "zircon"
        spec["partitions"][1]["type_guid"] = "vbmeta"
        spec["partitions"][2]["type_guid"] = "fvm"

        result, _ = create_and_read_gpt(spec)

        self.assertEqual(
            result[0].type_guid, "9B37FFF6-2E58-466A-983A-F7926D0B04E0"
        )
        self.assertEqual(
            result[1].type_guid, "421A8BFC-85D9-4D85-ACDA-B64EEC0133E9"
        )
        self.assertEqual(
            result[2].type_guid, "49FD7CB8-DF15-4E73-B9D9-992070127F0F"
        )

    def test_create_gpt_random_unique_guid(self):
        spec = default_test_spec(num_partitions=2)
        spec["partitions"][0].pop("unique_guid")
        spec["partitions"][1].pop("unique_guid")

        result, _ = create_and_read_gpt(spec)

        self.assertTrue(re.match(gpt.GUID_RE, result[0].unique_guid))
        self.assertTrue(re.match(gpt.GUID_RE, result[1].unique_guid))
        self.assertNotEqual(result[0].unique_guid, result[1].unique_guid)


if __name__ == "__main__":
    unittest.main()