<commit_before>from __future__ import absolute_import, division, print_function, unicode_literals

import os, sys, json, time, base64
from argparse import Namespace
from collections import OrderedDict
from io import open

from . import register_parser, logger, config, __version__
from .util.aws import (locate_ami, get_user_data, ensure_vpc, ensure_subnet, ensure_ingress_rule,
                       ensure_security_group, add_tags, get_bdm, resolve_instance_id, resources, clients,
                       gzip_compress_bytes)
from .util.crypto import ensure_ssh_key, new_ssh_key, add_ssh_host_key_to_known_hosts, get_ssh_key_filename
from .launch import launch, parser as launch_parser

def get_bootstrap_files(rootfs_skel_dirs):
    manifest = OrderedDict()
    aegea_conf = os.getenv("AEGEA_CONFIG_FILE")

    for rootfs_skel_dir in rootfs_skel_dirs:
        if rootfs_skel_dir == "auto":
            fn = os.path.join(os.path.dirname(__file__), "rootfs.skel")
        elif aegea_conf:
            # FIXME: not compatible with colon-separated AEGEA_CONFIG_FILE
            fn = os.path.join(os.path.dirname(aegea_conf), rootfs_skel_dir)
        elif os.path.exists(rootfs_skel_dir):
            fn = os.path.abspath(os.path.normpath(rootfs_skel_dir))
        else:
            raise Exception("rootfs_skel directory {} not found".format(fn))
        logger.debug("Trying rootfs.skel: %s" % fn)
        if not os.path.exists(fn):
            raise Exception("rootfs_skel directory {} not found".format(fn))
        for root, dirs, files in os.walk(fn):
            for file_ in files:
                path = os.path.join("/", os.path.relpath(root, fn), file_)
                with open(os.path.join(root, file_), "rb") as fh:
                    content = fh.read()
                    manifest[path] = dict(path=path,
                                          permissions=oct(os.stat(os.path.join(root, file_)).st_mode)[-3:])
                    try:
                        manifest[path].update(content=content.decode())
                    except UnicodeDecodeError:
                        manifest[path].update(content=base64.b64encode(gzip_compress_bytes(content)), encoding="gz+b64")
    return list(manifest.values())

def build_ami(args):
    from .util.ssh import AegeaSSHClient
    ssh_key_filename = get_ssh_key_filename(args, base_name=__name__)
    if args.snapshot_existing_host:
        instance = resources.ec2.Instance(resolve_instance_id(args.snapshot_existing_host))
        args.ami = instance.image_id
    else:
        if args.base_ami == "auto":
            args.ami = locate_ami(product=args.base_ami_product)
        else:
            args.ami = args.base_ami
        hostname = "{}-{}-{}".format(__name__, args.name, int(time.time())).replace(".", "-").replace("_", "-")
        launch_args = launch_parser.parse_args(args=[hostname], namespace=args)
        launch_args.wait_for_ssh = True
        launch_args.iam_role = None
        launch_args.cloud_config_data.update(files=get_bootstrap_files(args.rootfs_skel_dirs))
        instance = resources.ec2.Instance(launch(launch_args)["instance_id"])
    ssh_client = AegeaSSHClient()
    ssh_client.load_system_host_keys()
    ssh_client.connect(instance.public_dns_name, username="ubuntu", key_filename=ssh_key_filename)
    for i in range(900):
        try:
            ssh_client.check_call("ls /var/lib/cloud/data/result.json")
            if ssh_client.check_output("sudo jq .v1.errors /var/lib/cloud/data/result.json").strip() != "[]":
                raise Exception("cloud-init encountered errors")
            break
        except Exception as e:
            if "ENOENT" in str(e) or "EPERM" in str(e) or "No such file or directory" in str(e):
                time.sleep(1)
            else:
                raise
    else:
        raise Exception("cloud-init encountered errors")

    description = "Built by {} for {}".format(__name__, resources.iam.CurrentUser().user.name)
    image = instance.create_image(Name=args.name, Description=description, BlockDeviceMappings=get_bdm())
    tags = dict(tag.split("=", 1) for tag in args.tags)
    # FIXME: add base AMI name
    tags.update(Owner=resources.iam.CurrentUser().user.name, Base=args.ami, AegeaVersion=__version__)
    add_tags(image, **tags)
    logger.info("Waiting for %s to become available...", image.id)
    clients.ec2.get_waiter("image_available").wait(ImageIds=[image.id])<commit_msg>Propagate base AMI metadata in build_image<commit_after>from __future__ import absolute_import, division, print_function, unicode_literals

import os, sys, json, time, base64
from argparse import Namespace
from collections import OrderedDict
from io import open

from . import register_parser, logger, config, __version__
from .util.aws import (locate_ami, get_user_data, ensure_vpc, ensure_subnet, ensure_ingress_rule,
                       ensure_security_group, add_tags, get_bdm, resolve_instance_id, resources, clients,
                       gzip_compress_bytes)
from .util.crypto import ensure_ssh_key, new_ssh_key, add_ssh_host_key_to_known_hosts, get_ssh_key_filename
from .launch import launch, parser as launch_parser

def get_bootstrap_files(rootfs_skel_dirs):
    manifest = OrderedDict()
    aegea_conf = os.getenv("AEGEA_CONFIG_FILE")

    for rootfs_skel_dir in rootfs_skel_dirs:
        if rootfs_skel_dir == "auto":
            fn = os.path.join(os.path.dirname(__file__), "rootfs.skel")
        elif aegea_conf:
            # FIXME: not compatible with colon-separated AEGEA_CONFIG_FILE
            fn = os.path.join(os.path.dirname(aegea_conf), rootfs_skel_dir)
        elif os.path.exists(rootfs_skel_dir):
            fn = os.path.abspath(os.path.normpath(rootfs_skel_dir))
        else:
            raise Exception("rootfs_skel directory {} not found".format(fn))
        logger.debug("Trying rootfs.skel: %s" % fn)
        if not os.path.exists(fn):
            raise Exception("rootfs_skel directory {} not found".format(fn))
        for root, dirs, files in os.walk(fn):
            for file_ in files:
                path = os.path.join("/", os.path.relpath(root, fn), file_)
                with open(os.path.join(root, file_), "rb") as fh:
                    content = fh.read()
                    manifest[path] = dict(path=path,
                                          permissions=oct(os.stat(os.path.join(root, file_)).st_mode)[-3:])
                    try:
                        manifest[path].update(content=content.decode())
                    except UnicodeDecodeError:
                        manifest[path].update(content=base64.b64encode(gzip_compress_bytes(content)), encoding="gz+b64")
    return list(manifest.values())

def build_ami(args):
    from .util.ssh import AegeaSSHClient
    ssh_key_filename = get_ssh_key_filename(args, base_name=__name__)
    if args.snapshot_existing_host:
        instance = resources.ec2.Instance(resolve_instance_id(args.snapshot_existing_host))
        args.ami = instance.image_id
    else:
        if args.base_ami == "auto":
            args.ami = locate_ami(product=args.base_ami_product)
        else:
            args.ami = args.base_ami
        hostname = "{}-{}-{}".format(__name__, args.name, int(time.time())).replace(".", "-").replace("_", "-")
        launch_args = launch_parser.parse_args(args=[hostname], namespace=args)
        launch_args.wait_for_ssh = True
        launch_args.iam_role = None
        launch_args.cloud_config_data.update(files=get_bootstrap_files(args.rootfs_skel_dirs))
        instance = resources.ec2.Instance(launch(launch_args)["instance_id"])
    ssh_client = AegeaSSHClient()
    ssh_client.load_system_host_keys()
    ssh_client.connect(instance.public_dns_name, username="ubuntu", key_filename=ssh_key_filename)
    for i in range(900):
        try:
            ssh_client.check_call("ls /var/lib/cloud/data/result.json")
            if ssh_client.check_output("sudo jq .v1.errors /var/lib/cloud/data/result.json").strip() != "[]":
                raise Exception("cloud-init encountered errors")
            break
        except Exception as e:
            if "ENOENT" in str(e) or "EPERM" in str(e) or "No such file or directory" in str(e):
                time.sleep(1)
            else:
                raise
    else:
        raise Exception("cloud-init encountered errors")

    description = "Built by {} for {}".format(__name__, resources.iam.CurrentUser().user.name)
    image = instance.create_image(Name=args.name, Description=description, BlockDeviceMappings=get_bdm())
    tags = dict(tag.split("=", 1) for tag in args.tags)
    base_ami = resources.ec2.Image(args.ami)
    tags.update(Owner=resources.iam.CurrentUser().user.name, AegeaVersion=__version__,
                Base=base_ami.id, BaseName=base_ami.name, BaseDescription=base_ami.description or "")
    add_tags(image, **tags)
    logger.info("Waiting for %s to become available...", image.id)
    clients.ec2.get_waiter("image_available").wait(ImageIds=[image.id])
{'fraction_non_alphanumeric': 0.11273857079449623, 'fraction_numerical': 0.0041056369285397245, 'mean_word_length': 2.898356401384083, 'pattern_counts': {'":': 6, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4966101', 'n_tokens_mistral': 2848, 'n_tokens_neox': 2757, 'n_words': 542}