contrib/automation/hgautomation/aws.py

43075:57875cf423c9

import time

import boto3
import botocore.exceptions

from .linux import (
    BOOTSTRAP_DEBIAN,
)
from .ssh import (
    exec_command as ssh_exec_command,
    wait_for_ssh,
)
from .winrm import (
    run_powershell,
    wait_for_winrm,
)


SOURCE_ROOT = pathlib.Path(os.path.abspath(__file__)).parent.parent.parent.parent

INSTALL_WINDOWS_DEPENDENCIES = (SOURCE_ROOT / 'contrib' /
                                'install-windows-dependencies.ps1')


INSTANCE_TYPES_WITH_STORAGE = {
    'c5d',
    'd2',

                    {
                        'CidrIp': '0.0.0.0/0',
                        'Description': 'RDP from entire Internet',
                    },
                ],

            },
            {
                'FromPort': 5985,
                'ToPort': 5986,
                'IpProtocol': 'tcp',

                    {
                        'CidrIp': '0.0.0.0/0',
                        'Description': 'PowerShell Remoting (Windows Remote Management)',
                    },
                ],
            }
        ],
    },
}

}
'''.strip()


IAM_INSTANCE_PROFILES = {
    'ephemeral-ec2-1': {
        'roles': [
            'ephemeral-ec2-role-1',
        ],
    }
}


# User Data for Windows EC2 instance. Mainly used to set the password
# and configure WinRM.

class AWSConnection:
    """Manages the state of a connection with AWS."""

    def __init__(self, automation, region: str, ensure_ec2_state: bool=True):
        self.automation = automation
        self.local_state_path = automation.state_path

        self.prefix = 'hg-'

def rsa_key_fingerprint(p: pathlib.Path):
    """Compute the fingerprint of an RSA private key."""

    # TODO use rsa package.
    res = subprocess.run(
        ['openssl', 'pkcs8', '-in', str(p), '-nocrypt', '-topk8',
         '-outform', 'DER'],
        capture_output=True,
        check=True)

    sha1 = hashlib.sha1(res.stdout).hexdigest()
    return ':'.join(a + b for a, b in zip(sha1[::2], sha1[1::2]))


def ensure_key_pairs(state_path: pathlib.Path, ec2resource, prefix='hg-'):
    remote_existing = {}

    for kpi in ec2resource.key_pairs.all():
        if kpi.name.startswith(prefix):
            remote_existing[kpi.name[len(prefix):]] = kpi.key_fingerprint

    # Validate that we have these keys locally.
    key_path = state_path / 'keys'
    key_path.mkdir(exist_ok=True, mode=0o700)

    for f in sorted(os.listdir(key_path)):
        if not f.startswith('keypair-') or not f.endswith('.pub'):
            continue

        name = f[len('keypair-'):-len('.pub')]

        pub_full = key_path / f
        priv_full = key_path / ('keypair-%s' % name)

        with open(pub_full, 'r', encoding='ascii') as fh:
            data = fh.read()

        if not data.startswith('ssh-rsa '):
            print('unexpected format for key pair file: %s; removing' %
                  pub_full)
            pub_full.unlink()
            priv_full.unlink()
            continue

        local_existing[name] = rsa_key_fingerprint(priv_full)

            print('local key %s does not exist remotely' % name)
            remove_local(name)
            del local_existing[name]

        elif remote_existing[name] != local_existing[name]:
            print('key fingerprint mismatch for %s; '
                  'removing from local and remote' % name)
            remove_local(name)
            remove_remote('%s%s' % (prefix, name))
            del local_existing[name]
            del remote_existing[name]

        # SSH public key can be extracted via `ssh-keygen`.
        with pub_full.open('w', encoding='ascii') as fh:
            subprocess.run(
                ['ssh-keygen', '-y', '-f', str(priv_full)],
                stdout=fh,
                check=True)

        pub_full.chmod(0o0600)


def delete_instance_profile(profile):
    for role in profile.roles:
        print('removing role %s from instance profile %s' % (role.name,
                                                             profile.name))
        profile.remove_role(RoleName=role.name)

    print('deleting instance profile %s' % profile.name)
    profile.delete()

    remote_profiles = {}

    for profile in iamresource.instance_profiles.all():
        if profile.name.startswith(prefix):
            remote_profiles[profile.name[len(prefix):]] = profile

    for name in sorted(set(remote_profiles) - set(IAM_INSTANCE_PROFILES)):
        delete_instance_profile(remote_profiles[name])
        del remote_profiles[name]

    remote_roles = {}

    for role in iamresource.roles.all():
        if role.name.startswith(prefix):
            remote_roles[role.name[len(prefix):]] = role

    for name in sorted(set(remote_roles) - set(IAM_ROLES)):
        role = remote_roles[name]

        print('removing role %s' % role.name)

    for name in sorted(set(IAM_INSTANCE_PROFILES) - set(remote_profiles)):
        actual = '%s%s' % (prefix, name)
        print('creating IAM instance profile %s' % actual)

        profile = iamresource.create_instance_profile(
            InstanceProfileName=actual)
        remote_profiles[name] = profile

        waiter = iamclient.get_waiter('instance_profile_exists')
        waiter.wait(InstanceProfileName=actual)
        print('IAM instance profile %s is available' % actual)

def find_image(ec2resource, owner_id, name):
    """Find an AMI by its owner ID and name."""

    images = ec2resource.images.filter(
        Filters=[
            {
                'Name': 'owner-id',
                'Values': [owner_id],
            },
            {
                'Name': 'state',
                'Values': ['available'],
            },
            {
                'Name': 'image-type',
                'Values': ['machine'],
            },
            {
                'Name': 'name',
                'Values': [name],
            },
        ])

    for image in images:
        return image

    raise Exception('unable to find image for %s' % name)

    """
    existing = {}

    for group in ec2resource.security_groups.all():
        if group.group_name.startswith(prefix):
            existing[group.group_name[len(prefix):]] = group

    purge = set(existing) - set(SECURITY_GROUPS)

    for name in sorted(purge):
        group = existing[name]

        actual = '%s%s' % (prefix, name)
        print('adding security group %s' % actual)

        group_res = ec2resource.create_security_group(
            Description=group['description'],
            GroupName=actual,
        )

        group_res.authorize_ingress(
            IpPermissions=group['ingress'],
        )

        security_groups[name] = group_res

    return security_groups

            if not instance.public_ip_address:
                time.sleep(2)
                instance.reload()
                continue

            print('public IP address for %s: %s' % (
                instance.id, instance.public_ip_address))
            break


def remove_ami(ec2resource, image):
    """Remove an AMI and its underlying snapshots."""

def wait_for_ssm(ssmclient, instances):
    """Wait for SSM to come online for an iterable of instance IDs."""
    while True:
        res = ssmclient.describe_instance_information(
            Filters=[
                {
                    'Key': 'InstanceIds',
                    'Values': [i.id for i in instances],
                },
            ],
        )

        available = len(res['InstanceInformationList'])
        wanted = len(instances)

    res = ssmclient.send_command(
        InstanceIds=[i.id for i in instances],
        DocumentName=document_name,
        Parameters=parameters,
        CloudWatchOutputConfig={
            'CloudWatchOutputEnabled': True,
        },
    )

    command_id = res['Command']['CommandId']

    for instance in instances:
        while True:
            try:
                res = ssmclient.get_command_invocation(
                    CommandId=command_id,
                    InstanceId=instance.id,
                )
            except botocore.exceptions.ClientError as e:
                if e.response['Error']['Code'] == 'InvocationDoesNotExist':
                    print('could not find SSM command invocation; waiting')
                    time.sleep(1)

            if res['Status'] == 'Success':
                break
            elif res['Status'] in ('Pending', 'InProgress', 'Delayed'):
                time.sleep(2)
            else:
                raise Exception('command failed on %s: %s' % (
                    instance.id, res['Status']))


@contextlib.contextmanager
def temporary_ec2_instances(ec2resource, config):
    """Create temporary EC2 instances.

    config = copy.deepcopy(config)
    config['IamInstanceProfile'] = {
        'Name': 'hg-ephemeral-ec2-1',
    }
    config.setdefault('TagSpecifications', []).append({
        'ResourceType': 'instance',
        'Tags': [{'Key': 'Name', 'Value': 'hg-temp-windows'}],
    })
    config['UserData'] = WINDOWS_USER_DATA % password

    with temporary_ec2_instances(c.ec2resource, config) as instances:
        wait_for_ip_addresses(instances)

        print('waiting for Windows Remote Management service...')

        for instance in instances:
            client = wait_for_winrm(instance.public_ip_address, 'Administrator', password)
            print('established WinRM connection to %s' % instance.id)
            instance.winrm_client = client

        yield instances

    """
    # Find existing AMIs with this name and delete the ones that are invalid.
    # Store a reference to a good image so it can be returned one the
    # image state is reconciled.
    images = ec2resource.images.filter(
        Filters=[{'Name': 'name', 'Values': [name]}])

    existing_image = None

    for image in images:
        if image.tags is None:
            print('image %s for %s lacks required tags; removing' % (
                image.id, image.name))
            remove_ami(ec2resource, image)
        else:
            tags = {t['Key']: t['Value'] for t in image.tags}

            if tags.get('HGIMAGEFINGERPRINT') == fingerprint:
                existing_image = image
            else:
                print('image %s for %s has wrong fingerprint; removing' % (
                      image.id, image.name))
                remove_ami(ec2resource, image)

    return existing_image


def create_ami_from_instance(ec2client, instance, name, description,
                             fingerprint):
    """Create an AMI from a running instance.

    Returns the ``ec2resource.Image`` representing the created AMI.
    """
    instance.stop()

    ec2client.get_waiter('instance_stopped').wait(
        InstanceIds=[instance.id],
        WaiterConfig={
            'Delay': 5,
        })
    print('%s is stopped' % instance.id)

    image = instance.create_image(
        Name=name,
        Description=description,
    )

    image.create_tags(Tags=[
        {
            'Key': 'HGIMAGEFINGERPRINT',
            'Value': fingerprint,
        },
    ])

    print('waiting for image %s' % image.id)

    ec2client.get_waiter('image_available').wait(
        ImageIds=[image.id],
    )

    print('image %s available as %s' % (image.id, image.name))

    return image

            'debian-stretch-hvm-x86_64-gp2-2019-09-08-17994',
        )
        ssh_username = 'admin'
    elif distro == 'debian10':
        image = find_image(
            ec2resource,
            DEBIAN_ACCOUNT_ID_2,
            'debian-10-amd64-20190909-10',
        )
        ssh_username = 'admin'
    elif distro == 'ubuntu18.04':
        image = find_image(
            ec2resource,

        'MaxCount': 1,
        'MinCount': 1,
        'SecurityGroupIds': [c.security_groups['linux-dev-1'].id],
    }

    requirements2_path = (pathlib.Path(__file__).parent.parent /
                          'linux-requirements-py2.txt')
    requirements3_path = (pathlib.Path(__file__).parent.parent /
                          'linux-requirements-py3.txt')
    with requirements2_path.open('r', encoding='utf-8') as fh:
        requirements2 = fh.read()
    with requirements3_path.open('r', encoding='utf-8') as fh:
        requirements3 = fh.read()

    # Compute a deterministic fingerprint to determine whether image needs to
    # be regenerated.
    fingerprint = resolve_fingerprint({
        'instance_config': config,
        'bootstrap_script': BOOTSTRAP_DEBIAN,
        'requirements_py2': requirements2,
        'requirements_py3': requirements3,
    })

    existing_image = find_and_reconcile_image(ec2resource, name, fingerprint)

    if existing_image:
        return existing_image

        wait_for_ip_addresses(instances)

        instance = instances[0]

        client = wait_for_ssh(
            instance.public_ip_address, 22,
            username=ssh_username,
            key_filename=str(c.key_pair_path_private('automation')))

        home = '/home/%s' % ssh_username

        with client:
            print('connecting to SSH server')

            with sftp.open('%s/requirements-py3.txt' % home, 'wb') as fh:
                fh.write(requirements3)
                fh.chmod(0o0700)

            print('executing bootstrap')
            chan, stdin, stdout = ssh_exec_command(client,
                                                   '%s/bootstrap' % home)
            stdin.close()

            for line in stdout:
                print(line, end='')

            res = chan.recv_exit_status()
            if res:
                raise Exception('non-0 exit from bootstrap: %d' % res)

            print('bootstrap completed; stopping %s to create %s' % (
                  instance.id, name))

        return create_ami_from_instance(ec2client, instance, name,
                                        'Mercurial Linux development environment',
                                        fingerprint)


@contextlib.contextmanager
def temporary_linux_dev_instances(c: AWSConnection, image, instance_type,
                                  prefix='hg-', ensure_extra_volume=False):
    """Create temporary Linux development EC2 instances.

    Context manager resolves to a list of ``ec2.Instance`` that were created
    and are running.

        }
    ]

    # This is not an exhaustive list of instance types having instance storage.
    # But
    if (ensure_extra_volume
        and not instance_type.startswith(tuple(INSTANCE_TYPES_WITH_STORAGE))):
        main_device = block_device_mappings[0]['DeviceName']

        if main_device == 'xvda':
            second_device = 'xvdb'
        elif main_device == '/dev/sda1':
            second_device = '/dev/sdb'
        else:
            raise ValueError('unhandled primary EBS device name: %s' %
                             main_device)

        block_device_mappings.append({
            'DeviceName': second_device,
            'Ebs': {
                'DeleteOnTermination': True,
                'VolumeSize': 8,
                'VolumeType': 'gp2',
            }
        })

    config = {
        'BlockDeviceMappings': block_device_mappings,
        'EbsOptimized': True,
        'ImageId': image.id,

        ssh_private_key_path = str(c.key_pair_path_private('automation'))

        for instance in instances:
            client = wait_for_ssh(
                instance.public_ip_address, 22,
                username='hg',
                key_filename=ssh_private_key_path)

            instance.ssh_client = client
            instance.ssh_private_key_path = ssh_private_key_path

        try:

        finally:
            for instance in instances:
                instance.ssh_client.close()


def ensure_windows_dev_ami(c: AWSConnection, prefix='hg-',
                           base_image_name=WINDOWS_BASE_IMAGE_NAME):
    """Ensure Windows Development AMI is available and up-to-date.

    If necessary, a modern AMI will be built by starting a temporary EC2
    instance and bootstrapping it.

    commands.insert(0, 'Set-MpPreference -DisableRealtimeMonitoring $true')
    commands.append('Set-MpPreference -DisableRealtimeMonitoring $false')

    # Compute a deterministic fingerprint to determine whether image needs
    # to be regenerated.
    fingerprint = resolve_fingerprint({
        'instance_config': config,
        'user_data': WINDOWS_USER_DATA,
        'initial_bootstrap': WINDOWS_BOOTSTRAP_POWERSHELL,
        'bootstrap_commands': commands,
        'base_image_name': base_image_name,
    })

    existing_image = find_and_reconcile_image(ec2resource, name, fingerprint)

    if existing_image:
        return existing_image

        print('installing Windows features...')
        run_ssm_command(
            ssmclient,
            [instance],
            'AWS-RunPowerShellScript',
            {
                'commands': WINDOWS_BOOTSTRAP_POWERSHELL.split('\n'),
            },
        )

        # Reboot so all updates are fully applied.
        #
        # We don't use instance.reboot() here because it is asynchronous and

        # a while to stop and we may start trying to interact with the instance
        # before it has rebooted.
        print('rebooting instance %s' % instance.id)
        instance.stop()
        ec2client.get_waiter('instance_stopped').wait(
            InstanceIds=[instance.id],
            WaiterConfig={
                'Delay': 5,
            })

        instance.start()
        wait_for_ip_addresses([instance])

        # There is a race condition here between the User Data PS script running
        # and us connecting to WinRM. This can manifest as
        # "AuthorizationManager check failed" failures during run_powershell().
        # TODO figure out a workaround.

        print('waiting for Windows Remote Management to come back...')
        client = wait_for_winrm(instance.public_ip_address, 'Administrator',
                                c.automation.default_password())
        print('established WinRM connection to %s' % instance.id)
        instance.winrm_client = client

        print('bootstrapping instance...')
        run_powershell(instance.winrm_client, '\n'.join(commands))

        print('bootstrap completed; stopping %s to create image' % instance.id)
        return create_ami_from_instance(ec2client, instance, name,
                                        'Mercurial Windows development environment',
                                        fingerprint)


@contextlib.contextmanager
def temporary_windows_dev_instances(c: AWSConnection, image, instance_type,
                                    prefix='hg-', disable_antivirus=False):
    """Create a temporary Windows development EC2 instance.

    Context manager resolves to the list of ``EC2.Instance`` that were created.
    """
    config = {

    with create_temp_windows_ec2_instances(c, config) as instances:
        if disable_antivirus:
            for instance in instances:
                run_powershell(
                    instance.winrm_client,
                    'Set-MpPreference -DisableRealtimeMonitoring $true')

        yield instances

43076:2372284d9457

import time

import boto3
import botocore.exceptions

from .linux import BOOTSTRAP_DEBIAN
from .ssh import (
    exec_command as ssh_exec_command,
    wait_for_ssh,
)
from .winrm import (
    run_powershell,
    wait_for_winrm,
)


SOURCE_ROOT = pathlib.Path(
    os.path.abspath(__file__)
).parent.parent.parent.parent

INSTALL_WINDOWS_DEPENDENCIES = (
    SOURCE_ROOT / 'contrib' / 'install-windows-dependencies.ps1'
)


INSTANCE_TYPES_WITH_STORAGE = {
    'c5d',
    'd2',

                    {
                        'CidrIp': '0.0.0.0/0',
                        'Description': 'RDP from entire Internet',
                    },
                ],
            },
            {
                'FromPort': 5985,
                'ToPort': 5986,
                'IpProtocol': 'tcp',

                    {
                        'CidrIp': '0.0.0.0/0',
                        'Description': 'PowerShell Remoting (Windows Remote Management)',
                    },
                ],
            },
        ],
    },
}

}
'''.strip()


IAM_INSTANCE_PROFILES = {
    'ephemeral-ec2-1': {'roles': ['ephemeral-ec2-role-1',],}
}


# User Data for Windows EC2 instance. Mainly used to set the password
# and configure WinRM.

class AWSConnection:
    """Manages the state of a connection with AWS."""

    def __init__(self, automation, region: str, ensure_ec2_state: bool = True):
        self.automation = automation
        self.local_state_path = automation.state_path

        self.prefix = 'hg-'

def rsa_key_fingerprint(p: pathlib.Path):
    """Compute the fingerprint of an RSA private key."""

    # TODO use rsa package.
    res = subprocess.run(
        [
            'openssl',
            'pkcs8',
            '-in',
            str(p),
            '-nocrypt',
            '-topk8',
            '-outform',
            'DER',
        ],
        capture_output=True,
        check=True,
    )

    sha1 = hashlib.sha1(res.stdout).hexdigest()
    return ':'.join(a + b for a, b in zip(sha1[::2], sha1[1::2]))


def ensure_key_pairs(state_path: pathlib.Path, ec2resource, prefix='hg-'):
    remote_existing = {}

    for kpi in ec2resource.key_pairs.all():
        if kpi.name.startswith(prefix):
            remote_existing[kpi.name[len(prefix) :]] = kpi.key_fingerprint

    # Validate that we have these keys locally.
    key_path = state_path / 'keys'
    key_path.mkdir(exist_ok=True, mode=0o700)

    for f in sorted(os.listdir(key_path)):
        if not f.startswith('keypair-') or not f.endswith('.pub'):
            continue

        name = f[len('keypair-') : -len('.pub')]

        pub_full = key_path / f
        priv_full = key_path / ('keypair-%s' % name)

        with open(pub_full, 'r', encoding='ascii') as fh:
            data = fh.read()

        if not data.startswith('ssh-rsa '):
            print(
                'unexpected format for key pair file: %s; removing' % pub_full
            )
            pub_full.unlink()
            priv_full.unlink()
            continue

        local_existing[name] = rsa_key_fingerprint(priv_full)

            print('local key %s does not exist remotely' % name)
            remove_local(name)
            del local_existing[name]

        elif remote_existing[name] != local_existing[name]:
            print(
                'key fingerprint mismatch for %s; '
                'removing from local and remote' % name
            )
            remove_local(name)
            remove_remote('%s%s' % (prefix, name))
            del local_existing[name]
            del remote_existing[name]

        # SSH public key can be extracted via `ssh-keygen`.
        with pub_full.open('w', encoding='ascii') as fh:
            subprocess.run(
                ['ssh-keygen', '-y', '-f', str(priv_full)],
                stdout=fh,
                check=True,
            )

        pub_full.chmod(0o0600)


def delete_instance_profile(profile):
    for role in profile.roles:
        print(
            'removing role %s from instance profile %s'
            % (role.name, profile.name)
        )
        profile.remove_role(RoleName=role.name)

    print('deleting instance profile %s' % profile.name)
    profile.delete()

    remote_profiles = {}

    for profile in iamresource.instance_profiles.all():
        if profile.name.startswith(prefix):
            remote_profiles[profile.name[len(prefix) :]] = profile

    for name in sorted(set(remote_profiles) - set(IAM_INSTANCE_PROFILES)):
        delete_instance_profile(remote_profiles[name])
        del remote_profiles[name]

    remote_roles = {}

    for role in iamresource.roles.all():
        if role.name.startswith(prefix):
            remote_roles[role.name[len(prefix) :]] = role

    for name in sorted(set(remote_roles) - set(IAM_ROLES)):
        role = remote_roles[name]

        print('removing role %s' % role.name)

    for name in sorted(set(IAM_INSTANCE_PROFILES) - set(remote_profiles)):
        actual = '%s%s' % (prefix, name)
        print('creating IAM instance profile %s' % actual)

        profile = iamresource.create_instance_profile(
            InstanceProfileName=actual
        )
        remote_profiles[name] = profile

        waiter = iamclient.get_waiter('instance_profile_exists')
        waiter.wait(InstanceProfileName=actual)
        print('IAM instance profile %s is available' % actual)

def find_image(ec2resource, owner_id, name):
    """Find an AMI by its owner ID and name."""

    images = ec2resource.images.filter(
        Filters=[
            {'Name': 'owner-id', 'Values': [owner_id],},
            {'Name': 'state', 'Values': ['available'],},
            {'Name': 'image-type', 'Values': ['machine'],},
            {'Name': 'name', 'Values': [name],},
        ]
    )

    for image in images:
        return image

    raise Exception('unable to find image for %s' % name)

    """
    existing = {}

    for group in ec2resource.security_groups.all():
        if group.group_name.startswith(prefix):
            existing[group.group_name[len(prefix) :]] = group

    purge = set(existing) - set(SECURITY_GROUPS)

    for name in sorted(purge):
        group = existing[name]

        actual = '%s%s' % (prefix, name)
        print('adding security group %s' % actual)

        group_res = ec2resource.create_security_group(
            Description=group['description'], GroupName=actual,
        )

        group_res.authorize_ingress(IpPermissions=group['ingress'],)

        security_groups[name] = group_res

    return security_groups

            if not instance.public_ip_address:
                time.sleep(2)
                instance.reload()
                continue

            print(
                'public IP address for %s: %s'
                % (instance.id, instance.public_ip_address)
            )
            break


def remove_ami(ec2resource, image):
    """Remove an AMI and its underlying snapshots."""

def wait_for_ssm(ssmclient, instances):
    """Wait for SSM to come online for an iterable of instance IDs."""
    while True:
        res = ssmclient.describe_instance_information(
            Filters=[
                {'Key': 'InstanceIds', 'Values': [i.id for i in instances],},
            ],
        )

        available = len(res['InstanceInformationList'])
        wanted = len(instances)

    res = ssmclient.send_command(
        InstanceIds=[i.id for i in instances],
        DocumentName=document_name,
        Parameters=parameters,
        CloudWatchOutputConfig={'CloudWatchOutputEnabled': True,},
    )

    command_id = res['Command']['CommandId']

    for instance in instances:
        while True:
            try:
                res = ssmclient.get_command_invocation(
                    CommandId=command_id, InstanceId=instance.id,
                )
            except botocore.exceptions.ClientError as e:
                if e.response['Error']['Code'] == 'InvocationDoesNotExist':
                    print('could not find SSM command invocation; waiting')
                    time.sleep(1)

            if res['Status'] == 'Success':
                break
            elif res['Status'] in ('Pending', 'InProgress', 'Delayed'):
                time.sleep(2)
            else:
                raise Exception(
                    'command failed on %s: %s' % (instance.id, res['Status'])
                )


@contextlib.contextmanager
def temporary_ec2_instances(ec2resource, config):
    """Create temporary EC2 instances.

    config = copy.deepcopy(config)
    config['IamInstanceProfile'] = {
        'Name': 'hg-ephemeral-ec2-1',
    }
    config.setdefault('TagSpecifications', []).append(
        {
            'ResourceType': 'instance',
            'Tags': [{'Key': 'Name', 'Value': 'hg-temp-windows'}],
        }
    )
    config['UserData'] = WINDOWS_USER_DATA % password

    with temporary_ec2_instances(c.ec2resource, config) as instances:
        wait_for_ip_addresses(instances)

        print('waiting for Windows Remote Management service...')

        for instance in instances:
            client = wait_for_winrm(
                instance.public_ip_address, 'Administrator', password
            )
            print('established WinRM connection to %s' % instance.id)
            instance.winrm_client = client

        yield instances

    """
    # Find existing AMIs with this name and delete the ones that are invalid.
    # Store a reference to a good image so it can be returned one the
    # image state is reconciled.
    images = ec2resource.images.filter(
        Filters=[{'Name': 'name', 'Values': [name]}]
    )

    existing_image = None

    for image in images:
        if image.tags is None:
            print(
                'image %s for %s lacks required tags; removing'
                % (image.id, image.name)
            )
            remove_ami(ec2resource, image)
        else:
            tags = {t['Key']: t['Value'] for t in image.tags}

            if tags.get('HGIMAGEFINGERPRINT') == fingerprint:
                existing_image = image
            else:
                print(
                    'image %s for %s has wrong fingerprint; removing'
                    % (image.id, image.name)
                )
                remove_ami(ec2resource, image)

    return existing_image


def create_ami_from_instance(
    ec2client, instance, name, description, fingerprint
):
    """Create an AMI from a running instance.

    Returns the ``ec2resource.Image`` representing the created AMI.
    """
    instance.stop()

    ec2client.get_waiter('instance_stopped').wait(
        InstanceIds=[instance.id], WaiterConfig={'Delay': 5,}
    )
    print('%s is stopped' % instance.id)

    image = instance.create_image(Name=name, Description=description,)

    image.create_tags(
        Tags=[{'Key': 'HGIMAGEFINGERPRINT', 'Value': fingerprint,},]
    )

    print('waiting for image %s' % image.id)

    ec2client.get_waiter('image_available').wait(ImageIds=[image.id],)

    print('image %s available as %s' % (image.id, image.name))

    return image

            'debian-stretch-hvm-x86_64-gp2-2019-09-08-17994',
        )
        ssh_username = 'admin'
    elif distro == 'debian10':
        image = find_image(
            ec2resource, DEBIAN_ACCOUNT_ID_2, 'debian-10-amd64-20190909-10',
        )
        ssh_username = 'admin'
    elif distro == 'ubuntu18.04':
        image = find_image(
            ec2resource,

        'MaxCount': 1,
        'MinCount': 1,
        'SecurityGroupIds': [c.security_groups['linux-dev-1'].id],
    }

    requirements2_path = (
        pathlib.Path(__file__).parent.parent / 'linux-requirements-py2.txt'
    )
    requirements3_path = (
        pathlib.Path(__file__).parent.parent / 'linux-requirements-py3.txt'
    )
    with requirements2_path.open('r', encoding='utf-8') as fh:
        requirements2 = fh.read()
    with requirements3_path.open('r', encoding='utf-8') as fh:
        requirements3 = fh.read()

    # Compute a deterministic fingerprint to determine whether image needs to
    # be regenerated.
    fingerprint = resolve_fingerprint(
        {
            'instance_config': config,
            'bootstrap_script': BOOTSTRAP_DEBIAN,
            'requirements_py2': requirements2,
            'requirements_py3': requirements3,
        }
    )

    existing_image = find_and_reconcile_image(ec2resource, name, fingerprint)

    if existing_image:
        return existing_image

        wait_for_ip_addresses(instances)

        instance = instances[0]

        client = wait_for_ssh(
            instance.public_ip_address,
            22,
            username=ssh_username,
            key_filename=str(c.key_pair_path_private('automation')),
        )

        home = '/home/%s' % ssh_username

        with client:
            print('connecting to SSH server')

            with sftp.open('%s/requirements-py3.txt' % home, 'wb') as fh:
                fh.write(requirements3)
                fh.chmod(0o0700)

            print('executing bootstrap')
            chan, stdin, stdout = ssh_exec_command(
                client, '%s/bootstrap' % home
            )
            stdin.close()

            for line in stdout:
                print(line, end='')

            res = chan.recv_exit_status()
            if res:
                raise Exception('non-0 exit from bootstrap: %d' % res)

            print(
                'bootstrap completed; stopping %s to create %s'
                % (instance.id, name)
            )

        return create_ami_from_instance(
            ec2client,
            instance,
            name,
            'Mercurial Linux development environment',
            fingerprint,
        )


@contextlib.contextmanager
def temporary_linux_dev_instances(
    c: AWSConnection,
    image,
    instance_type,
    prefix='hg-',
    ensure_extra_volume=False,
):
    """Create temporary Linux development EC2 instances.

    Context manager resolves to a list of ``ec2.Instance`` that were created
    and are running.

        }
    ]

    # This is not an exhaustive list of instance types having instance storage.
    # But
    if ensure_extra_volume and not instance_type.startswith(
        tuple(INSTANCE_TYPES_WITH_STORAGE)
    ):
        main_device = block_device_mappings[0]['DeviceName']

        if main_device == 'xvda':
            second_device = 'xvdb'
        elif main_device == '/dev/sda1':
            second_device = '/dev/sdb'
        else:
            raise ValueError(
                'unhandled primary EBS device name: %s' % main_device
            )

        block_device_mappings.append(
            {
                'DeviceName': second_device,
                'Ebs': {
                    'DeleteOnTermination': True,
                    'VolumeSize': 8,
                    'VolumeType': 'gp2',
                },
            }
        )

    config = {
        'BlockDeviceMappings': block_device_mappings,
        'EbsOptimized': True,
        'ImageId': image.id,

        ssh_private_key_path = str(c.key_pair_path_private('automation'))

        for instance in instances:
            client = wait_for_ssh(
                instance.public_ip_address,
                22,
                username='hg',
                key_filename=ssh_private_key_path,
            )

            instance.ssh_client = client
            instance.ssh_private_key_path = ssh_private_key_path

        try:

        finally:
            for instance in instances:
                instance.ssh_client.close()


def ensure_windows_dev_ami(
    c: AWSConnection, prefix='hg-', base_image_name=WINDOWS_BASE_IMAGE_NAME
):
    """Ensure Windows Development AMI is available and up-to-date.

    If necessary, a modern AMI will be built by starting a temporary EC2
    instance and bootstrapping it.

    commands.insert(0, 'Set-MpPreference -DisableRealtimeMonitoring $true')
    commands.append('Set-MpPreference -DisableRealtimeMonitoring $false')

    # Compute a deterministic fingerprint to determine whether image needs
    # to be regenerated.
    fingerprint = resolve_fingerprint(
        {
            'instance_config': config,
            'user_data': WINDOWS_USER_DATA,
            'initial_bootstrap': WINDOWS_BOOTSTRAP_POWERSHELL,
            'bootstrap_commands': commands,
            'base_image_name': base_image_name,
        }
    )

    existing_image = find_and_reconcile_image(ec2resource, name, fingerprint)

    if existing_image:
        return existing_image

        print('installing Windows features...')
        run_ssm_command(
            ssmclient,
            [instance],
            'AWS-RunPowerShellScript',
            {'commands': WINDOWS_BOOTSTRAP_POWERSHELL.split('\n'),},
        )

        # Reboot so all updates are fully applied.
        #
        # We don't use instance.reboot() here because it is asynchronous and

        # a while to stop and we may start trying to interact with the instance
        # before it has rebooted.
        print('rebooting instance %s' % instance.id)
        instance.stop()
        ec2client.get_waiter('instance_stopped').wait(
            InstanceIds=[instance.id], WaiterConfig={'Delay': 5,}
        )

        instance.start()
        wait_for_ip_addresses([instance])

        # There is a race condition here between the User Data PS script running
        # and us connecting to WinRM. This can manifest as
        # "AuthorizationManager check failed" failures during run_powershell().
        # TODO figure out a workaround.

        print('waiting for Windows Remote Management to come back...')
        client = wait_for_winrm(
            instance.public_ip_address,
            'Administrator',
            c.automation.default_password(),
        )
        print('established WinRM connection to %s' % instance.id)
        instance.winrm_client = client

        print('bootstrapping instance...')
        run_powershell(instance.winrm_client, '\n'.join(commands))

        print('bootstrap completed; stopping %s to create image' % instance.id)
        return create_ami_from_instance(
            ec2client,
            instance,
            name,
            'Mercurial Windows development environment',
            fingerprint,
        )


@contextlib.contextmanager
def temporary_windows_dev_instances(
    c: AWSConnection,
    image,
    instance_type,
    prefix='hg-',
    disable_antivirus=False,
):
    """Create a temporary Windows development EC2 instance.

    Context manager resolves to the list of ``EC2.Instance`` that were created.
    """
    config = {

    with create_temp_windows_ec2_instances(c, config) as instances:
        if disable_antivirus:
            for instance in instances:
                run_powershell(
                    instance.winrm_client,
                    'Set-MpPreference -DisableRealtimeMonitoring $true',
                )

        yield instances