# SPDX-License-Identifier: GPL-3.0-or-later
"""
Tests for checking Freedom Maker's internal library of actions.
"""

import contextlib
import logging
import os
import random
import stat
import string
import subprocess
import tempfile
from unittest.mock import Mock, call, patch

import pytest

from .. import library, releases


@pytest.fixture(name='random_string')
def fixture_random_string():
    """Generate a random string."""

    def _random_string():
        return ''.join(
            [random.choice(string.ascii_lowercase) for _ in range(8)])

    return _random_string


@pytest.fixture(name='image')
def fixture_image(random_string):
    """Return an image file name."""
    return random_string()


@pytest.fixture(name='state')
def fixture_state(image):
    """Create a temporary mount point and return mid-build state."""
    with tempfile.TemporaryDirectory() as mount_point_directory:
        os.makedirs(mount_point_directory + '/usr/sbin/')
        os.makedirs(mount_point_directory + '/tmp/')
        os.makedirs(mount_point_directory + '/etc/apt')

        yield {
            'image_file': image,
            'mount_point': mount_point_directory,
            'success': True,
        }


@contextlib.contextmanager
def assert_file_change(path, content, expected_content):
    """Context manager to verify that file contents changed as expected."""
    if content is not None:
        with open(path, 'w', encoding='utf-8') as file_handle:
            file_handle.write(content)

    yield

    with open(path, 'r', encoding='utf-8') as file_handle:
        changed_content = file_handle.read()

    assert expected_content == changed_content


def test_run(caplog):
    """Test the run utility."""

    # Environment passing
    expected_environ = {
        'LC_ALL': 'C',
        'LANGUAGE': 'C',
        'LANG': 'C',
        'DEBIAN_FRONTEND': 'noninteractive',
        'DEBCONF_NONINTERACTIVE_SEEN': 'true'
    }
    output = library.run(['bash', '-c', 'set'])
    lines = output.splitlines()
    for key, value in expected_environ.items():
        assert f'{key}={value}'.encode() in lines

    # Argument passing
    output = library.run(['bash', '-c', 'echo $0 $@', '1', '2', '3', '4', '5'])
    assert output == b'1 2 3 4 5\n'

    # Stdout printing
    caplog.clear()
    caplog.set_level(logging.DEBUG)
    library.run(['bash', '-c', 'echo "Hello"'])
    assert caplog.record_tuples == [
        ('freedommaker.library', logging.INFO, 'Executing command - '
         "['bash', '-c', 'echo \"Hello\"'] {}"),
        ('freedommaker.library', logging.DEBUG, '> Hello')
    ]

    # Stderr printing
    caplog.clear()
    library.run(['bash', '-c', 'echo "Hello" 1>&2'])
    assert caplog.record_tuples == [
        ('freedommaker.library', logging.INFO, 'Executing command - '
         "['bash', '-c', 'echo \"Hello\" 1>&2'] {}"),
        ('freedommaker.library', logging.DEBUG, '2> Hello')
    ]

    # Return value
    library.run(['bash', '-c', 'exit 0'], check=True)
    library.run(['bash', '-c', 'exit 0'], check=False)
    library.run(['bash', '-c', 'exit 1'], check=False)
    with pytest.raises(subprocess.CalledProcessError):
        library.run(['bash', '-c', 'exit 1'])


@patch('freedommaker.library.run')
def test_run_in_chroot(run, state):
    """Test executing inside a chroot environment."""
    args = ['1', '2', '3']
    kwargs = {'a': 'x', 'b': 'y'}

    library.run_in_chroot(state, args, **kwargs)
    expected_args = ['chroot', state['mount_point']] + args
    assert run.call_args == call(expected_args, **kwargs)


def test_path_in_mount(state):
    """Test returning a sub-directory in mount point."""
    output = library.path_in_mount(state, 'boot')
    assert output == state['mount_point'] + '/boot'

    output = library.path_in_mount(state, '/boot')
    assert output == '/boot'


def test_schedule_clean(state):
    """Test scheduling cleanup jobs."""
    cleanup = Mock()
    library.schedule_cleanup(state, cleanup, 2, b=3)
    assert state['cleanup'][-1] == [cleanup, (2, ), {'b': 3}]


def test_cleanup(state):
    """Test the cleanup works."""
    method = Mock()
    library.schedule_cleanup(state, method, 1, a=1)
    library.schedule_cleanup(state, method, 2, a=2)
    library.schedule_cleanup(state, method, 3, a=3)
    library.cleanup(state)
    assert method.call_args_list == [((3, ), {
        'a': 3
    }), ((2, ), {
        'a': 2
    }), ((1, ), {
        'a': 1
    })]


@patch('freedommaker.library.run')
def test_create_ram_directory_image(run, image, state):
    """Test that RAM directory is properly created."""
    library.create_ram_directory_image(state, image, '4G')
    run.assert_called_once_with([
        'mount', '-o', 'size=4G', '-t', 'tmpfs', 'tmpfs',
        state['ram_directory'].name
    ])
    assert state['image_file'] == (state['ram_directory'].name + '/' + image)
    assert state['cleanup'] == [[
        library.remove_ram_directory, (state['ram_directory'], ), {}
    ], [library.copy_image, (state, state['image_file'], image), {}]]
    state['ram_directory'].cleanup()


@patch('freedommaker.library.run')
def test_remove_ram_directory(run, random_string):
    """Test removing a RAM directory."""
    directory = Mock(name=random_string())
    library.remove_ram_directory(directory)
    run.assert_called_once_with(['umount', directory.name])
    directory.cleanup.assert_called()


@patch('freedommaker.library.run')
def test_copy_image(run, random_string, image, state):
    """Test copying temp image to final image."""
    temp_image = random_string()
    library.copy_image(state, temp_image, image)
    run.assert_called_once_with(['cp', '--sparse=always', temp_image, image])

    run.reset_mock()
    state['success'] = False
    library.copy_image(state, temp_image, image)
    run.assert_called_once_with(
        ['cp', '--sparse=always', temp_image, image + '.failed'])


def test_create_temp_image(image, state):
    """Test creating a temporary image file on disk."""
    library.create_temp_image(state, image)
    assert state['image_file'] == image + '.temp'
    assert state['cleanup'] == [[
        library.move_image, (state, image + '.temp', image), {}
    ]]


@patch('freedommaker.library.run')
def test_move_image(run, random_string, image, state):
    """Test moving temp image to final image."""
    source_image = random_string()
    library.move_image(state, source_image, image)
    run.assert_called_once_with(['mv', source_image, image])

    run.reset_mock()
    state['success'] = False
    library.move_image(state, source_image, image)
    run.assert_called_once_with(['mv', source_image, image + '.failed'])


@patch('freedommaker.library.run')
def test_create_image(run, image, state):
    """Test creating an image."""
    library.create_image(state, '4G')
    run.assert_called_once_with(
        ['qemu-img', 'create', '-f', 'raw', image, '4G'])


@patch('freedommaker.library.run')
def test_create_partition_table(run, image, state):
    """Test creating a partition table."""
    library.create_partition_table(state, 'msdos')
    run.assert_called_once_with(['parted', '-s', image, 'mklabel', 'msdos'])


@patch('freedommaker.library.run')
def test_create_partition(run, image, state):
    """Test creating a partition table."""
    library.create_partition(state, 'root', '10mib', '50%', 'f2fs')
    run.assert_called_once_with(
        ['parted', '-s', image, 'mkpart', 'primary', 'f2fs', '10mib', '50%'])

    assert state['partitions'] == ['root']

    library.create_partition(state, 'root', '10mib', '50%', 'vfat')
    run.assert_called_with(
        ['parted', '-s', image, 'mkpart', 'primary', 'fat32', '10mib', '50%'])


@patch('freedommaker.library.run')
def test_set_boot_flag(run, image, state):
    """Test that boot flag is properly set."""
    library.set_boot_flag(state, 3)
    run.assert_called_with(['parted', '-s', image, 'set', '3', 'boot', 'on'])


@patch('freedommaker.library.run')
def test_loopback_setup(run, image, state):
    """Test that loopback device is properly setup."""
    state['partitions'] = ['firmware', 'boot', 'root']

    run.return_value = b'''remove x x
add x loop99p1
add x loop99p2
add x loop99p3
modify x x
'''
    library.loopback_setup(state)
    run.assert_called_with(['kpartx', '-asv', image])
    assert state['devices'] == {
        'firmware': '/dev/mapper/loop99p1',
        'boot': '/dev/mapper/loop99p2',
        'root': '/dev/mapper/loop99p3'
    }
    assert state['loop_device'] == '/dev/loop99'
    assert state['cleanup'] == [
        [library.force_release_loop_device, ('/dev/loop99', ), {}],
        [library.force_release_partition_loop, ('/dev/mapper/loop99p1', ), {}],
        [library.force_release_partition_loop, ('/dev/mapper/loop99p2', ), {}],
        [library.force_release_partition_loop, ('/dev/mapper/loop99p3', ), {}],
        [library.loopback_teardown, (image, ), {}]
    ]


@patch('freedommaker.library.run')
def test_force_release_partition_loop(run):
    """Test loop device is forcefully released."""
    library.force_release_partition_loop('/dev/test/loop99')
    run.assert_called_with(['dmsetup', 'remove', '/dev/test/loop99'],
                           check=False)


@patch('freedommaker.library.run')
def test_force_release_loop_device(run):
    """Test loop device is forcefully released."""
    library.force_release_partition_loop('/dev/test/loop99')
    run.assert_called_with(['dmsetup', 'remove', '/dev/test/loop99'],
                           check=False)


@patch('freedommaker.library.run')
def test_loopback_teardown(run):
    """Test tearing down of loopback."""
    library.loopback_teardown('/dev/test/loop99')
    run.assert_called_with(['kpartx', '-dsv', '/dev/test/loop99'])


@patch('freedommaker.library.run')
def test_create_filesystem(run):
    """Test creating filesystem."""
    library.create_filesystem('/dev/test/loop99p1', 'btrfs')
    assert run.call_args_list == [
        call(['mkfs', '-t', 'btrfs', '/dev/test/loop99p1']),
        call(['udevadm', 'trigger', '/dev/test/loop99p1']),
        call(['udevadm', 'settle'])
    ]


@patch('freedommaker.library.run')
def test_mount_filesystem(run, state):
    """Test mounting a filesystem and setting proper state."""
    state['devices'] = {
        'root': '/dev/test/loop99p1',
        'firmware': '/dev/test/loop99p2'
    }
    library.mount_filesystem(state, 'root', None)
    run.assert_called_with(
        ['mount', '/dev/test/loop99p1', state['mount_point']])

    library.mount_filesystem(state, 'firmware', 'boot/firmware')
    run.assert_called_with([
        'mount', '/dev/test/loop99p2', state['mount_point'] + '/boot/firmware'
    ])

    library.mount_filesystem(state, '/dev/pts', 'dev/pts', is_bind_mount=True)
    run.assert_called_with(
        ['mount', '/dev/pts', state['mount_point'] + '/dev/pts', '-o', 'bind'])

    sub_mount_points = {
        'root': None,
        'firmware': 'boot/firmware',
        '/dev/pts': 'dev/pts'
    }
    assert state['sub_mount_points'] == sub_mount_points


@patch('freedommaker.library.run')
def test_unmount_filesystem(run, state):
    """Test unmounting a filesystem."""
    library.unmount_filesystem('/dev/', state['mount_point'], True)
    assert run.call_args_list == [
        call(['umount', state['mount_point']], check=True)
    ]


@patch('freedommaker.library.run')
def test_process_cleanup(run, state):
    """Test cleaning up processes."""
    library.process_cleanup(state)
    assert run.call_args_list == [
        call(['fuser', '-mvk', state['mount_point']], check=False),
        call(['fuser', '-mvk', state['mount_point']], check=False)
    ]


@patch('freedommaker.library.run')
def test_setup_extra_storage(run, image, state, random_string):
    """Test that setting up extra storage works."""
    extra_storage_file = image + '.extra'
    size = '100M'
    mount_point = state['mount_point']
    loop_device = random_string()
    run.return_value = f'{loop_device}\n'.encode()

    library.setup_extra_storage(state, 'ext4', '100M')
    run.assert_not_called()

    library.setup_extra_storage(state, 'btrfs', '100M')
    run.assert_has_calls([
        call(['qemu-img', 'create', '-f', 'raw', extra_storage_file, size]),
        call(['losetup', '--show', '--find', extra_storage_file]),
        call(['btrfs', 'device', 'add', loop_device, mount_point])
    ])

    assert state['cleanup'] == [[
        library.cleanup_extra_storage,
        (state, loop_device, extra_storage_file), {}
    ]]


@patch('freedommaker.library.run')
def test_cleanup_extra_storage(run, image, state, random_string):
    """Test that extra storage will be cleaned properly."""
    extra_storage_file = image + '.extra'
    loop_device = random_string()
    mount_point = state['mount_point']

    library.cleanup_extra_storage(state, loop_device, extra_storage_file)

    run.assert_has_calls([
        call(['btrfs', 'balance', 'start', '-musage=0', mount_point],
             check=False),
        call(['btrfs', 'balance', 'start', '-dusage=0', mount_point],
             check=False),
        call(['btrfs', 'balance', 'start', '-musage=20', mount_point],
             check=False),
        call(['btrfs', 'balance', 'start', '-dusage=20', mount_point],
             check=False),
        call(['btrfs', 'balance', 'start', '-musage=40', mount_point],
             check=False),
        call(['btrfs', 'balance', 'start', '-dusage=40', mount_point],
             check=False),
        call(['btrfs', 'balance', 'start', '-musage=60', mount_point],
             check=False),
        call(['btrfs', 'balance', 'start', '-dusage=60', mount_point],
             check=False),
        call(['btrfs', 'balance', 'start', '-musage=80', mount_point],
             check=False),
        call(['btrfs', 'balance', 'start', '-dusage=80', mount_point],
             check=False),
        call(['btrfs', 'balance', 'start', '-mconvert=dup', mount_point],
             check=False),
        call(['btrfs', 'device', 'remove', loop_device, mount_point]),
        call(['losetup', '--detach', loop_device]),
        call(['rm', '-f', extra_storage_file]),
        call(['btrfs', 'balance', 'start', '-musage=0', mount_point],
             check=False),
        call(['btrfs', 'balance', 'start', '-dusage=0', mount_point],
             check=False),
        call(['btrfs', 'balance', 'start', '-musage=20', mount_point],
             check=False),
        call(['btrfs', 'balance', 'start', '-dusage=20', mount_point],
             check=False),
        call(['btrfs', 'balance', 'start', '-musage=40', mount_point],
             check=False),
        call(['btrfs', 'balance', 'start', '-dusage=40', mount_point],
             check=False),
        call(['btrfs', 'balance', 'start', '-musage=60', mount_point],
             check=False),
        call(['btrfs', 'balance', 'start', '-dusage=60', mount_point],
             check=False),
        call(['btrfs', 'balance', 'start', '-musage=80', mount_point],
             check=False),
        call(['btrfs', 'balance', 'start', '-dusage=80', mount_point],
             check=False),
    ])


@patch('freedommaker.library.run')
def test_debootstrap(run, state):
    """Test debootstrapping."""
    library.debootstrap(state, 'amd64', 'stretch', 'minbase',
                        ['main', 'contrib'], ['p1', 'p2'],
                        'http://deb.debian.org/debian')
    run.assert_called_with([
        'debootstrap', '--arch=amd64', '--variant=minbase',
        '--components=main,contrib', '--include=p1,p2', 'stretch',
        state['mount_point'], 'http://deb.debian.org/debian'
    ])

    assert state['cleanup'] == [[
        library.unmount_filesystem,
        (None, state['mount_point'] + '/etc/machine-id'), {
            'check': False
        }
    ]]


def test_no_daemon_policy(state):
    """Test that no daemon run policy is properly set."""
    file_path = state['mount_point'] + '/usr/sbin/policy-rc.d'
    with library.no_run_daemon_policy(state):
        with open(file_path, 'r', encoding='utf-8') as file_handle:
            contents = file_handle.read()

        assert contents == '#!/bin/sh\nexit 101\n'
        assert oct(os.stat(file_path)[stat.ST_MODE])[-3:] == '755'

        assert os.path.isfile(file_path)


@patch('freedommaker.library.run_in_chroot')
def test_install_package(run, state):
    """Test installing a package."""
    library.install_package(state, 'nmap')
    run.assert_called_with(state, ['apt-get', 'install', '-y', 'nmap'])


@patch('freedommaker.library.run_in_chroot')
def test_install_from_backports(run, state):
    """Test installing a package from backports."""
    library.install_from_backports(
        state, ('freedombox', 'freedombox-doc-en', 'freedombox-doc-es'))

    assert run.call_args_list == [
        call(state, [
            'apt-get', 'install', '-y',
            f'freedombox/{releases.STABLE_CODENAME}-backports',
            f'freedombox-doc-en/{releases.STABLE_CODENAME}-backports',
            f'freedombox-doc-es/{releases.STABLE_CODENAME}-backports'
        ]),
        call(state, ['apt-get', 'autoremove', '-y'])
    ]


@patch('freedommaker.library.install_package')
@patch('freedommaker.library.run_in_chroot')
def test_install_custom_package(run, _install_package, state):
    """Test installing a custom package."""
    with tempfile.NamedTemporaryFile() as file_path:
        library.install_custom_package(state, file_path.name)

        run.assert_called_with(state, [
            'apt', 'install', '-y', '--allow-downgrades',
            '/tmp/' + os.path.basename(file_path.name)
        ])


def test_set_hostname(state):
    """Test that hostname is properly set."""
    hosts_path = state['mount_point'] + '/etc/hosts'
    hostname_path = state['mount_point'] + '/etc/hostname'

    content = '''127.0.0.1 localhost
::1     localhost ip6-localhost ip6-loopback
ff02::1 ip6-allnodes
ff02::2 ip6-allrouters
'''
    expected_content = content + '''127.0.1.1 fbx\n'''
    with assert_file_change(hosts_path, content, expected_content):
        with assert_file_change(hostname_path, 'test', 'fbx\n'):
            library.set_hostname(state, 'fbx')


def test_get_fstab_options():
    """Test getting fstab options for a file system."""
    options = library.get_fstab_options('ext4')
    assert options == 'errors=remount-ro'

    options = library.get_fstab_options('btrfs')
    assert options == 'compress=zstd:1'


@patch('freedommaker.library.run', return_value=b'test-uuid')
def test_get_uuid_of_device(run):
    """Test retrieving UUID of device."""
    response = library.get_uuid_of_device('/dev/test/loop99p1')
    run.assert_called_with(
        ['blkid', '--output=value', '--match-tag=UUID', '/dev/test/loop99p1'])
    assert response == 'test-uuid'


@patch('freedommaker.library.get_uuid_of_device')
def test_add_fstab_entry(get_uuid, state):
    """Test adding entries to /etc/fstab."""
    fstab_path = state['mount_point'] + '/etc/fstab'
    state['devices'] = {
        'root': '/dev/test/loop99p1',
        'boot': '/dev/test/loop99p2'
    }
    state['sub_mount_points'] = {'root': None, 'boot': 'boot'}

    expected_content = 'UUID=root-uuid / btrfs compress=zstd:1 0 1\n'
    get_uuid.return_value = 'root-uuid'
    with assert_file_change(fstab_path, 'initial-trash', expected_content):
        library.add_fstab_entry(state, 'root', 'btrfs', 1, append=False)

    expected_content += 'UUID=boot-uuid /boot ext4 errors=remount-ro 0 2\n'
    get_uuid.return_value = 'boot-uuid'
    with assert_file_change(fstab_path, None, expected_content):
        library.add_fstab_entry(state, 'boot', 'ext4', 2, append=True)


@patch('freedommaker.library.run_in_chroot')
def test_install_grub(run, state):
    """Test installing grub boot loader."""
    state['loop_device'] = '/dev/test/loop99'
    library.install_grub(state)
    assert run.call_args_list == [
        call(state, ['update-grub']),
        call(state, ['grub-install', '/dev/test/loop99'])
    ]

    run.reset_mock()
    library.install_grub(state, target='arm-efi')
    assert run.call_args_list == [
        call(state, ['update-grub']),
        call(state, ['grub-install', '/dev/test/loop99', '--target=arm-efi'])
    ]

    run.reset_mock()
    library.install_grub(state, target='arm-efi', is_efi=True)
    assert run.call_args_list == [
        call(state, ['update-grub']),
        call(state, [
            'grub-install', '/dev/test/loop99', '--target=arm-efi',
            '--no-nvram'
        ])
    ]


@patch('freedommaker.library.run_in_chroot')
def test_setup_apt(run, state):
    """Test setting up apt."""
    sources_path = state['mount_point'] + '/etc/apt/sources.list'

    stable_content = '''
deb http://deb.debian.org/debian stable main
deb-src http://deb.debian.org/debian stable main

deb http://deb.debian.org/debian stable-updates main
deb-src http://deb.debian.org/debian stable-updates main

deb http://security.debian.org/debian-security/ stable-security main
deb-src http://security.debian.org/debian-security/ stable-security main
'''
    with assert_file_change(sources_path, None, stable_content):
        library.setup_apt(state, 'http://deb.debian.org/debian', 'stable',
                          ['main'])

    assert run.call_args_list == [
        call(state, ['apt-get', 'update']),
        call(state, ['apt-get', 'clean'])
    ]

    unstable_content = '''
deb http://ftp.us.debian.org/debian unstable main contrib non-free-firmware
deb-src http://ftp.us.debian.org/debian unstable main contrib non-free-firmware
'''
    with assert_file_change(sources_path, None, unstable_content):
        library.setup_apt(state, 'http://ftp.us.debian.org/debian', 'unstable',
                          ['main', 'contrib', 'non-free-firmware'])


@patch('freedommaker.library.run_in_chroot')
def test_setup_flash_kernel(run, state):
    """Test setting up flash kernel."""
    machine_path = state['mount_point'] + '/etc/flash-kernel/machine'
    expected_content = 'test-machine'
    with assert_file_change(machine_path, None, expected_content):
        library.setup_flash_kernel(state, 'test-machine', None, 'ext2')

    assert run.call_args_list == [
        call(state, ['apt-get', 'install', '-y', 'flash-kernel']),
        call(state, ['flash-kernel'])
    ]

    run.reset_mock()
    with assert_file_change(machine_path, None, expected_content):
        library.setup_flash_kernel(state, 'test-machine', 'debug', 'vfat')

    selection = b'flash-kernel flash-kernel/linux_cmdline string debug'
    assert run.call_args_list == [
        call(state, ['debconf-set-selections'], input=selection),
        call(state, ['apt-get', 'install', '-y', 'flash-kernel']),
    ]


@patch('freedommaker.library.run_in_chroot')
def test_update_initramfs(run, state):
    """Test updating initramfs."""
    library.update_initramfs(state)

    assert run.call_args_list == [call(state, ['update-initramfs', '-u'])]


@patch('freedommaker.library.run')
def test_install_boot_loader_path(run, image, state):
    """Test installing boot loader components using dd."""
    path = 'u-boot/path'
    full_path = state['mount_point'] + '/' + path
    library.install_boot_loader_part(state, path, '533', '515')
    run.assert_called_with([
        'dd', 'if=' + full_path, 'of=' + image, 'seek=533', 'bs=515',
        'conv=notrunc'
    ])

    library.install_boot_loader_part(state, path, '533', '515', '90')
    run.assert_called_with([
        'dd', 'if=' + full_path, 'of=' + image, 'seek=533', 'bs=515',
        'conv=notrunc', 'count=90'
    ])


@patch('freedommaker.library.run')
def test_fill_free_space_with_zeros(run, state):
    """Test filling free space with zeros."""
    zeros_path = state['mount_point'] + '/ZEROS'
    library.fill_free_space_with_zeros(state)
    assert run.call_args_list == [
        call(['dd', 'if=/dev/zero', 'of=' + zeros_path, 'bs=1M'], check=False),
        call(['rm', '-f', zeros_path])
    ]


@patch('freedommaker.library.run')
def test_compress(run, random_string):
    """Test compressing an image."""
    archive_file = random_string()
    image_file = random_string()

    library.compress(archive_file, image_file)
    run.assert_called_with([
        'xz', '--no-warn', '--threads=0', '--memlimit=50%', '-9', '--force',
        image_file
    ])


@patch('os.remove')
@patch('freedommaker.library.run')
def test_sign(run, remove, random_string):
    """Test running signing with GPG."""
    archive = random_string()
    signature = archive + '.sig'
    remove.side_effect = FileNotFoundError
    library.sign(archive)
    run.assert_called_with(
        ['gpg', '--output', signature, '--detach-sig', archive])
