#!/usr/bin/env python3
# group: rw quick
#
# This test covers the basic fleecing workflow, which provides a
# point-in-time snapshot of a node that can be queried over NBD.
#
# Copyright (C) 2018 Red Hat, Inc.
# John helped, too.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program 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 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/>.
#
# Creator/Owner: John Snow <jsnow@redhat.com>

from subprocess import CalledProcessError

import iotests
from iotests import log, qemu_img, qemu_io

iotests.script_initialize(
    supported_fmts=['qcow2'],
    supported_platforms=['linux'],
    required_fmts=['copy-before-write'],
    unsupported_imgopts=['compat']
)

patterns = [('0x5d', '0',         '64k'),
            ('0xd5', '1M',        '64k'),
            ('0xdc', '32M',       '64k'),
            ('0xcd', '0x3ff0000', '64k')]  # 64M - 64K

overwrite = [('0xab', '0',         '64k'), # Full overwrite
             ('0xad', '0x00f8000', '64k'), # Partial-left (1M-32K)
             ('0x1d', '0x2008000', '64k'), # Partial-right (32M+32K)
             ('0xea', '0x3fe0000', '64k')] # Adjacent-left (64M - 128K)

zeroes = [('0', '0x00f8000', '32k'), # Left-end of partial-left (1M-32K)
          ('0', '0x2010000', '32k'), # Right-end of partial-right (32M+64K)
          ('0', '0x3fe0000', '64k')] # overwrite[3]

remainder = [('0xd5', '0x108000',  '32k'), # Right-end of partial-left [1]
             ('0xdc', '32M',       '32k'), # Left-end of partial-right [2]
             ('0xcd', '0x3ff0000', '64k')] # patterns[3]

def do_test(vm, use_cbw, use_snapshot_access_filter, base_img_path,
            fleece_img_path, nbd_sock_path=None,
            target_img_path=None,
            bitmap=False):
    push_backup = target_img_path is not None
    assert (nbd_sock_path is not None) != push_backup
    if push_backup:
        assert use_cbw

    log('--- Setting up images ---')
    log('')

    qemu_img('create', '-f', iotests.imgfmt, base_img_path, '64M')
    if bitmap:
        qemu_img('bitmap', '--add', base_img_path, 'bitmap0')

    if use_snapshot_access_filter:
        assert use_cbw
        qemu_img('create', '-f', 'raw', fleece_img_path, '64M')
    else:
        qemu_img('create', '-f', 'qcow2', fleece_img_path, '64M')

    if push_backup:
        qemu_img('create', '-f', 'qcow2', target_img_path, '64M')

    for p in patterns:
        qemu_io('-f', iotests.imgfmt,
                '-c', 'write -P%s %s %s' % p, base_img_path)

    log('Done')

    log('')
    log('--- Launching VM ---')
    log('')

    src_node = 'source'
    tmp_node = 'temp'
    qom_path = '/machine/peripheral/sda'
    vm.add_blockdev(f'driver={iotests.imgfmt},file.driver=file,'
                    f'file.filename={base_img_path},node-name={src_node}')
    vm.add_device('virtio-scsi')
    vm.add_device(f'scsi-hd,id=sda,drive={src_node}')
    vm.launch()
    log('Done')

    log('')
    log('--- Setting up Fleecing Graph ---')
    log('')


    if use_snapshot_access_filter:
        log(vm.qmp('blockdev-add', {
            'node-name': tmp_node,
            'driver': 'file',
            'filename': fleece_img_path,
        }))
    else:
        # create tmp_node backed by src_node
        log(vm.qmp('blockdev-add', {
            'driver': 'qcow2',
            'node-name': tmp_node,
            'file': {
                'driver': 'file',
                'filename': fleece_img_path,
            },
            'backing': src_node,
        }))

    # Establish CBW from source to fleecing node
    if use_cbw:
        fl_cbw = {
            'driver': 'copy-before-write',
            'node-name': 'fl-cbw',
            'file': src_node,
            'target': tmp_node
        }

        if bitmap:
            fl_cbw['bitmap'] = {'node': src_node, 'name': 'bitmap0'}

        log(vm.qmp('blockdev-add', fl_cbw))

        log(vm.qmp('qom-set', path=qom_path, property='drive', value='fl-cbw'))

        if use_snapshot_access_filter:
            log(vm.qmp('blockdev-add', {
                'driver': 'snapshot-access',
                'node-name': 'fl-access',
                'file': 'fl-cbw',
            }))
    else:
        log(vm.qmp('blockdev-backup',
                   job_id='fleecing',
                   device=src_node,
                   target=tmp_node,
                   sync='none'))

    export_node = 'fl-access' if use_snapshot_access_filter else tmp_node

    if push_backup:
        log('')
        log('--- Starting actual backup ---')
        log('')

        log(vm.qmp('blockdev-add', **{
            'driver': iotests.imgfmt,
            'node-name': 'target',
            'file': {
                'driver': 'file',
                'filename': target_img_path
            }
        }))
        log(vm.qmp('blockdev-backup', device=export_node,
                   sync='full', target='target',
                   job_id='push-backup', speed=1))
    else:
        log('')
        log('--- Setting up NBD Export ---')
        log('')

        nbd_uri = 'nbd+unix:///%s?socket=%s' % (export_node, nbd_sock_path)
        log(vm.qmp('nbd-server-start',
                   {'addr': { 'type': 'unix',
                              'data': { 'path': nbd_sock_path } } }))

        log(vm.qmp('nbd-server-add', device=export_node))

        log('')
        log('--- Sanity Check ---')
        log('')

        for p in patterns + zeroes:
            cmd = 'read -P%s %s %s' % p
            log(cmd)

            try:
                qemu_io('-r', '-f', 'raw', '-c', cmd, nbd_uri)
            except CalledProcessError as exc:
                if bitmap and p in zeroes:
                    log(exc.stdout)
                else:
                    raise

    log('')
    log('--- Testing COW ---')
    log('')

    for p in overwrite:
        cmd = 'write -P%s %s %s' % p
        log(cmd)
        log(vm.hmp_qemu_io(qom_path, cmd, qdev=True))

    if push_backup:
        # Check that previous operations were done during backup, not after
        # If backup is already finished, it's possible that it was finished
        # even before hmp qemu_io write, and we didn't actually test
        # copy-before-write operation. This should not happen, as we use
        # speed=1. But worth checking.
        result = vm.qmp('query-block-jobs')
        assert len(result['return']) == 1

        result = vm.qmp('block-job-set-speed', device='push-backup', speed=0)
        assert result == {'return': {}}

        log(vm.event_wait(name='BLOCK_JOB_COMPLETED',
                          match={'data': {'device': 'push-backup'}}),
            filters=[iotests.filter_qmp_event])
        log(vm.qmp('blockdev-del', node_name='target'))

    log('')
    log('--- Verifying Data ---')
    log('')

    for p in patterns + zeroes:
        cmd = 'read -P%s %s %s' % p
        log(cmd)
        args = ['-r', '-c', cmd]
        if push_backup:
            args += [target_img_path]
        else:
            args += ['-f', 'raw', nbd_uri]

        try:
            qemu_io(*args)
        except CalledProcessError as exc:
            if bitmap and p in zeroes:
                log(exc.stdout)
            else:
                raise

    log('')
    log('--- Cleanup ---')
    log('')

    if not push_backup:
        log(vm.qmp('nbd-server-stop'))

    if use_cbw:
        if use_snapshot_access_filter:
            log(vm.qmp('blockdev-del', node_name='fl-access'))
        log(vm.qmp('qom-set', path=qom_path, property='drive', value=src_node))
        log(vm.qmp('blockdev-del', node_name='fl-cbw'))
    else:
        log(vm.qmp('block-job-cancel', device='fleecing'))
        e = vm.event_wait('BLOCK_JOB_CANCELLED')
        assert e is not None
        log(e, filters=[iotests.filter_qmp_event])

    log(vm.qmp('blockdev-del', node_name=tmp_node))
    vm.shutdown()

    log('')
    log('--- Confirming writes ---')
    log('')

    for p in overwrite + remainder:
        cmd = 'read -P%s %s %s' % p
        log(cmd)
        qemu_io(base_img_path, '-c', cmd)

    log('')
    log('Done')


def test(use_cbw, use_snapshot_access_filter,
         nbd_sock_path=None, target_img_path=None, bitmap=False):
    with iotests.FilePath('base.img') as base_img_path, \
         iotests.FilePath('fleece.img') as fleece_img_path, \
         iotests.VM() as vm:
        do_test(vm, use_cbw, use_snapshot_access_filter, base_img_path,
                fleece_img_path, nbd_sock_path, target_img_path,
                bitmap=bitmap)

def test_pull(use_cbw, use_snapshot_access_filter, bitmap=False):
    with iotests.FilePath('nbd.sock',
                          base_dir=iotests.sock_dir) as nbd_sock_path:
        test(use_cbw, use_snapshot_access_filter, nbd_sock_path, None,
             bitmap=bitmap)

def test_push():
    with iotests.FilePath('target.img') as target_img_path:
        test(True, True, None, target_img_path)


log('=== Test backup(sync=none) based fleecing ===\n')
test_pull(False, False)

log('=== Test cbw-filter based fleecing ===\n')
test_pull(True, False)

log('=== Test fleecing-format based fleecing ===\n')
test_pull(True, True)

log('=== Test fleecing-format based fleecing with bitmap ===\n')
test_pull(True, True, bitmap=True)

log('=== Test push backup with fleecing ===\n')
test_push()
