prune unused vagrant-libvirt images

prune_libvirt_images.py

#!/usr/bin/python3
"""
Script to remove unused libvirt images left lying around by libvirt-vagrant.

This is a workaround for
https://github.com/vagrant-libvirt/vagrant-libvirt/issues/85
"""

import argparse
import functools
import re
import subprocess
import typing


__author__ = 'Marius Gedminas <[email protected]>'
__version__ = '1.0'
__license__ = 'GPL v2 or v3'


class Error(SystemExit):
    pass


class LibvirtImage(typing.NamedTuple):
    name: str
    filename: str


class VagrantBox(typing.NamedTuple):
    name: str
    provider: str
    version: str

    @property
    def libvirt_name(self):
        name = self.name.replace('/', '-VAGRANTSLASH-')
        return f'{name}_vagrant_box_image_{self.version}.img'


def collect(fn):
    @functools.wraps(fn)
    def wrapper(*args, **kw):
        return list(fn(*args, **kw))
    return wrapper


@collect
def list_libvirt_images() -> typing.List[LibvirtImage]:
    lines = iter(subprocess.run(
        "virsh -c qemu:///system vol-list default".split(),
        stdout=subprocess.PIPE,
        encoding='UTF-8',
        check=True,
    ).stdout.splitlines())
    for line in lines:
        if line.startswith('-----'):
            break
    for line in lines:
        if not line.strip():
            break
        try:
            name, filename = line.split()
        except ValueError:
            raise Error(f'Could not parse virsh vol-list output:\n{line}')
        yield LibvirtImage(name, filename)


@collect
def list_vagrant_boxes() -> typing.List[VagrantBox]:
    for line in subprocess.run(
        "vagrant box list".split(),
        stdout=subprocess.PIPE,
        encoding='UTF-8',
        check=True,
    ).stdout.splitlines():
        try:
            name, provider, version = line.split()
        except ValueError:
            raise Error(f'Could not parse vagrant box list output:\n{line}')
        yield VagrantBox(name, provider.strip('(,'), version.strip(')'))


def delete_libvirt_image(name: str) -> None:
    subprocess.run(
        ["virsh", "-c", "qemu:///system", "vol-delete", name, "default"],
        check=True,
    )


def main():
    parser = argparse.ArgumentParser(
        description="prune libvirt images not used by any vagrant boxes")
    parser.add_argument(
        "--version", action="version",
        version="%(prog)s version " + __version__)
    parser.add_argument(
        "-v", "--verbose", action="store_true",
        help="print more information")
    parser.add_argument(
        "-m", "--matching", metavar='REGEXP', dest='match',
        help="act only on images with names matching a regular expression")
    parser.add_argument(
        "-n", "--dry-run", action="store_true", default=True,
        help="print names of unused images (default)")
    parser.add_argument(
        "-d", "--delete", action="store_false", dest='dry_run',
        help="delete unused images (DANGEROUS)")
    args = parser.parse_args()

    info = print if args.verbose else lambda *args, **kw: None
    match = re.compile(args.match).search if args.match else lambda n: True

    used_boxes = list_vagrant_boxes()
    info(f"Found {len(used_boxes)} vagrant boxes")
    expected_images = {}
    for box in used_boxes:
        if box.provider == 'libvirt':
            info(f"{box.name} {box.version} -> {box.libvirt_name}")
            expected_images[box.libvirt_name] = box
        else:
            info(f"{box.name} {box.version} uses a different provider"
                 f" ({box.provider})")

    available_images = list_libvirt_images()
    info(f"Found {len(available_images)} libvirt images")
    for image in available_images:
        if not match(image.name):
            info(f"{image.name} does not match {args.match}")
        elif image.name in expected_images:
            info(f"{image.name} is used")
        else:
            info(f"{image.name} is NOT used")
            if args.dry_run:
                if not args.verbose:
                    print(image.name)
            else:
                delete_libvirt_image(image.name)


if __name__ == "__main__":
    main()