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

import argparse
import os
import shutil
import sys
import tempfile

from lxcmgr import lxcmgr
from lxcmgr.paths import LXC_ROOT

def get_layers(container):
    with open(os.path.join(LXC_ROOT, container, 'config')) as f:
        for line in f.read().splitlines():
            if line.startswith('lxc.hook.pre-start'):
                return line.split()[-1].split(',')

def copy(source, destination):
    if os.path.isdir(source):
        shutil.copytree(source, destination, True)
    else:
        shutil.copy2(source, destination)

def extract(args):
    with tempfile.TemporaryDirectory() as tmp_rootfs:
        layers = get_layers(args.container)
        lxcmgr.mount_rootfs(args.container, layers, tmp_rootfs)
        source = os.path.join(tmp_rootfs, args.source.lstrip('/'))
        try:
            copy(source, args.destination)
        except:
            lxcmgr.unmount_rootfs(tmp_rootfs)
            raise
        lxcmgr.unmount_rootfs(tmp_rootfs)

def main(args):
    if args.action == 'prepare':
        # Used with LXC hooks on container startup
        lxcmgr.prepare_container(args.container, args.layers)
    elif args.action == 'cleanup':
        # Used with LXC hooks on container stop
        lxcmgr.cleanup_container(args.container)
    elif args.action == 'extract':
        # Used in install.sh scripts to get files or directories from containers rootfs (excluding persistent mounts)
        extract(args)

parser = argparse.ArgumentParser(description='Collection of auxiliary LXC tools')
subparsers = parser.add_subparsers()

parser_prepare = subparsers.add_parser('prepare', help='Perform pre-start steps for LXC')
parser_prepare.set_defaults(action='prepare')
parser_prepare.add_argument('layers', help='OverlayFS LXC rootfs layers')
parser_prepare.add_argument('container', help='Container name')
parser_prepare.add_argument('lxc', nargs=argparse.REMAINDER)

parser_cleanup = subparsers.add_parser('cleanup', help='Perform post-stop steps for LXC')
parser_cleanup.set_defaults(action='cleanup')
parser_cleanup.add_argument('container', help='Container name')
parser_cleanup.add_argument('lxc', nargs=argparse.REMAINDER)

parser_extract = subparsers.add_parser('extract', help='Extracts files or directories from containers rootfs (excluding persistent mounts)')
parser_extract.set_defaults(action='extract')
parser_extract.add_argument('container', help='Container name')
parser_extract.add_argument('source', help='Source file or directory within the container')
parser_extract.add_argument('destination', help='Destination file or directory on the host')

main(parser.parse_args())