spoc/usr/bin/spoc-image

#!/usr/bin/python3
# -*- coding: utf-8 -*-

import argparse
import os
import sys

from spoc import repo_local
from spoc import repo_online
from spoc import repo_publish
from spoc.cli import ActionQueue, print_lock, readable_size
from spoc.config import LOCK_FILE
from spoc.depsolver import DepSolver
from spoc.exceptions import ImageNotFoundError
from spoc.flock import locked
from spoc.image import Image
from spoc.imagebuilder import ImageBuilder

def get_image_name(file_path):
    # Read and return image name from image file
    with open(file_path) as f:
        for line in f:
            if line.startswith('IMAGE '):
                return line.split()[1]
    return None

def listing(list_type):
    # Lists images in particular state
    if list_type == 'installed':
        images = repo_local.get_images()
    elif list_type == 'online':
        images = repo_online.get_images()
    elif list_type == 'published':
        images = repo_publish.get_images()
    for image in images:
        print(image)

@locked(LOCK_FILE, print_lock)
def download(image_name):
    # Download and unpack image from online repository
    queue = ActionQueue()
    local_images = repo_local.get_images()
    for layer in repo_online.get_image(image_name)['layers']:
        if layer not in local_images:
            queue.download_image(Image(layer, False))
    queue.process()

@locked(LOCK_FILE, print_lock)
def delete(image_name):
    # Remove the image including all images that have it as one of its parents
    # Check if image is in use
    used_by = [c for c,d in repo_local.get_containers().items() if image_name in d['layers']]
    if used_by:
        sys.exit(f'Error: Image {image_name} is used by container{"s" if len(used_by) > 1 else ""} {", ".join(used_by)}')
    used_images = set()
    for definition in repo_local.get_containers().values():
        used_images.update(definition['layers'])
    # Build dependency tree to safely remove the images in order of dependency
    depsolver = DepSolver()
    for image,definition in repo_local.get_images().items():
        if image_name in definition['layers']:
            image = Image(image)
            depsolver.add(image.name, set(image.layers) - used_images, image)
    # Enqueue and run the removal actions
    queue = ActionQueue()
    for image in reversed(depsolver.solve()):
        queue.delete_image(image)
    queue.process()

@locked(LOCK_FILE, print_lock)
def clean():
    # Remove images which aren't used in any locally defined containers
    used_images = set()
    for definition in repo_local.get_containers().values():
        used_images.update(definition['layers'])
    # Build dependency tree to safely remove the images in order of dependency
    depsolver = DepSolver()
    for image in set(repo_local.get_images()) - used_images:
        image = Image(image)
        depsolver.add(image.name, set(image.layers) - used_images, image)
    # Enqueue and run the removal actions
    queue = ActionQueue()
    for image in reversed(depsolver.solve()):
        queue.delete_image(image)
    queue.process()

@locked(LOCK_FILE, print_lock)
def build(filename, force, do_publish):
    # Check if a build is needed and attempt to build the image from image file
    image_name = get_image_name(filename)
    if force or image_name not in repo_local.get_images():
        image = Image(image_name, False)
        print(f'Building image {image_name} from file {os.path.abspath(filename)}')
        image.delete()
        image.create(ImageBuilder(), filename)
        print(f'Image {image_name} built successfully')
        # If publishing was requested, force publish after successful build
        force = True
    else:
        print(f'Image {image_name} already built, skipping build task')
    if do_publish:
        publish(image_name, force)

def publish(image_name, force):
    # Check if publishing is needed and attempt to publish the image
    if force or image_name not in repo_publish.get_images():
        image = Image(image_name)
        print(f'Publishing image {image_name}')
        image.unpublish()
        size, dlsize = image.publish()
        print(f'Image {image_name} compressed from {readable_size(size)} to {readable_size(dlsize)} and published successfully')
    else:
        print(f'Image {image_name} already published, skipping publish task')

def unpublish(image_name):
    # Remove the image from publish repo
    Image(image_name, False).unpublish()

parser = argparse.ArgumentParser(description='SPOC image manager')
parser.set_defaults(action=None)
subparsers = parser.add_subparsers()

parser_list = subparsers.add_parser('list')
parser_list.set_defaults(action=listing)
parser_list.add_argument('type', choices=('installed', 'online', 'published'), default='installed', const='installed', nargs='?')

parser_download = subparsers.add_parser('download')
parser_download.set_defaults(action=download)
parser_download.add_argument('image')

parser_delete = subparsers.add_parser('delete')
parser_delete.set_defaults(action=delete)
parser_delete.add_argument('image')

parser_clean = subparsers.add_parser('clean')
parser_clean.set_defaults(action=clean)

parser_build = subparsers.add_parser('build')
parser_build.set_defaults(action=build)
parser_build.add_argument('-f', '--force', action='store_true', help='Force rebuild already existing image')
parser_build.add_argument('-p', '--publish', action='store_true', help='Publish the image after successful build')
parser_build.add_argument('filename')

parser_publish = subparsers.add_parser('publish')
parser_publish.set_defaults(action=publish)
parser_publish.add_argument('-f', '--force', action='store_true', help='Force republish already published image')
parser_publish.add_argument('image')

parser_unpublish = subparsers.add_parser('unpublish')
parser_unpublish.set_defaults(action=unpublish)
parser_unpublish.add_argument('image')

args = parser.parse_args()

if args.action is listing:
    listing(args.type)
elif args.action is download:
    download(args.image)
elif args.action is delete:
    delete(args.image)
elif args.action is clean:
    clean()
elif args.action is build:
    build(args.filename, args.force, args.publish)
elif args.action is publish:
    publish(args.image, args.force)
elif args.action is unpublish:
    unpublish(args.image)
else:
    parser.print_usage()